From fc5ebc81847be51446e491abff3d56bcc84e6e5a Mon Sep 17 00:00:00 2001 From: Hazmi Date: Thu, 2 Jul 2026 13:14:49 +0000 Subject: [PATCH] feat(iceberg): Added dictionary size bytes write config --- .../execution/AbstractIcebergWriteExec.scala | 5 +- .../enhanced/VeloxIcebergSuite.scala | 110 ++++++++++++++++++ .../apache/gluten/config/VeloxConfig.scala | 6 + cpp/velox/config/VeloxConfig.h | 2 + cpp/velox/utils/ConfigExtractor.cc | 2 + docs/velox-configuration.md | 1 + .../gluten/execution/IcebergWriteExec.scala | 9 +- 7 files changed, 132 insertions(+), 3 deletions(-) diff --git a/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/AbstractIcebergWriteExec.scala b/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/AbstractIcebergWriteExec.scala index 2b4d50a73f0..37e6f30937d 100644 --- a/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/AbstractIcebergWriteExec.scala +++ b/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/AbstractIcebergWriteExec.scala @@ -17,7 +17,7 @@ package org.apache.gluten.execution import org.apache.gluten.IcebergNestedFieldVisitor -import org.apache.gluten.config.VeloxConfig.{MAX_TARGET_FILE_SIZE_SESSION, PARQUET_PAGE_SIZE_BYTES} +import org.apache.gluten.config.VeloxConfig.{MAX_TARGET_FILE_SIZE_SESSION, PARQUET_DICT_SIZE_BYTES, PARQUET_PAGE_SIZE_BYTES} import org.apache.gluten.connector.write.{ColumnarBatchDataWriterFactory, ColumnarStreamingDataWriterFactory, IcebergDataWriteFactory} import org.apache.spark.sql.internal.SQLConf @@ -48,7 +48,8 @@ abstract class AbstractIcebergWriteExec extends IcebergWriteExec { Seq( PARQUET_PAGE_SIZE_BYTES.key -> getParquetPageSizeBytes, - MAX_TARGET_FILE_SIZE_SESSION.key -> getTargetFileSizeBytes + MAX_TARGET_FILE_SIZE_SESSION.key -> getTargetFileSizeBytes, + PARQUET_DICT_SIZE_BYTES.key -> getDictSizeBytes ).foreach { case (key, value) => if (SQLConf.get.getConfString(key, null) == null) { diff --git a/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala b/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala index 0cc6cddda2d..de485499934 100644 --- a/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala +++ b/backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala @@ -29,6 +29,8 @@ import org.apache.spark.sql.gluten.TestUtils import org.apache.hadoop.fs.Path import org.apache.iceberg.shaded.org.apache.parquet.ParquetReadOptions +import org.apache.iceberg.shaded.org.apache.parquet.column.Encoding +import org.apache.iceberg.shaded.org.apache.parquet.column.page.{DataPage, DataPageV1, DataPageV2} import org.apache.iceberg.shaded.org.apache.parquet.hadoop.ParquetFileReader import org.apache.iceberg.shaded.org.apache.parquet.hadoop.util.HadoopInputFile @@ -543,6 +545,114 @@ class VeloxIcebergSuite extends IcebergSuite { } } } + + test("iceberg parquet writer respects dictionary page size bytes") { + val table = "iceberg_dict_page_size_tbl" + + def parquetFiles(table: String): Seq[String] = { + spark.sql(s""" + |SELECT file_path + |FROM default.$table.files + |""".stripMargin).collect().map(_.getString(0)).toSeq + } + + def pageEncoding(page: DataPage): Encoding = { + page.accept(new DataPage.Visitor[Encoding] { + override def visit(dataPageV1: DataPageV1): Encoding = dataPageV1.getValueEncoding + override def visit(dataPageV2: DataPageV2): Encoding = dataPageV2.getDataEncoding + }) + } + + def dataPageEncodings(table: String, columnName: String): Seq[Encoding] = { + val conf = spark.sparkContext.hadoopConfiguration + + parquetFiles(table).flatMap { + file => + val inputFile = HadoopInputFile.fromPath(new Path(file), conf) + val reader = ParquetFileReader.open(inputFile, ParquetReadOptions.builder().build()) + + try { + val column = reader + .getFooter + .getFileMetaData + .getSchema + .getColumns + .asScala + .find(_.getPath.toSeq == Seq(columnName)) + .getOrElse { + fail(s"Column $columnName was not found in Parquet file $file") + } + + val encodings = scala.collection.mutable.ArrayBuffer.empty[Encoding] + + var rowGroup = reader.readNextRowGroup() + while (rowGroup != null) { + val pageReader = rowGroup.getPageReader(column) + pageReader.readDictionaryPage() + + var page = pageReader.readPage() + while (page != null) { + encodings += pageEncoding(page) + page = pageReader.readPage() + } + + rowGroup = reader.readNextRowGroup() + } + + encodings + } finally { + reader.close() + } + } + } + + withSQLConf( + "spark.sql.shuffle.partitions" -> "1" + ) { + withTable(table) { + spark.sql(s""" + |CREATE TABLE $table ( + | value SMALLINT + |) USING iceberg + |TBLPROPERTIES ( + | 'write.format.default' = 'parquet', + | 'write.parquet.compression-codec' = 'uncompressed', + | 'write.parquet.dict-size-bytes' = '1B' + |) + |""".stripMargin) + + val df = spark.sql(s""" + |INSERT INTO $table + |SELECT CAST(id + 1 AS SMALLINT) + |FROM range(0, 10000, 1, 1) + |""".stripMargin) + + assert( + df.queryExecution.executedPlan + .asInstanceOf[CommandResultExec] + .commandPhysicalPlan + .isInstanceOf[VeloxIcebergAppendDataExec]) + + checkAnswer( + spark.sql(s"SELECT count(*) FROM $table"), + Seq(Row(10000L))) + + val encodings = dataPageEncodings(table, "value") + + assert(encodings.nonEmpty, "Expected at least one Parquet data page") + assert( + encodings.head == Encoding.RLE_DICTIONARY, + s"Expected the first data page to use dictionary encoding, " + + s"but got encodings=${encodings.mkString("[", ", ", "]")}" + ) + assert( + encodings.contains(Encoding.PLAIN), + s"Expected write.parquet.dict-size-bytes=1B to make later data pages fall back " + + s"to PLAIN, but got encodings=${encodings.mkString("[", ", ", "]")}" + ) + } + } + } test("iceberg parquet writer default row group size test") { val table = "iceberg_default_row_group_size" val defaultRowGroupBytes = 128L * 1024 * 1024 diff --git a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala index dacb835b8c5..58645290586 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala @@ -862,6 +862,12 @@ object VeloxConfig extends ConfigRegistry { .bytesConf(ByteUnit.BYTE) .createWithDefaultString("1MB") + val PARQUET_DICT_SIZE_BYTES = + buildConf("spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes") + .doc("The maximum size in bytes for a Parquet dictionary page") + .bytesConf(ByteUnit.BYTE) + .createWithDefaultString("2MB") + val ENABLE_TIMESTAMP_NTZ_VALIDATION = buildConf("spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation") .doc( diff --git a/cpp/velox/config/VeloxConfig.h b/cpp/velox/config/VeloxConfig.h index f16f48c40de..5b5e129c342 100644 --- a/cpp/velox/config/VeloxConfig.h +++ b/cpp/velox/config/VeloxConfig.h @@ -191,6 +191,8 @@ const uint32_t kGlogSeverityLevelDefault = 1; // Iceberg write configs const std::string kWriteParquetPageSizeBytes = "spark.gluten.sql.columnar.backend.velox.parquet.pageSizeBytes"; +const std::string kWriteParquetDictSizeBytes = + "spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes"; // Query trace /// Enable query tracing flag. diff --git a/cpp/velox/utils/ConfigExtractor.cc b/cpp/velox/utils/ConfigExtractor.cc index f398c2588d3..ecb3d9e7ada 100644 --- a/cpp/velox/utils/ConfigExtractor.cc +++ b/cpp/velox/utils/ConfigExtractor.cc @@ -252,6 +252,8 @@ std::shared_ptr createHiveConnectorSessionC conf->get(kOrcUseColumnNames, true) ? "true" : "false"; configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kWriterPageSizeSession)] = conf->get(kWriteParquetPageSizeBytes, "1MB"); + configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kWriterDictionaryPageSizeLimitSession)] = + conf->get(kWriteParquetDictSizeBytes, "2MB"); configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kNullStructIfAllFieldsMissingSession)] = "true"; diff --git a/docs/velox-configuration.md b/docs/velox-configuration.md index 97720808878..5f7c49afc46 100644 --- a/docs/velox-configuration.md +++ b/docs/velox-configuration.md @@ -56,6 +56,7 @@ nav_order: 16 | spark.gluten.sql.columnar.backend.velox.memoryUseHugePages | 🔄 Dynamic | false | Use explicit huge pages for Velox memory allocation. | | spark.gluten.sql.columnar.backend.velox.orc.scan.enabled | 🔄 Dynamic | true | Enable velox orc scan. If disabled, vanilla spark orc scan will be used. | | spark.gluten.sql.columnar.backend.velox.orcUseColumnNames | 🔄 Dynamic | true | Maps table field names to file field names using names, not indices for ORC files. | +| spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes | 🔄 Dynamic | 2MB | The maximum size in bytes for a Parquet dictionary page | | spark.gluten.sql.columnar.backend.velox.parquet.pageSizeBytes | 🔄 Dynamic | 1MB | The page size in bytes is for compression. | | spark.gluten.sql.columnar.backend.velox.parquetMaxTargetFileSize | 🔄 Dynamic | 0b | The target file size for each output file when writing data. 0 means no limit on target file size, and the actual file size will be determined by other factors such as max partition number and shuffle batch size. | | spark.gluten.sql.columnar.backend.velox.parquetUseColumnNames | 🔄 Dynamic | true | Maps table field names to file field names using names, not indices for Parquet files. | diff --git a/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergWriteExec.scala b/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergWriteExec.scala index 0544e47c1dd..270745063af 100644 --- a/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergWriteExec.scala +++ b/gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergWriteExec.scala @@ -19,7 +19,7 @@ package org.apache.gluten.execution import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.iceberg.{FileFormat, PartitionField, PartitionSpec, Schema, TableProperties} -import org.apache.iceberg.TableProperties.{ORC_COMPRESSION, ORC_COMPRESSION_DEFAULT, PARQUET_COMPRESSION, PARQUET_COMPRESSION_DEFAULT, PARQUET_PAGE_SIZE_BYTES, PARQUET_PAGE_SIZE_BYTES_DEFAULT} +import org.apache.iceberg.TableProperties.{ORC_COMPRESSION, ORC_COMPRESSION_DEFAULT, PARQUET_COMPRESSION, PARQUET_COMPRESSION_DEFAULT, PARQUET_DICT_SIZE_BYTES, PARQUET_DICT_SIZE_BYTES_DEFAULT, PARQUET_PAGE_SIZE_BYTES, PARQUET_PAGE_SIZE_BYTES_DEFAULT} import org.apache.iceberg.avro.AvroSchemaUtil import org.apache.iceberg.spark.source.IcebergWriteUtil import org.apache.iceberg.types.Type.TypeID @@ -60,6 +60,13 @@ trait IcebergWriteExec extends ColumnarV2TableWriteExec { IcebergWriteUtil.getWriteConf(write).targetDataFileSize().toString } + protected def getDictSizeBytes: String = { + val tableProps = IcebergWriteUtil.getTable(write).properties() + tableProps.getOrDefault( + normalizeCapacityString(PARQUET_DICT_SIZE_BYTES), + normalizeCapacityString(PARQUET_DICT_SIZE_BYTES_DEFAULT.toString)) + } + protected def getPartitionSpec: PartitionSpec = { IcebergWriteUtil.getPartitionSpec(write) }