Skip to content

Commit fc5ebc8

Browse files
infvginfvg
committed
feat(iceberg): Added dictionary size bytes write config
1 parent d666ceb commit fc5ebc8

7 files changed

Lines changed: 132 additions & 3 deletions

File tree

‎backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/AbstractIcebergWriteExec.scala‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
package org.apache.gluten.execution
1818

1919
import org.apache.gluten.IcebergNestedFieldVisitor
20-
import org.apache.gluten.config.VeloxConfig.{MAX_TARGET_FILE_SIZE_SESSION, PARQUET_PAGE_SIZE_BYTES}
20+
import org.apache.gluten.config.VeloxConfig.{MAX_TARGET_FILE_SIZE_SESSION, PARQUET_DICT_SIZE_BYTES, PARQUET_PAGE_SIZE_BYTES}
2121
import org.apache.gluten.connector.write.{ColumnarBatchDataWriterFactory, ColumnarStreamingDataWriterFactory, IcebergDataWriteFactory}
2222

2323
import org.apache.spark.sql.internal.SQLConf
@@ -48,7 +48,8 @@ abstract class AbstractIcebergWriteExec extends IcebergWriteExec {
4848

4949
Seq(
5050
PARQUET_PAGE_SIZE_BYTES.key -> getParquetPageSizeBytes,
51-
MAX_TARGET_FILE_SIZE_SESSION.key -> getTargetFileSizeBytes
51+
MAX_TARGET_FILE_SIZE_SESSION.key -> getTargetFileSizeBytes,
52+
PARQUET_DICT_SIZE_BYTES.key -> getDictSizeBytes
5253
).foreach {
5354
case (key, value) =>
5455
if (SQLConf.get.getConfString(key, null) == null) {

‎backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ import org.apache.spark.sql.gluten.TestUtils
2929

3030
import org.apache.hadoop.fs.Path
3131
import org.apache.iceberg.shaded.org.apache.parquet.ParquetReadOptions
32+
import org.apache.iceberg.shaded.org.apache.parquet.column.Encoding
33+
import org.apache.iceberg.shaded.org.apache.parquet.column.page.{DataPage, DataPageV1, DataPageV2}
3234
import org.apache.iceberg.shaded.org.apache.parquet.hadoop.ParquetFileReader
3335
import org.apache.iceberg.shaded.org.apache.parquet.hadoop.util.HadoopInputFile
3436

@@ -543,6 +545,114 @@ class VeloxIcebergSuite extends IcebergSuite {
543545
}
544546
}
545547
}
548+
549+
test("iceberg parquet writer respects dictionary page size bytes") {
550+
val table = "iceberg_dict_page_size_tbl"
551+
552+
def parquetFiles(table: String): Seq[String] = {
553+
spark.sql(s"""
554+
|SELECT file_path
555+
|FROM default.$table.files
556+
|""".stripMargin).collect().map(_.getString(0)).toSeq
557+
}
558+
559+
def pageEncoding(page: DataPage): Encoding = {
560+
page.accept(new DataPage.Visitor[Encoding] {
561+
override def visit(dataPageV1: DataPageV1): Encoding = dataPageV1.getValueEncoding
562+
override def visit(dataPageV2: DataPageV2): Encoding = dataPageV2.getDataEncoding
563+
})
564+
}
565+
566+
def dataPageEncodings(table: String, columnName: String): Seq[Encoding] = {
567+
val conf = spark.sparkContext.hadoopConfiguration
568+
569+
parquetFiles(table).flatMap {
570+
file =>
571+
val inputFile = HadoopInputFile.fromPath(new Path(file), conf)
572+
val reader = ParquetFileReader.open(inputFile, ParquetReadOptions.builder().build())
573+
574+
try {
575+
val column = reader
576+
.getFooter
577+
.getFileMetaData
578+
.getSchema
579+
.getColumns
580+
.asScala
581+
.find(_.getPath.toSeq == Seq(columnName))
582+
.getOrElse {
583+
fail(s"Column $columnName was not found in Parquet file $file")
584+
}
585+
586+
val encodings = scala.collection.mutable.ArrayBuffer.empty[Encoding]
587+
588+
var rowGroup = reader.readNextRowGroup()
589+
while (rowGroup != null) {
590+
val pageReader = rowGroup.getPageReader(column)
591+
pageReader.readDictionaryPage()
592+
593+
var page = pageReader.readPage()
594+
while (page != null) {
595+
encodings += pageEncoding(page)
596+
page = pageReader.readPage()
597+
}
598+
599+
rowGroup = reader.readNextRowGroup()
600+
}
601+
602+
encodings
603+
} finally {
604+
reader.close()
605+
}
606+
}
607+
}
608+
609+
withSQLConf(
610+
"spark.sql.shuffle.partitions" -> "1"
611+
) {
612+
withTable(table) {
613+
spark.sql(s"""
614+
|CREATE TABLE $table (
615+
| value SMALLINT
616+
|) USING iceberg
617+
|TBLPROPERTIES (
618+
| 'write.format.default' = 'parquet',
619+
| 'write.parquet.compression-codec' = 'uncompressed',
620+
| 'write.parquet.dict-size-bytes' = '1B'
621+
|)
622+
|""".stripMargin)
623+
624+
val df = spark.sql(s"""
625+
|INSERT INTO $table
626+
|SELECT CAST(id + 1 AS SMALLINT)
627+
|FROM range(0, 10000, 1, 1)
628+
|""".stripMargin)
629+
630+
assert(
631+
df.queryExecution.executedPlan
632+
.asInstanceOf[CommandResultExec]
633+
.commandPhysicalPlan
634+
.isInstanceOf[VeloxIcebergAppendDataExec])
635+
636+
checkAnswer(
637+
spark.sql(s"SELECT count(*) FROM $table"),
638+
Seq(Row(10000L)))
639+
640+
val encodings = dataPageEncodings(table, "value")
641+
642+
assert(encodings.nonEmpty, "Expected at least one Parquet data page")
643+
assert(
644+
encodings.head == Encoding.RLE_DICTIONARY,
645+
s"Expected the first data page to use dictionary encoding, " +
646+
s"but got encodings=${encodings.mkString("[", ", ", "]")}"
647+
)
648+
assert(
649+
encodings.contains(Encoding.PLAIN),
650+
s"Expected write.parquet.dict-size-bytes=1B to make later data pages fall back " +
651+
s"to PLAIN, but got encodings=${encodings.mkString("[", ", ", "]")}"
652+
)
653+
}
654+
}
655+
}
546656
test("iceberg parquet writer default row group size test") {
547657
val table = "iceberg_default_row_group_size"
548658
val defaultRowGroupBytes = 128L * 1024 * 1024

‎backends-velox/src/main/scala/org/apache/gluten/config/VeloxConfig.scala‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -862,6 +862,12 @@ object VeloxConfig extends ConfigRegistry {
862862
.bytesConf(ByteUnit.BYTE)
863863
.createWithDefaultString("1MB")
864864

865+
val PARQUET_DICT_SIZE_BYTES =
866+
buildConf("spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes")
867+
.doc("The maximum size in bytes for a Parquet dictionary page")
868+
.bytesConf(ByteUnit.BYTE)
869+
.createWithDefaultString("2MB")
870+
865871
val ENABLE_TIMESTAMP_NTZ_VALIDATION =
866872
buildConf("spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation")
867873
.doc(

‎cpp/velox/config/VeloxConfig.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,8 @@ const uint32_t kGlogSeverityLevelDefault = 1;
191191

192192
// Iceberg write configs
193193
const std::string kWriteParquetPageSizeBytes = "spark.gluten.sql.columnar.backend.velox.parquet.pageSizeBytes";
194+
const std::string kWriteParquetDictSizeBytes =
195+
"spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes";
194196

195197
// Query trace
196198
/// Enable query tracing flag.

‎cpp/velox/utils/ConfigExtractor.cc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,8 @@ std::shared_ptr<facebook::velox::config::ConfigBase> createHiveConnectorSessionC
252252
conf->get<bool>(kOrcUseColumnNames, true) ? "true" : "false";
253253
configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kWriterPageSizeSession)] =
254254
conf->get<std::string>(kWriteParquetPageSizeBytes, "1MB");
255+
configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kWriterDictionaryPageSizeLimitSession)] =
256+
conf->get<std::string>(kWriteParquetDictSizeBytes, "2MB");
255257
configs[parquetSessionProperty(facebook::velox::parquet::ParquetConfig::kNullStructIfAllFieldsMissingSession)] =
256258
"true";
257259

