From f033e4d32ccbc79399f0aa966a0cc2754005bfd6 Mon Sep 17 00:00:00 2001 From: Fokko Date: Mon, 16 Jun 2025 13:32:04 +0200 Subject: [PATCH 1/2] Spark: Support Parquet dictionary encoded UUIDs While fixing some issues on the PyIceberg ends to fully support UUIDs: https://github.com/apache/iceberg-python/pull/2007 I noticed this issue, and was suprised since UUID used to work with Spark, but it turns out that the dictionary encoded UUID was not implemented yet. For PyIceberg we only generate little data, so therefore this wasn't caught previously. --- .../GenericArrowVectorAccessorFactory.java | 22 +++++++++++++++++-- .../ArrowVectorAccessorFactory.java | 8 +++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/GenericArrowVectorAccessorFactory.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/GenericArrowVectorAccessorFactory.java index ecbc1cf1d993..2e24ce2e87a8 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/GenericArrowVectorAccessorFactory.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/GenericArrowVectorAccessorFactory.java @@ -156,7 +156,8 @@ public ArrowVectorAccessor getVecto switch (primitive.getPrimitiveTypeName()) { case FIXED_LEN_BYTE_ARRAY: case BINARY: - return new DictionaryBinaryAccessor<>((IntVector) vector, dictionary); + return new DictionaryBinaryAccessor<>( + (IntVector) vector, dictionary, stringFactorySupplier.get()); case FLOAT: return new DictionaryFloatAccessor<>((IntVector) vector, dictionary); case INT64: @@ -452,17 +453,27 @@ private static class DictionaryBinaryAccessor< extends ArrowVectorAccessor { private final IntVector offsetVector; private final Dictionary dictionary; + private final StringFactory stringFactory; - DictionaryBinaryAccessor(IntVector vector, Dictionary dictionary) { + DictionaryBinaryAccessor( + IntVector vector, Dictionary dictionary, StringFactory stringFactory) { super(vector); this.offsetVector = vector; this.dictionary = dictionary; + this.stringFactory = stringFactory; } @Override public final byte[] getBinary(int rowId) { return dictionary.decodeToBinary(offsetVector.get(rowId)).getBytes(); } + + @Override + public Utf8StringT getUTF8String(int rowId) { + return null == stringFactory + ? super.getUTF8String(rowId) + : stringFactory.ofRow(offsetVector, dictionary, rowId); + } } private static class DictionaryTimestampInt96Accessor< @@ -815,6 +826,13 @@ default Utf8StringT ofRow(FixedSizeBinaryVector vector, int rowId) { getGenericClass().getSimpleName())); } + /** Create a UTF8 String from the row value in the Dictionary. */ + default Utf8StringT ofRow(IntVector offsetVector, Dictionary dictionary, int rowId) { + throw new UnsupportedOperationException( + String.format( + "Creating %s from a Dictionary is not supported", getGenericClass().getSimpleName())); + } + /** Create a UTF8 String from the byte array. */ Utf8StringT ofBytes(byte[] bytes); diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ArrowVectorAccessorFactory.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ArrowVectorAccessorFactory.java index 29e938bb092e..b4bb9a918732 100644 --- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ArrowVectorAccessorFactory.java +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ArrowVectorAccessorFactory.java @@ -22,11 +22,13 @@ import java.nio.ByteBuffer; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.FixedSizeBinaryVector; +import org.apache.arrow.vector.IntVector; import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.complex.ListVector; import org.apache.iceberg.arrow.vectorized.GenericArrowVectorAccessorFactory; import org.apache.iceberg.util.UUIDUtil; +import org.apache.parquet.column.Dictionary; import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.vectorized.ArrowColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; @@ -81,6 +83,12 @@ public UTF8String ofRow(FixedSizeBinaryVector vector, int rowId) { return UTF8String.fromString(UUIDUtil.convert(vector.get(rowId)).toString()); } + @Override + public UTF8String ofRow(IntVector offsetVector, Dictionary dictionary, int rowId) { + byte[] bytes = dictionary.decodeToBinary(offsetVector.get(rowId)).getBytes(); + return UTF8String.fromString(UUIDUtil.convert(bytes).toString()); + } + @Override public UTF8String ofBytes(byte[] bytes) { return UTF8String.fromBytes(bytes); From cc50da4060e23399564f97a44f0ac5748e2eaa8f Mon Sep 17 00:00:00 2001 From: Fokko Date: Wed, 18 Jun 2025 10:13:27 +0200 Subject: [PATCH 2/2] Add another test --- .../vectorized/TestParquetVectorizedReads.java | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/parquet/vectorized/TestParquetVectorizedReads.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/parquet/vectorized/TestParquetVectorizedReads.java index ad8f80caf873..ffc9cf162485 100644 --- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/parquet/vectorized/TestParquetVectorizedReads.java +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/parquet/vectorized/TestParquetVectorizedReads.java @@ -400,4 +400,19 @@ public void testUnsupportedReadsForParquetV2() throws Exception { .hasMessageStartingWith("Cannot support vectorized reads for column") .hasMessageEndingWith("Disable vectorized reads to read this table/file"); } + + @Test + public void testUuidReads() throws Exception { + // Just one row to maintain dictionary encoding + int numRows = 1; + Schema schema = new Schema(optional(100, "uuid", Types.UUIDType.get())); + + File dataFile = File.createTempFile("junit", null, temp.toFile()); + assertThat(dataFile.delete()).as("Delete should succeed").isTrue(); + Iterable data = generateData(schema, numRows, 0L, 0, IDENTITY); + try (FileAppender writer = getParquetV2Writer(schema, dataFile)) { + writer.addAll(data); + } + assertRecordsMatch(schema, numRows, data, dataFile, false, BATCH_SIZE); + } }