This is an automated email from the ASF dual-hosted git repository.
amoghj pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new 73b179c3c1 Spark 3.5, Arrow: Support for Row lineage when using the
Parquet Vectorized reader (#12928)
73b179c3c1 is described below
commit 73b179c3c130e54499d45a9203f63b58cc38e552
Author: Amogh Jahagirdar <[email protected]>
AuthorDate: Thu Jun 5 16:00:36 2025 -0700
Spark 3.5, Arrow: Support for Row lineage when using the Parquet Vectorized
reader (#12928)
---
.../iceberg/arrow/vectorized/VectorHolder.java | 7 +-
.../arrow/vectorized/VectorizedArrowReader.java | 211 ++++++++++++++++++++-
.../arrow/vectorized/VectorizedReaderBuilder.java | 41 ++--
.../TestRowLevelOperationsWithLineage.java | 79 +++++++-
.../apache/iceberg/spark/data/GenericsHelpers.java | 12 +-
.../org/apache/iceberg/spark/data/TestHelpers.java | 27 ---
.../vectorized/TestParquetVectorizedReads.java | 79 +++++++-
7 files changed, 388 insertions(+), 68 deletions(-)
diff --git
a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java
b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java
index 8919c3b6f7..0245b6bba2 100644
--- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java
+++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java
@@ -73,7 +73,7 @@ public class VectorHolder {
icebergField = field;
}
- private VectorHolder(FieldVector vec, Types.NestedField field,
NullabilityHolder nulls) {
+ VectorHolder(FieldVector vec, Types.NestedField field, NullabilityHolder
nulls) {
columnDescriptor = null;
vector = vec;
isDictionaryEncoded = false;
@@ -167,6 +167,11 @@ public class VectorHolder {
}
}
+ public static VectorHolder vectorHolder(
+ FieldVector vector, Types.NestedField icebergField, NullabilityHolder
nulls) {
+ return new VectorHolder(vector, icebergField, nulls);
+ }
+
public static class DeletedVectorHolder extends VectorHolder {
private final int numRows;
diff --git
a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
index c71c8b1df5..b9509f1999 100644
---
a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
+++
b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java
@@ -461,6 +461,53 @@ public class VectorizedArrowReader implements
VectorizedReader<VectorHolder> {
return new PositionVectorReader(true);
}
+ public static VectorizedArrowReader rowIds(Long baseRowId,
VectorizedArrowReader idReader) {
+ if (baseRowId != null) {
+ return new RowIdVectorReader(baseRowId, idReader);
+ } else {
+ return nulls();
+ }
+ }
+
+ public static VectorizedArrowReader lastUpdated(
+ Long baseRowId, Long fileLastUpdated, VectorizedArrowReader seqReader) {
+ if (fileLastUpdated != null && baseRowId != null) {
+ return new LastUpdatedSeqVectorReader(fileLastUpdated, seqReader);
+ } else {
+ return nulls();
+ }
+ }
+
+ public static VectorizedReader<?> replaceWithMetadataReader(
+ Types.NestedField icebergField,
+ VectorizedReader<?> reader,
+ Map<Integer, ?> idToConstant,
+ boolean setArrowValidityVector) {
+ int id = icebergField.fieldId();
+ if (id == MetadataColumns.ROW_ID.fieldId()) {
+ Long baseRowId = (Long) idToConstant.get(id);
+ return rowIds(baseRowId, (VectorizedArrowReader) reader);
+ } else if (id == MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId()) {
+ Long baseRowId = (Long)
idToConstant.get(MetadataColumns.ROW_ID.fieldId());
+ Long fileSeqNumber = (Long) idToConstant.get(id);
+ return VectorizedArrowReader.lastUpdated(
+ baseRowId, fileSeqNumber, (VectorizedArrowReader) reader);
+ } else if (idToConstant.containsKey(id)) {
+ // containsKey is used because the constant may be null
+ return new ConstantVectorReader<>(icebergField, idToConstant.get(id));
+ } else if (id == MetadataColumns.ROW_POSITION.fieldId()) {
+ if (setArrowValidityVector) {
+ return positionsWithSetArrowValidityVector();
+ } else {
+ return VectorizedArrowReader.positions();
+ }
+ } else if (id == MetadataColumns.IS_DELETED.fieldId()) {
+ return new DeletedVectorReader();
+ }
+
+ return reader;
+ }
+
private static final class NullVectorReader extends VectorizedArrowReader {
private static final NullVectorReader INSTANCE = new NullVectorReader();
@@ -530,12 +577,6 @@ public class VectorizedArrowReader implements
VectorizedReader<VectorHolder> {
return vector;
}
- private static NullabilityHolder newNullabilityHolder(int size) {
- NullabilityHolder nullabilityHolder = new NullabilityHolder(size);
- nullabilityHolder.setNotNulls(0, size);
- return nullabilityHolder;
- }
-
@Override
public void setRowGroupInfo(
PageReadStore source, Map<ColumnPath, ColumnChunkMetaData> metadata) {
@@ -567,6 +608,164 @@ public class VectorizedArrowReader implements
VectorizedReader<VectorHolder> {
}
}
+ private static final class RowIdVectorReader extends VectorizedArrowReader {
+ private static final Field ROW_ID_ARROW_FIELD =
ArrowSchemaUtil.convert(MetadataColumns.ROW_ID);
+
+ private final long firstRowId;
+ private final VectorizedReader<VectorHolder> idReader;
+ private final VectorizedReader<VectorHolder> posReader;
+ private NullabilityHolder nulls;
+
+ private RowIdVectorReader(long firstRowId, VectorizedArrowReader idReader)
{
+ this.firstRowId = firstRowId;
+ this.idReader = idReader != null ? idReader : nulls();
+ this.posReader = new PositionVectorReader(true);
+ }
+
+ @Override
+ public VectorHolder read(VectorHolder reuse, int numValsToRead) {
+ FieldVector positions = null;
+ FieldVector ids = null;
+
+ try {
+ positions = posReader.read(null, numValsToRead).vector();
+ VectorHolder idsHolder = idReader.read(null, numValsToRead);
+ ids = idsHolder.vector();
+ ArrowVectorAccessor<?, String, ?, ?> idsAccessor =
+ ids == null ? null :
ArrowVectorAccessors.getVectorAccessor(idsHolder);
+
+ BigIntVector rowIds = allocateBigIntVector(ROW_ID_ARROW_FIELD,
numValsToRead);
+ ArrowBuf dataBuffer = rowIds.getDataBuffer();
+ for (int i = 0; i < numValsToRead; i += 1) {
+ long bufferOffset = (long) i * Long.BYTES;
+ if (idsAccessor == null || isNull(idsHolder, i)) {
+ long rowId = firstRowId + (Long) positions.getObject(i);
+ dataBuffer.setLong(bufferOffset, rowId);
+ } else {
+ long materializedRowId = idsAccessor.getLong(i);
+ dataBuffer.setLong(bufferOffset, materializedRowId);
+ }
+ }
+
+ rowIds.setValueCount(numValsToRead);
+ return VectorHolder.vectorHolder(rowIds, MetadataColumns.ROW_ID,
nulls);
+ } finally {
+ if (positions != null) {
+ positions.close();
+ }
+
+ if (ids != null) {
+ ids.close();
+ }
+ }
+ }
+
+ @Override
+ public void setRowGroupInfo(
+ PageReadStore source, Map<ColumnPath, ColumnChunkMetaData> metadata) {
+ idReader.setRowGroupInfo(source, metadata);
+ posReader.setRowGroupInfo(source, metadata);
+ }
+
+ @Override
+ public void setBatchSize(int batchSize) {
+ if (nulls == null || nulls.size() < batchSize) {
+ this.nulls = newNullabilityHolder(batchSize);
+ }
+
+ idReader.setBatchSize(batchSize);
+ posReader.setBatchSize(batchSize);
+ }
+
+ @Override
+ public void close() {
+ // don't close result vectors as they are not owned by readers
+ }
+ }
+
+ private static final class LastUpdatedSeqVectorReader extends
VectorizedArrowReader {
+ private static final Field LAST_UPDATED_SEQ =
+ ArrowSchemaUtil.convert(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER);
+
+ private final long lastUpdatedSeq;
+ private final VectorizedReader<VectorHolder> seqReader;
+ private NullabilityHolder nulls;
+
+ private LastUpdatedSeqVectorReader(
+ long lastUpdatedSeq, VectorizedReader<VectorHolder> seqReader) {
+ this.lastUpdatedSeq = lastUpdatedSeq;
+ this.seqReader = seqReader == null ? nulls() : seqReader;
+ }
+
+ @Override
+ public VectorHolder read(VectorHolder reuse, int numValsToRead) {
+ FieldVector seqNumbers = null;
+ try {
+ VectorHolder seqNumbersHolder = seqReader.read(null, numValsToRead);
+ seqNumbers = seqNumbersHolder.vector();
+ ArrowVectorAccessor<?, String, ?, ?> seqAccessor =
+ seqNumbers == null ? null :
ArrowVectorAccessors.getVectorAccessor(seqNumbersHolder);
+
+ BigIntVector lastUpdatedSequenceNumbers =
+ allocateBigIntVector(LAST_UPDATED_SEQ, numValsToRead);
+ ArrowBuf dataBuffer = lastUpdatedSequenceNumbers.getDataBuffer();
+ for (int i = 0; i < numValsToRead; i += 1) {
+ long bufferOffset = (long) i * Long.BYTES;
+ if (seqAccessor == null || isNull(seqNumbersHolder, i)) {
+ dataBuffer.setLong(bufferOffset, lastUpdatedSeq);
+ } else {
+ long materializedSeqNumber = seqAccessor.getLong(i);
+ dataBuffer.setLong(bufferOffset, materializedSeqNumber);
+ }
+ }
+
+ lastUpdatedSequenceNumbers.setValueCount(numValsToRead);
+ return VectorHolder.vectorHolder(
+ lastUpdatedSequenceNumbers,
MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER, nulls);
+ } finally {
+ if (seqNumbers != null) {
+ seqNumbers.close();
+ }
+ }
+ }
+
+ @Override
+ public void setRowGroupInfo(
+ PageReadStore source, Map<ColumnPath, ColumnChunkMetaData> metadata) {
+ seqReader.setRowGroupInfo(source, metadata);
+ }
+
+ @Override
+ public void setBatchSize(int batchSize) {
+ if (nulls == null || nulls.size() < batchSize) {
+ this.nulls = newNullabilityHolder(batchSize);
+ }
+
+ seqReader.setBatchSize(batchSize);
+ }
+
+ @Override
+ public void close() {
+ // don't close result vectors as they are not owned by readers
+ }
+ }
+
+ private static boolean isNull(VectorHolder holder, int index) {
+ return holder.nullabilityHolder().isNullAt(index) == 1;
+ }
+
+ private static BigIntVector allocateBigIntVector(Field field, int
valueCount) {
+ BigIntVector vector = (BigIntVector)
field.createVector(ArrowAllocation.rootAllocator());
+ vector.allocateNew(valueCount);
+ return vector;
+ }
+
+ private static NullabilityHolder newNullabilityHolder(int size) {
+ NullabilityHolder nullabilityHolder = new NullabilityHolder(size);
+ nullabilityHolder.setNotNulls(0, size);
+ return nullabilityHolder;
+ }
+
/**
* A Dummy Vector Reader which doesn't actually read files, instead it
returns a dummy
* VectorHolder which indicates the constant value which should be used for
this column.
diff --git
a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java
b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java
index 398f42eb1c..391c356197 100644
---
a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java
+++
b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java
@@ -24,11 +24,9 @@ import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.IntStream;
import org.apache.arrow.memory.BufferAllocator;
-import org.apache.iceberg.MetadataColumns;
import org.apache.iceberg.Schema;
import org.apache.iceberg.arrow.ArrowAllocation;
import
org.apache.iceberg.arrow.vectorized.VectorizedArrowReader.ConstantVectorReader;
-import
org.apache.iceberg.arrow.vectorized.VectorizedArrowReader.DeletedVectorReader;
import org.apache.iceberg.parquet.TypeWithSchemaVisitor;
import org.apache.iceberg.parquet.VectorizedReader;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
@@ -101,33 +99,26 @@ public class VectorizedReaderBuilder extends
TypeWithSchemaVisitor<VectorizedRea
Lists.newArrayListWithExpectedSize(icebergFields.size());
for (Types.NestedField field : icebergFields) {
- int id = field.fieldId();
- VectorizedReader<?> reader = readersById.get(id);
- if (idToConstant.containsKey(id)) {
- reorderedFields.add(constantReader(field, idToConstant.get(id)));
- } else if (id == MetadataColumns.ROW_POSITION.fieldId()) {
- if (setArrowValidityVector) {
-
reorderedFields.add(VectorizedArrowReader.positionsWithSetArrowValidityVector());
- } else {
- reorderedFields.add(VectorizedArrowReader.positions());
- }
- } else if (id == MetadataColumns.IS_DELETED.fieldId()) {
- reorderedFields.add(new DeletedVectorReader());
- } else if (reader != null) {
- reorderedFields.add(reader);
- } else if (field.initialDefault() != null) {
- reorderedFields.add(
- constantReader(field, convert.apply(field.type(),
field.initialDefault())));
- } else if (field.isOptional()) {
- reorderedFields.add(VectorizedArrowReader.nulls());
- } else {
- throw new IllegalArgumentException(
- String.format("Missing required field: %s", field.name()));
- }
+ VectorizedReader<?> reader =
+ VectorizedArrowReader.replaceWithMetadataReader(
+ field, readersById.get(field.fieldId()), idToConstant,
setArrowValidityVector);
+ reorderedFields.add(defaultReader(field, reader));
}
return vectorizedReader(reorderedFields);
}
+ private VectorizedReader<?> defaultReader(Types.NestedField field,
VectorizedReader<?> reader) {
+ if (reader != null) {
+ return reader;
+ } else if (field.initialDefault() != null) {
+ return constantReader(field, convert.apply(field.type(),
field.initialDefault()));
+ } else if (field.isOptional()) {
+ return VectorizedArrowReader.nulls();
+ }
+
+ throw new IllegalArgumentException(String.format("Missing required field:
%s", field.name()));
+ }
+
private <T> ConstantVectorReader<T> constantReader(Types.NestedField field,
T constant) {
return new ConstantVectorReader<>(field, constant);
}
diff --git
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
index 62c7690641..34f1f5d341 100644
---
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
+++
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRowLevelOperationsWithLineage.java
@@ -91,7 +91,6 @@ public abstract class TestRowLevelOperationsWithLineage
extends SparkRowLevelOpe
assumeThat(formatVersion).isGreaterThanOrEqualTo(3);
// ToDo: Remove these as row lineage inheritance gets implemented in the
other readers
assumeThat(fileFormat).isEqualTo(FileFormat.PARQUET);
- assumeThat(vectorized).isFalse();
}
@AfterEach
@@ -396,6 +395,84 @@ public abstract class TestRowLevelOperationsWithLineage
extends SparkRowLevelOpe
rowsWithLineage());
}
+ @TestTemplate
+ public void testMergeWithManyRecords() throws NoSuchTableException,
ParseException, IOException {
+ createAndInitTable("id INT, data STRING", null);
+ createBranchIfNeeded();
+ Table table = loadIcebergTable(spark, tableName);
+
+ int numRecords = 25000;
+ int startingId = 100;
+
+ List<Record> initialRecords = Lists.newArrayList();
+ int rowId = 0;
+ for (int id = 100; id < startingId + numRecords; id++) {
+ initialRecords.add(createRecord(SCHEMA, id, "data_" + id, rowId++, 1L));
+ }
+
+ appendUnpartitionedRecords(table, initialRecords);
+ createOrReplaceView(
+ "source",
+ "id int, data string",
+ "{ \"id\": 101, \"data\": \"updated_data_101\" }\n "
+ + "{ \"id\": 26000, \"data\": \"data_26000\" }\n");
+ sql(
+ "MERGE INTO %s AS t USING source AS s "
+ + "ON t.id == s.id "
+ + "WHEN MATCHED THEN "
+ + " UPDATE SET t.data = s.data "
+ + "WHEN NOT MATCHED THEN "
+ + " INSERT *",
+ commitTarget());
+
+ Snapshot updateSnapshot = latestSnapshot(table);
+ long updateSnapshotFirstRowId = updateSnapshot.firstRowId();
+ List<Object[]> allRows = rowsWithLineageAndFilePos();
+ List<Object[]> carriedOverAndUpdatedRows =
+ allRows.stream()
+ .filter(row -> (long) row[3] < updateSnapshotFirstRowId)
+ .collect(Collectors.toList());
+
+ int newlyInsertedId = 26000;
+ int updatedId = 101;
+ List<Object[]> expectedCarriedOverAndUpdatedRows =
+ ImmutableList.<Object[]>builder()
+ .add(row(1L, 100, "data_100", 0L, ANY, ANY))
+ .add(row(updateSnapshot.sequenceNumber(), updatedId,
"updated_data_101", 1L, ANY, ANY))
+ .addAll(
+ // Every record with higher ids than the updated excluding the
new row should be a
+ // carry over
+ initialRecords.stream()
+ .filter(
+ initialRecord -> {
+ int id = initialRecord.get(0, Integer.class);
+ return id > updatedId && id != newlyInsertedId;
+ })
+ .map(this::recordToExpectedRow)
+ .collect(Collectors.toList()))
+ .build();
+
+ assertEquals(
+ "Rows which are carried over or updated should have expected lineage",
+ expectedCarriedOverAndUpdatedRows,
+ carriedOverAndUpdatedRows);
+
+ Object[] newRow =
+ Iterables.getOnlyElement(
+ allRows.stream()
+ .filter(row -> (long) row[3] >= updateSnapshotFirstRowId)
+ .collect(Collectors.toList()));
+ assertAddedRowLineage(row(updateSnapshot.sequenceNumber(), 26000,
"data_26000"), newRow);
+ }
+
+ private Object[] recordToExpectedRow(Record record) {
+ int id = record.get(0, Integer.class);
+ String data = record.get(1, String.class);
+ long rowId = record.get(2, Long.class);
+ long lastUpdated = record.get(3, Long.class);
+ return row(lastUpdated, id, data, rowId, ANY, ANY);
+ }
+
private List<Object[]> rowsWithLineageAndFilePos() {
return sql(
"SELECT s._last_updated_sequence_number, s.id, s.data, s._row_id,
files.first_row_id, s._pos FROM %s"
diff --git
a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/GenericsHelpers.java
b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/GenericsHelpers.java
index aed9f2c2af..e03b827cbe 100644
---
a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/GenericsHelpers.java
+++
b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/GenericsHelpers.java
@@ -73,11 +73,15 @@ public class GenericsHelpers {
}
public static void assertEqualsBatch(
- Types.StructType struct, Iterator<Record> expectedRecords, ColumnarBatch
batch) {
- for (int rowId = 0; rowId < batch.numRows(); rowId++) {
- InternalRow row = batch.getRow(rowId);
+ Types.StructType struct,
+ Iterator<Record> expectedRecords,
+ ColumnarBatch batch,
+ Map<Integer, Object> idToConstant,
+ Integer batchFirstRowPos) {
+ for (int rowPos = 0; rowPos < batch.numRows(); rowPos++) {
+ InternalRow row = batch.getRow(rowPos);
Record expectedRecord = expectedRecords.next();
- assertEqualsUnsafe(struct, expectedRecord, row);
+ assertEqualsUnsafe(struct, expectedRecord, row, idToConstant,
batchFirstRowPos + rowPos);
}
}
diff --git
a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestHelpers.java
b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestHelpers.java
index 7baad225eb..82ebb1d950 100644
---
a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestHelpers.java
+++
b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/TestHelpers.java
@@ -113,33 +113,6 @@ public class TestHelpers {
}
}
- public static void assertEqualsBatch(
- Types.StructType struct, Iterator<Record> expected, ColumnarBatch batch)
{
- for (int rowId = 0; rowId < batch.numRows(); rowId++) {
- InternalRow row = batch.getRow(rowId);
- Record rec = expected.next();
-
- List<Types.NestedField> fields = struct.fields();
- for (int readPos = 0; readPos < fields.size(); readPos += 1) {
- Types.NestedField field = fields.get(readPos);
- Field writeField = rec.getSchema().getField(field.name());
-
- Type fieldType = field.type();
- Object actualValue = row.isNullAt(readPos) ? null : row.get(readPos,
convert(fieldType));
-
- Object expectedValue;
- if (writeField != null) {
- int writePos = writeField.pos();
- expectedValue = rec.get(writePos);
- } else {
- expectedValue = field.initialDefault();
- }
-
- assertEqualsUnsafe(fieldType, expectedValue, actualValue);
- }
- }
- }
-
public static void assertEqualsBatchWithRows(
Types.StructType struct, Iterator<Row> expected, ColumnarBatch batch) {
for (int rowId = 0; rowId < batch.numRows(); rowId++) {
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 9d443de748..ad8f80caf8 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
@@ -27,6 +27,8 @@ import static org.assertj.core.api.Assumptions.assumeThat;
import java.io.File;
import java.io.IOException;
import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
import org.apache.iceberg.Files;
import org.apache.iceberg.Schema;
import org.apache.iceberg.data.RandomGenericData;
@@ -37,6 +39,7 @@ import org.apache.iceberg.io.FileAppender;
import org.apache.iceberg.parquet.Parquet;
import org.apache.iceberg.relocated.com.google.common.base.Function;
import org.apache.iceberg.relocated.com.google.common.base.Strings;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
@@ -77,6 +80,13 @@ public class TestParquetVectorizedReads extends AvroDataTest
{
IDENTITY);
}
+ @Override
+ protected void writeAndValidate(Schema writeSchema, Schema expectedSchema,
List<Record> records)
+ throws IOException {
+ writeAndValidate(
+ writeSchema, expectedSchema, true, BATCH_SIZE, records.size(),
records, ID_TO_CONSTANT);
+ }
+
@Override
protected boolean supportsDefaultValues() {
return true;
@@ -87,6 +97,11 @@ public class TestParquetVectorizedReads extends AvroDataTest
{
return false;
}
+ @Override
+ protected boolean supportsRowLineage() {
+ return true;
+ }
+
private void writeAndValidate(
Schema schema, int numRecords, long seed, float nullPercentage, boolean
reuseContainers)
throws IOException {
@@ -104,6 +119,29 @@ public class TestParquetVectorizedReads extends
AvroDataTest {
int batchSize,
Function<Record, Record> transform)
throws IOException {
+ writeAndValidate(
+ writeSchema,
+ expectedSchema,
+ numRecords,
+ seed,
+ nullPercentage,
+ reuseContainers,
+ batchSize,
+ transform,
+ ImmutableMap.of());
+ }
+
+ private void writeAndValidate(
+ Schema writeSchema,
+ Schema expectedSchema,
+ int numRecords,
+ long seed,
+ float nullPercentage,
+ boolean reuseContainers,
+ int batchSize,
+ Function<Record, Record> transform,
+ Map<Integer, Object> idToConstant)
+ throws IOException {
// Write test data
assumeThat(
TypeUtil.find(
@@ -115,6 +153,25 @@ public class TestParquetVectorizedReads extends
AvroDataTest {
Iterable<Record> expected =
generateData(writeSchema, numRecords, seed, nullPercentage, transform);
+ writeAndValidate(
+ writeSchema,
+ expectedSchema,
+ reuseContainers,
+ batchSize,
+ numRecords,
+ expected,
+ idToConstant);
+ }
+
+ private void writeAndValidate(
+ Schema writeSchema,
+ Schema expectedSchema,
+ boolean reuseContainers,
+ int batchSize,
+ int numRecords,
+ Iterable<Record> expected,
+ Map<Integer, Object> idToConstant)
+ throws IOException {
// write a test parquet file using iceberg writer
File testFile = File.createTempFile("junit", null, temp.toFile());
assertThat(testFile.delete()).as("Delete should succeed").isTrue();
@@ -123,7 +180,8 @@ public class TestParquetVectorizedReads extends
AvroDataTest {
writer.addAll(expected);
}
- assertRecordsMatch(expectedSchema, numRecords, expected, testFile,
reuseContainers, batchSize);
+ assertRecordsMatch(
+ expectedSchema, numRecords, expected, testFile, reuseContainers,
batchSize, idToConstant);
}
protected int getNumRows() {
@@ -165,14 +223,26 @@ public class TestParquetVectorizedReads extends
AvroDataTest {
boolean reuseContainers,
int batchSize)
throws IOException {
+ assertRecordsMatch(
+ schema, expectedSize, expected, testFile, reuseContainers, batchSize,
ImmutableMap.of());
+ }
+
+ void assertRecordsMatch(
+ Schema schema,
+ int expectedSize,
+ Iterable<Record> expected,
+ File testFile,
+ boolean reuseContainers,
+ int batchSize,
+ Map<Integer, Object> idToConstant)
+ throws IOException {
Parquet.ReadBuilder readBuilder =
Parquet.read(Files.localInput(testFile))
.project(schema)
.recordsPerBatch(batchSize)
.createBatchedReaderFunc(
type ->
- VectorizedSparkParquetReaders.buildReader(
- schema, type, Maps.newHashMap(), null));
+ VectorizedSparkParquetReaders.buildReader(schema, type,
idToConstant, null));
if (reuseContainers) {
readBuilder.reuseContainers();
}
@@ -182,8 +252,9 @@ public class TestParquetVectorizedReads extends
AvroDataTest {
int numRowsRead = 0;
while (batches.hasNext()) {
ColumnarBatch batch = batches.next();
+ GenericsHelpers.assertEqualsBatch(
+ schema.asStruct(), expectedIter, batch, idToConstant, numRowsRead);
numRowsRead += batch.numRows();
- GenericsHelpers.assertEqualsBatch(schema.asStruct(), expectedIter,
batch);
}
assertThat(numRowsRead).isEqualTo(expectedSize);
}