‎docs/velox-configuration.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ nav_order: 16
5656
| spark.gluten.sql.columnar.backend.velox.memoryUseHugePages | 🔄 Dynamic | false | Use explicit huge pages for Velox memory allocation. |
5757
| spark.gluten.sql.columnar.backend.velox.orc.scan.enabled | 🔄 Dynamic | true | Enable velox orc scan. If disabled, vanilla spark orc scan will be used. |
5858
| spark.gluten.sql.columnar.backend.velox.orcUseColumnNames | 🔄 Dynamic | true | Maps table field names to file field names using names, not indices for ORC files. |
59+
| spark.gluten.sql.columnar.backend.velox.parquet.dictionaryPageSizeBytes | 🔄 Dynamic | 2MB | The maximum size in bytes for a Parquet dictionary page |
5960
| spark.gluten.sql.columnar.backend.velox.parquet.pageSizeBytes | 🔄 Dynamic | 1MB | The page size in bytes is for compression. |
6061
| 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. |
6162
| spark.gluten.sql.columnar.backend.velox.parquetUseColumnNames | 🔄 Dynamic | true | Maps table field names to file field names using names, not indices for Parquet files. |

‎gluten-iceberg/src/main/scala/org/apache/gluten/execution/IcebergWriteExec.scala‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ package org.apache.gluten.execution
1919
import org.apache.gluten.backendsapi.BackendsApiManager
2020

2121
import org.apache.iceberg.{FileFormat, PartitionField, PartitionSpec, Schema, TableProperties}
22-
import org.apache.iceberg.TableProperties.{ORC_COMPRESSION, ORC_COMPRESSION_DEFAULT, PARQUET_COMPRESSION, PARQUET_COMPRESSION_DEFAULT, PARQUET_PAGE_SIZE_BYTES, PARQUET_PAGE_SIZE_BYTES_DEFAULT}
22+
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}
2323
import org.apache.iceberg.avro.AvroSchemaUtil
2424
import org.apache.iceberg.spark.source.IcebergWriteUtil
2525
import org.apache.iceberg.types.Type.TypeID
@@ -60,6 +60,13 @@ trait IcebergWriteExec extends ColumnarV2TableWriteExec {
6060
IcebergWriteUtil.getWriteConf(write).targetDataFileSize().toString
6161
}
6262

63+
protected def getDictSizeBytes: String = {
64+
val tableProps = IcebergWriteUtil.getTable(write).properties()
65+
tableProps.getOrDefault(
66+
normalizeCapacityString(PARQUET_DICT_SIZE_BYTES),
67+
normalizeCapacityString(PARQUET_DICT_SIZE_BYTES_DEFAULT.toString))
68+
}
69+
6370
protected def getPartitionSpec: PartitionSpec = {
6471
IcebergWriteUtil.getPartitionSpec(write)
6572
}

0 commit comments

Comments
 (0)