From 1e95954895c4620ae15fb589516e51349ed1ab4a Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Thu, 27 Aug 2026 23:58:09 +0530 Subject: [PATCH 1/9] Spark: Add format model TCK coverage for vectorized reads Co-authored-by: Joy Haldar --- .../apache/iceberg/data/DataGenerator.java | 2 +- .../apache/iceberg/data/DataGenerators.java | 4 +- .../iceberg/data/ReadFormatModelTests.java | 50 ++++---- .../data/TestSparkVectorizedFormatModel.java | 110 ++++++++++++++++++ 4 files changed, 140 insertions(+), 26 deletions(-) create mode 100644 spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java diff --git a/data/src/test/java/org/apache/iceberg/data/DataGenerator.java b/data/src/test/java/org/apache/iceberg/data/DataGenerator.java index ad068fbc686c..7c969f44e2cd 100644 --- a/data/src/test/java/org/apache/iceberg/data/DataGenerator.java +++ b/data/src/test/java/org/apache/iceberg/data/DataGenerator.java @@ -22,7 +22,7 @@ import org.apache.iceberg.Schema; /** Interface for generating test data with different schema types. */ -interface DataGenerator { +public interface DataGenerator { /** Returns the Iceberg schema for this test data. */ Schema schema(); diff --git a/data/src/test/java/org/apache/iceberg/data/DataGenerators.java b/data/src/test/java/org/apache/iceberg/data/DataGenerators.java index 8e423ecabbac..f42f813c8ad3 100644 --- a/data/src/test/java/org/apache/iceberg/data/DataGenerators.java +++ b/data/src/test/java/org/apache/iceberg/data/DataGenerators.java @@ -32,9 +32,9 @@ * Test data generators for different schema types. Add new generators to ALL array to include them * in format model tests. */ -class DataGenerators { +public class DataGenerators { - static final DataGenerator[] ALL = + public static final DataGenerator[] ALL = new DataGenerator[] { new Primitives(), new UUID(), diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index 34712be7e704..bd772b1fa91f 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -117,6 +117,11 @@ protected boolean supportsBatchReads() { return false; } + protected void assertRecordsEqual(Schema schema, List expected, List actual) { + assertThat(actual).hasSize(expected.size()); + assertEquals(schema, convertToEngineRecords(expected, schema), actual); + } + @TempDir private File tableDir; /** @@ -259,9 +264,8 @@ void testDataWriterGenericWriteEngineRead(FileFormat fileFormat, DataGenerator d .project(schema) .build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(schema, genericRecords, readRecords); } - - assertEquals(schema, convertToEngineRecords(genericRecords, schema), readRecords); } /** Write with Generic Record, read with engine type T */ @@ -301,9 +305,8 @@ void testEqualityDeleteWriterGenericWriteEngineRead( .project(schema) .build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(schema, genericRecords, readRecords); } - - assertEquals(schema, convertToEngineRecords(genericRecords, schema), readRecords); } /** Write position deletes, read with Generic Record */ @@ -362,8 +365,6 @@ void testReaderBuilderProjection(FileFormat fileFormat) throws IOException { writeGenericRecords(fileFormat, fullSchema, genericRecords); List projectedGenericRecords = project(genericRecords, projectedSchema); - List expectedEngineRecords = - convertToEngineRecords(projectedGenericRecords, projectedSchema); InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); List readRecords; @@ -372,9 +373,8 @@ void testReaderBuilderProjection(FileFormat fileFormat) throws IOException { .project(projectedSchema) .build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(projectedSchema, projectedGenericRecords, readRecords); } - - assertEquals(projectedSchema, expectedEngineRecords, readRecords); } @ParameterizedTest @@ -488,6 +488,7 @@ void testReaderBuilderCaseSensitive(FileFormat fileFormat) throws IOException { @FieldSource("FILE_FORMATS") void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { + assumeFalse(supportsBatchReads(), "Batch reads return batches, not rows"); assumeSupports(fileFormat, FEATURE_SPLIT); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -543,7 +544,7 @@ void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReaderBuilderReuseContainers(FileFormat fileFormat) throws IOException { - + assumeFalse(supportsBatchReads(), "Batch reads reuse differently"); assumeSupports(fileFormat, FEATURE_REUSE_CONTAINERS); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -678,6 +679,7 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testNestedDefaultValue(FileFormat fileFormat) throws IOException { + assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -744,6 +746,7 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testMapNestedDefaultValue(FileFormat fileFormat) throws IOException { + assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -831,6 +834,7 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testListNestedDefaultValue(FileFormat fileFormat) throws IOException { + assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -1798,7 +1802,7 @@ void testReadMetadataColumnRowPosition(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { - + assumeFalse(supportsBatchReads(), "DeletedColumnVector requires a delete filter"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema schema = dataGenerator.schema(); List genericRecords = dataGenerator.generateRecords(); @@ -1818,7 +1822,10 @@ void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { - + // TODO: ORC's vectorized reader treats _row_id as a plain constant + // (VectorizedSparkOrcReaders.StructConverter) rather than applying row-lineage rules, + // unlike Arrow's RowIdVectorReader. + assumeFalse(supportsBatchReads(), "Row lineage is incorrect for vectorized ORC reads"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema schema = dataGenerator.schema(); List genericRecords = dataGenerator.generateRecords(); @@ -1851,7 +1858,10 @@ void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) throws IOException { - + // TODO: ORC's vectorized reader treats _row_id as a plain constant + // (VectorizedSparkOrcReaders.StructConverter) rather than applying row-lineage rules, + // unlike Arrow's RowIdVectorReader. + assumeFalse(supportsBatchReads(), "Row lineage is incorrect for vectorized ORC reads"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema dataSchema = dataGenerator.schema(); @@ -2311,6 +2321,7 @@ void testSchemaEvolutionRequiredToOptional(FileFormat fileFormat) throws IOExcep @ParameterizedTest @FieldSource("FILE_FORMATS") void testSchemaEvolutionEmptyProjection(FileFormat fileFormat) throws IOException { + assumeFalse(supportsBatchReads(), "Row count assertion does not apply to batches"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema writeSchema = dataGenerator.schema(); @@ -2354,9 +2365,8 @@ void testReadFileWithoutFieldIdsUsingNameMapping(FileFormat fileFormat) throws I .withNameMapping(nameMapping) .build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(icebergSchema, genericRecords, readRecords); } - - assertEquals(icebergSchema, convertToEngineRecords(genericRecords, icebergSchema), readRecords); } protected void readAndAssertGenericRecords( @@ -2411,7 +2421,7 @@ protected static void assumeSupports(FileFormat fileFormat, String feature) { assumeThat(MISSING_FEATURES.getOrDefault(fileFormat, new String[] {})).doesNotContain(feature); } - private static boolean supportsGenerator(FileFormat fileFormat, DataGenerator generator) { + protected static boolean supportsGenerator(FileFormat fileFormat, DataGenerator generator) { boolean hasVariant = TypeUtil.find(generator.schema(), type -> type.typeId() == Type.TypeID.VARIANT) != null; return !hasVariant || supportsFeature(fileFormat, FEATURE_VARIANT); @@ -2816,11 +2826,8 @@ private void readAndAssertMetadataColumn( try (CloseableIterable reader = readerBuilder.build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(projectionSchema, expectedRecords, readRecords); } - - assertThat(readRecords).hasSize(expectedRecords.size()); - assertEquals( - projectionSchema, convertToEngineRecords(expectedRecords, projectionSchema), readRecords); } private static Record copy(Record source, Schema sourceSchema, Schema targetSchema) { @@ -2867,10 +2874,7 @@ protected void readAndAssertEngineRecords( .project(readSchema) .build()) { readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(readSchema, expectedGenericRecords, readRecords); } - - assertThat(readRecords).hasSize(expectedGenericRecords.size()); - assertEquals( - readSchema, convertToEngineRecords(expectedGenericRecords, readSchema), readRecords); } } diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java new file mode 100644 index 000000000000..8a2bacceb9f0 --- /dev/null +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.spark.data; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import java.util.Set; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Schema; +import org.apache.iceberg.data.DataGenerators; +import org.apache.iceberg.data.ReadFormatModelTests; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.types.Type; +import org.apache.spark.sql.vectorized.ColumnarBatch; +import org.junit.jupiter.params.provider.Arguments; + +class TestSparkVectorizedFormatModel extends ReadFormatModelTests { + + private static final FileFormat[] FILE_FORMATS = {FileFormat.ORC, FileFormat.PARQUET}; + + private static final List FORMAT_AND_GENERATOR = + Arrays.stream(FILE_FORMATS) + .flatMap( + format -> + Arrays.stream(DataGenerators.ALL) + .filter(generator -> supportsGenerator(format, generator)) + .map(generator -> Arguments.of(format, generator))) + .toList(); + + private static final Set UNSUPPORTED_TYPE_IDS = + Set.of( + Type.TypeID.TIME, + Type.TypeID.TIMESTAMP_NANO, + Type.TypeID.STRUCT, + Type.TypeID.LIST, + Type.TypeID.MAP, + Type.TypeID.UUID); + + @Override + protected Collection unsupportedTypeIds() { + return UNSUPPORTED_TYPE_IDS; + } + + @Override + protected boolean supportsBatchReads() { + return true; + } + + @Override + protected Object engineSchema(Schema schema) { + return SparkSchemaUtil.convert(schema); + } + + @Override + protected Object convertConstantToEngine(Type type, Object value) { + return SparkUtil.internalToSpark(type, value); + } + + @Override + protected Class engineType() { + return ColumnarBatch.class; + } + + @Override + protected ColumnarBatch convertToEngine(Record record, Schema schema) { + throw new UnsupportedOperationException("Vectorized reads are read-only"); + } + + @Override + protected void assertEquals( + Schema schema, List expected, List actual) { + throw new UnsupportedOperationException("Use assertRecordsEqual for batch reads"); + } + + @Override + protected void assertRecordsEqual( + Schema schema, List expected, List actual) { + int rowId = 0; + for (ColumnarBatch batch : actual) { + for (int i = 0; i < batch.numRows(); i++) { + TestHelpers.assertEquals( + schema, InternalRowConverter.convert(schema, expected.get(rowId)), batch.getRow(i)); + rowId++; + } + } + + assertThat(rowId).isEqualTo(expected.size()); + } +} From e8b97bac09d5bbe049ced2ecda64f182c1800360 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Fri, 16 Oct 2026 22:25:13 +0530 Subject: [PATCH 2/9] Spark 4.2: skip nested type tests via unsupportedTypeIds, move test to 4.2 Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 26 ++++++++++++++++--- .../data/TestSparkVectorizedFormatModel.java | 0 2 files changed, 23 insertions(+), 3 deletions(-) rename spark/{v4.1 => v4.2}/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java (100%) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index bd772b1fa91f..a8ba5702089d 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -148,6 +148,14 @@ private Schema filterUnsupported(Schema schema) { return kept.size() == schema.columns().size() ? schema : new Schema(kept); } + protected void assumeSupportedTypes(Schema schema) { + Collection unsupportedTypeIds = unsupportedTypeIds(); + + Type unsupportedType = + TypeUtil.find(schema, type -> unsupportedTypeIds.contains(type.typeId())); + assumeThat(unsupportedType).as("Schema contains a type unsupported by this engine").isNull(); + } + protected Schema supportedSchema(DataGenerator generator) { Schema filtered = filterUnsupported(generator.schema()); assumeThat(filtered.columns()) @@ -679,7 +687,6 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testNestedDefaultValue(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -703,6 +710,7 @@ void testNestedDefaultValue(FileFormat fileFormat) throws IOException { .withDoc("Used to test nested field defaults") .build()); + assumeSupportedTypes(writeSchema); List genericRecords = RandomGenericData.generate(writeSchema, 10, 1L); writeGenericRecords(fileFormat, writeSchema, genericRecords); @@ -746,7 +754,6 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testMapNestedDefaultValue(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -774,6 +781,7 @@ void testMapNestedDefaultValue(FileFormat fileFormat) throws IOException { .withDoc("Used to test nested map value field defaults") .build()); + assumeSupportedTypes(writeSchema); List genericRecords = RandomGenericData.generate(writeSchema, 10, 1L); writeGenericRecords(fileFormat, writeSchema, genericRecords); @@ -834,7 +842,6 @@ record -> { @ParameterizedTest @FieldSource("FILE_FORMATS") void testListNestedDefaultValue(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "Vectorized reads do not support nested types"); assumeSupports(fileFormat, FEATURE_READER_DEFAULT); Types.NestedField idField = Types.NestedField.required(1, "id", Types.LongType.get()); @@ -860,6 +867,7 @@ void testListNestedDefaultValue(FileFormat fileFormat) throws IOException { .withDoc("Used to test nested field defaults") .build()); + assumeSupportedTypes(writeSchema); List genericRecords = RandomGenericData.generate(writeSchema, 10, 1L); writeGenericRecords(fileFormat, writeSchema, genericRecords); @@ -934,6 +942,8 @@ void testNestedDefaultValueWhenParentStructIsNull(FileFormat fileFormat) throws Types.NestedField.optional("nested").withId(2).ofType(writeNested).build(); Schema writeSchema = new Schema(idField, nestedField); + assumeSupportedTypes(writeSchema); + Record presentNested = GenericRecord.create(writeNested); presentNested.setField("a", "x"); presentNested.setField("b", "y"); @@ -997,6 +1007,7 @@ void testNestedDefaultValueWhenParentStructWithListIsNull(FileFormat fileFormat) Types.NestedField.required(5, "inner", Types.StringType.get()))) .build(); Schema writeSchema = new Schema(idField, nestedField); + assumeSupportedTypes(writeSchema); Types.StructType nestedType = nestedField.type().asStructType(); // alternate present and null structs, with a different number of list elements per row so @@ -1072,6 +1083,7 @@ void testNestedDefaultValueWhenParentStructWithOnlyListIsNull(FileFormat fileFor 3, "tags", Types.ListType.ofRequired(4, Types.StringType.get())))) .build(); Schema writeSchema = new Schema(idField, nestedField); + assumeSupportedTypes(writeSchema); Types.StructType nestedType = nestedField.type().asStructType(); List genericRecords = Lists.newArrayList(); @@ -1109,6 +1121,8 @@ void testNestedProjectionWithoutDefaultWhenParentStructIsNull(FileFormat fileFor Types.StructType.of(Types.NestedField.required(3, "inner", Types.StringType.get()))); Schema writeSchema = new Schema(idField, nestedField); + assumeSupportedTypes(writeSchema); + Record present = GenericRecord.create(writeSchema); present.setField("id", 1L); Record presentNested = GenericRecord.create(nestedField.type().asStructType()); @@ -1166,6 +1180,7 @@ void testArrayElementDefault(FileFormat fileFormat) throws IOException { .ofType(Types.ListType.ofOptional(3, writeElement)) .build()); + assumeSupportedTypes(writeSchema); Record present = GenericRecord.create(writeElement); present.setField("inner", "a"); Record row = GenericRecord.create(writeSchema); @@ -1221,6 +1236,7 @@ void testArrayNullAndEmpty(FileFormat fileFormat) throws IOException { .ofType(Types.ListType.ofOptional(3, writeElement)) .build()); + assumeSupportedTypes(writeSchema); Record nullList = GenericRecord.create(writeSchema); nullList.setField("id", 1L); nullList.setField("arr", null); @@ -1275,6 +1291,8 @@ void testMapValueDefault(FileFormat fileFormat) throws IOException { .ofType(Types.MapType.ofOptional(3, 4, Types.StringType.get(), writeValue)) .build()); + assumeSupportedTypes(writeSchema); + Record present = GenericRecord.create(writeValue); present.setField("inner", "a"); Map writtenMap = Maps.newLinkedHashMap(); @@ -1336,6 +1354,8 @@ void testMapNullAndEmpty(FileFormat fileFormat) throws IOException { .ofType(Types.MapType.ofOptional(3, 4, Types.StringType.get(), writeValue)) .build()); + assumeSupportedTypes(writeSchema); + Record nullMap = GenericRecord.create(writeSchema); nullMap.setField("id", 1L); nullMap.setField("map", null); diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java similarity index 100% rename from spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java rename to spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java From 2240f754d1715af6813498bcea6f92e2b7a66875 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Fri, 16 Oct 2026 22:43:20 +0530 Subject: [PATCH 3/9] Count rows via countRows hook so split and empty projection tests run for batch readers Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 18 +++++++++++------- .../data/TestSparkVectorizedFormatModel.java | 5 +++++ 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index a8ba5702089d..ded0176f16ff 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -156,6 +156,13 @@ protected void assumeSupportedTypes(Schema schema) { assumeThat(unsupportedType).as("Schema contains a type unsupported by this engine").isNull(); } + /** + * Returns the number of rows in the read result; batch engines override to sum across batches. + */ + protected int countRows(List records) { + return records.size(); + } + protected Schema supportedSchema(DataGenerator generator) { Schema filtered = filterUnsupported(generator.schema()); assumeThat(filtered.columns()) @@ -495,8 +502,6 @@ void testReaderBuilderCaseSensitive(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { - - assumeFalse(supportsBatchReads(), "Batch reads return batches, not rows"); assumeSupports(fileFormat, FEATURE_SPLIT); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -522,7 +527,7 @@ void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { readRecords = ImmutableList.copyOf(reader); } - assertThat(readRecords).hasSizeGreaterThan(0).hasSizeLessThan(records.size()); + assertThat(countRows(readRecords)).isGreaterThan(0).isLessThan(records.size()); // split(fileLength, 0): empty range at the end of the file, so no records should be returned List emptyReadRecords; @@ -534,7 +539,7 @@ void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { emptyReadRecords = ImmutableList.copyOf(reader); } - assertThat(emptyReadRecords).isEmpty(); + assertThat(countRows(emptyReadRecords)).isZero(); // split(0, fileLength): full file range, so all records should be returned try (CloseableIterable reader = @@ -545,7 +550,7 @@ void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { readRecords = ImmutableList.copyOf(reader); } - assertThat(readRecords).hasSize(records.size()); + assertThat(countRows(readRecords)).isEqualTo(records.size()); } /** Verifies the contract of reuseContainers */ @@ -2341,7 +2346,6 @@ void testSchemaEvolutionRequiredToOptional(FileFormat fileFormat) throws IOExcep @ParameterizedTest @FieldSource("FILE_FORMATS") void testSchemaEvolutionEmptyProjection(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "Row count assertion does not apply to batches"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema writeSchema = dataGenerator.schema(); @@ -2359,7 +2363,7 @@ void testSchemaEvolutionEmptyProjection(FileFormat fileFormat) throws IOExceptio readRecords = ImmutableList.copyOf(reader); } - assertThat(readRecords).hasSameSizeAs(genericRecords); + assertThat(countRows(readRecords)).isEqualTo(genericRecords.size()); } @ParameterizedTest diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java index 8a2bacceb9f0..397f4937bbe6 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -67,6 +67,11 @@ protected boolean supportsBatchReads() { return true; } + @Override + protected int countRows(List batches) { + return batches.stream().mapToInt(ColumnarBatch::numRows).sum(); + } + @Override protected Object engineSchema(Schema schema) { return SparkSchemaUtil.convert(schema); From d2cae98ab77b7d8ffdf4995db9ca3f8e7eedd852 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Fri, 16 Oct 2026 23:19:12 +0530 Subject: [PATCH 4/9] Test recordsPerBatch for both batch and row readers Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index ded0176f16ff..5b53d2150380 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -1471,21 +1471,36 @@ void testPrimitiveDefaultValuesNotApplied(FileFormat fileFormat) throws IOExcept @ParameterizedTest @FieldSource("FILE_FORMATS") - void testReaderBuilderRecordsPerBatchNotSupported(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), engineType().getSimpleName() + " supports batch reads"); - + void testReaderBuilderRecordsPerBatch(FileFormat fileFormat) throws IOException { DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema schema = dataGenerator.schema(); - List genericRecords = dataGenerator.generateRecords(); + List genericRecords = dataGenerator.generateRecords(1000); writeGenericRecords(fileFormat, schema, genericRecords); InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - assertThatThrownBy( - () -> - FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) - .recordsPerBatch(100)) - .hasMessageContaining("Batch reading is not supported") - .isInstanceOf(UnsupportedOperationException.class); + if (!supportsBatchReads()) { + assertThatThrownBy( + () -> + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) + .recordsPerBatch(100)) + .hasMessageContaining("Batch reading is not supported") + .isInstanceOf(UnsupportedOperationException.class); + return; + } + + List readRecords; + try (CloseableIterable reader = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) + .project(schema) + .recordsPerBatch(100) + .build()) { + readRecords = ImmutableList.copyOf(reader); + } + + for (T batch : readRecords) { + assertThat(countRows(List.of(batch))).isLessThanOrEqualTo(100); + } + assertThat(countRows(readRecords)).isEqualTo(genericRecords.size()); } @ParameterizedTest From d01b937f2ebaaa47a6e80603df62e259e8fc3810 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Fri, 16 Oct 2026 23:30:19 +0530 Subject: [PATCH 5/9] Skip row lineage tests only for vectorized ORC, reference #18543 Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index 5b53d2150380..5c8388f94030 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -1862,10 +1862,10 @@ void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { - // TODO: ORC's vectorized reader treats _row_id as a plain constant - // (VectorizedSparkOrcReaders.StructConverter) rather than applying row-lineage rules, - // unlike Arrow's RowIdVectorReader. - assumeFalse(supportsBatchReads(), "Row lineage is incorrect for vectorized ORC reads"); + // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed + assumeFalse( + supportsBatchReads() && fileFormat == FileFormat.ORC, + "Row lineage is not applied by the vectorized ORC reader, see #18543"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema schema = dataGenerator.schema(); List genericRecords = dataGenerator.generateRecords(); @@ -1898,10 +1898,10 @@ void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) throws IOException { - // TODO: ORC's vectorized reader treats _row_id as a plain constant - // (VectorizedSparkOrcReaders.StructConverter) rather than applying row-lineage rules, - // unlike Arrow's RowIdVectorReader. - assumeFalse(supportsBatchReads(), "Row lineage is incorrect for vectorized ORC reads"); + // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed + assumeFalse( + supportsBatchReads() && fileFormat == FileFormat.ORC, + "Row lineage is not applied by the vectorized ORC reader, see #18543"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema dataSchema = dataGenerator.schema(); From 556ce35632490a4875590e3649979c4e02870ac6 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Fri, 16 Oct 2026 23:39:41 +0530 Subject: [PATCH 6/9] Reword _deleted skip message and limit it to vectorized Parquet Co-authored-by: Joy Haldar --- .../java/org/apache/iceberg/data/ReadFormatModelTests.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index 5c8388f94030..330cf3feea18 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -1842,7 +1842,9 @@ void testReadMetadataColumnRowPosition(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "DeletedColumnVector requires a delete filter"); + assumeFalse( + supportsBatchReads() && fileFormat == FileFormat.PARQUET, + "Reading _deleted without a delete filter is not supported by the vectorized reader"); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); Schema schema = dataGenerator.schema(); List genericRecords = dataGenerator.generateRecords(); From 0f643c1de4ff35b384bf8a644a6c942c118d3758 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Sun, 18 Oct 2026 09:21:01 +0530 Subject: [PATCH 7/9] Test container reuse for Spark vectorized Parquet reads Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 50 ++++++++++++------- .../data/TestSparkVectorizedFormatModel.java | 25 ++++++++++ 2 files changed, 58 insertions(+), 17 deletions(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index 330cf3feea18..8a5549daec10 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -163,6 +163,14 @@ protected int countRows(List records) { return records.size(); } + protected void assertContainersNotReused(T first, T second) { + assertThat(second).isNotSameAs(first); + } + + protected void assertContainersReused(T first, T second) { + assertThat(second).isSameAs(first); + } + protected Schema supportedSchema(DataGenerator generator) { Schema filtered = filterUnsupported(generator.schema()); assumeThat(filtered.columns()) @@ -557,7 +565,6 @@ void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReaderBuilderReuseContainers(FileFormat fileFormat) throws IOException { - assumeFalse(supportsBatchReads(), "Batch reads reuse differently"); assumeSupports(fileFormat, FEATURE_REUSE_CONTAINERS); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -570,31 +577,40 @@ void testReaderBuilderReuseContainers(FileFormat fileFormat) throws IOException InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); // Without reuseContainers: every record must be a distinct object instance - List noReuseRecords; - try (CloseableIterable reader = - FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) - .project(schema) - .build()) { - noReuseRecords = ImmutableList.copyOf(reader); + ReadBuilder noReuseBuilder = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile).project(schema); + + if (supportsBatchReads()) { + noReuseBuilder.recordsPerBatch(2); } - for (int i = 0; i < noReuseRecords.size(); i++) { - for (int j = i + 1; j < noReuseRecords.size(); j++) { - assertThat(noReuseRecords.get(i)).isNotSameAs(noReuseRecords.get(j)); + try (CloseableIterable reader = noReuseBuilder.build()) { + List noReuseRecords = ImmutableList.copyOf(reader); + + assertThat(noReuseRecords).hasSizeGreaterThan(1); + for (int i = 0; i < noReuseRecords.size(); i++) { + for (int j = i + 1; j < noReuseRecords.size(); j++) { + assertContainersNotReused(noReuseRecords.get(i), noReuseRecords.get(j)); + } } } - // With reuseContainers: all collected elements must be the same object instance - List reuseRecords; - try (CloseableIterable reader = + // With reuseContainers: verify that the underlying containers are reused + ReadBuilder builder = FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) .project(schema) - .reuseContainers() - .build()) { - reuseRecords = ImmutableList.copyOf(reader); + .reuseContainers(); + + if (supportsBatchReads()) { + builder.recordsPerBatch(2); } - reuseRecords.forEach(r -> assertThat(r).isSameAs(reuseRecords.get(0))); + try (CloseableIterable reader = builder.build()) { + List reuseRecords = ImmutableList.copyOf(reader); + + assertThat(reuseRecords).hasSizeGreaterThan(1); + reuseRecords.forEach(r -> assertContainersReused(reuseRecords.get(0), r)); + } } @ParameterizedTest diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java index 397f4937bbe6..0a9132d41320 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -24,6 +24,7 @@ import java.util.Collection; import java.util.List; import java.util.Set; +import org.apache.arrow.vector.ValueVector; import org.apache.iceberg.FileFormat; import org.apache.iceberg.Schema; import org.apache.iceberg.data.DataGenerators; @@ -31,6 +32,7 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.spark.data.vectorized.IcebergArrowColumnVector; import org.apache.iceberg.types.Type; import org.apache.spark.sql.vectorized.ColumnarBatch; import org.junit.jupiter.params.provider.Arguments; @@ -72,6 +74,29 @@ protected int countRows(List batches) { return batches.stream().mapToInt(ColumnarBatch::numRows).sum(); } + @Override + protected void assertContainersReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isSameAs(arrowVector(first, i)); + } + } + + private ValueVector arrowVector(ColumnarBatch batch, int col) { + IcebergArrowColumnVector vector = (IcebergArrowColumnVector) batch.column(col); + return vector.vectorAccessor().getVector(); + } + + @Override + protected void assertContainersNotReused(ColumnarBatch first, ColumnarBatch second) { + assertThat(second.numCols()).isEqualTo(first.numCols()); + + for (int i = 0; i < first.numCols(); i++) { + assertThat(arrowVector(second, i)).isNotSameAs(arrowVector(first, i)); + } + } + @Override protected Object engineSchema(Schema schema) { return SparkSchemaUtil.convert(schema); From 04bc92813de357f973668cb5f403027450cd4169 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Mon, 19 Oct 2026 15:49:29 +0530 Subject: [PATCH 8/9] Skip unsupported geospatial types in Spark vectorized format tests Co-authored-by: Joy Haldar --- .../iceberg/spark/data/TestSparkVectorizedFormatModel.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java index 0a9132d41320..dc9f86d67827 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -57,7 +57,9 @@ class TestSparkVectorizedFormatModel extends ReadFormatModelTests Type.TypeID.STRUCT, Type.TypeID.LIST, Type.TypeID.MAP, - Type.TypeID.UUID); + Type.TypeID.UUID, + Type.TypeID.GEOMETRY, + Type.TypeID.GEOGRAPHY); @Override protected Collection unsupportedTypeIds() { From 43466103a877e7d232faf2303af6ddc472c05358 Mon Sep 17 00:00:00 2001 From: Joy Haldar Date: Wed, 21 Oct 2026 11:58:48 +0530 Subject: [PATCH 9/9] Address review feedback for Spark vectorized format tests Co-authored-by: Joy Haldar --- .../iceberg/data/ReadFormatModelTests.java | 79 +++++++++++-------- .../data/TestSparkVectorizedFormatModel.java | 18 +++-- 2 files changed, 60 insertions(+), 37 deletions(-) diff --git a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java index 7c00e421bdca..037426f077ec 100644 --- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java +++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java @@ -119,17 +119,30 @@ protected boolean supportsBatchReads() { return false; } - protected void assertRecordsEqual(Schema schema, List expected, List actual) { - assertThat(actual).hasSize(expected.size()); - assertEquals(schema, convertToEngineRecords(expected, schema), actual); + protected void assertRecordsEqual( + Schema schema, List expected, CloseableIterable actual) { + int rowOffset = 0; + + for (T recordOrBatch : actual) { + int numRows = countRows(List.of(recordOrBatch)); + assertThat(rowOffset + numRows).isLessThanOrEqualTo(expected.size()); + assertEquals( + schema, + convertToEngineRecords(expected.subList(rowOffset, rowOffset + numRows), schema), + List.of(recordOrBatch)); + + rowOffset += numRows; + } + + assertThat(rowOffset).isEqualTo(expected.size()); } @TempDir private File tableDir; /** * Returns the set of type IDs that this engine cannot handle. The base implementation returns an - * empty collection; engines override to declare which top-level columns should be excluded from - * test schemas. + * empty collection, engines override to declare which types should be filtered from test schemas + * or cause tests to be skipped. */ protected Collection unsupportedTypeIds() { return Set.of(); @@ -159,7 +172,7 @@ protected void assumeSupportedTypes(Schema schema) { } /** - * Returns the number of rows in the read result; batch engines override to sum across batches. + * Returns the number of rows in the read result, batch engines override to sum across batches. */ protected int countRows(List records) { return records.size(); @@ -289,13 +302,16 @@ void testDataWriterGenericWriteEngineRead(FileFormat fileFormat, DataGenerator d // Read back and verify InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - List readRecords; - try (CloseableIterable reader = - FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) - .project(schema) - .build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(schema, genericRecords, readRecords); + + ReadBuilder builder = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile).project(schema); + + if (supportsBatchReads()) { + builder.recordsPerBatch(2); + } + + try (CloseableIterable reader = builder.build()) { + assertRecordsEqual(schema, genericRecords, reader); } } @@ -330,13 +346,16 @@ void testEqualityDeleteWriterGenericWriteEngineRead( // Read back and verify InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - List readRecords; - try (CloseableIterable reader = - FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) - .project(schema) - .build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(schema, genericRecords, readRecords); + + ReadBuilder builder = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile).project(schema); + + if (supportsBatchReads()) { + builder.recordsPerBatch(2); + } + + try (CloseableIterable reader = builder.build()) { + assertRecordsEqual(schema, genericRecords, reader); } } @@ -398,13 +417,12 @@ void testReaderBuilderProjection(FileFormat fileFormat) throws IOException { List projectedGenericRecords = project(genericRecords, projectedSchema); InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - List readRecords; + try (CloseableIterable reader = FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) .project(projectedSchema) .build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(projectedSchema, projectedGenericRecords, readRecords); + assertRecordsEqual(projectedSchema, projectedGenericRecords, reader); } } @@ -451,7 +469,7 @@ void testReaderBuilderFilter(FileFormat fileFormat) throws IOException { readRecords = ImmutableList.copyOf(reader); } - assertThat(readRecords).hasSizeGreaterThan(0).hasSizeLessThan(genericRecords.size()); + assertThat(countRows(readRecords)).isGreaterThan(0).isLessThan(genericRecords.size()); } /** @@ -2460,14 +2478,13 @@ void testReadFileWithoutFieldIdsUsingNameMapping(FileFormat fileFormat) throws I NameMapping nameMapping = MappingUtil.create(icebergSchema); InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - List readRecords; + try (CloseableIterable reader = FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) .project(icebergSchema) .withNameMapping(nameMapping) .build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(icebergSchema, genericRecords, readRecords); + assertRecordsEqual(icebergSchema, genericRecords, reader); } } @@ -2948,8 +2965,7 @@ private void readAndAssertMetadataColumn( } try (CloseableIterable reader = readerBuilder.build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(projectionSchema, expectedRecords, readRecords); + assertRecordsEqual(projectionSchema, expectedRecords, reader); } } @@ -2991,13 +3007,12 @@ protected void readAndAssertEngineRecords( throws IOException { List expectedGenericRecords = sourceRecords.stream().map(converter).toList(); InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile(); - List readRecords; + try (CloseableIterable reader = FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) .project(readSchema) .build()) { - readRecords = ImmutableList.copyOf(reader); - assertRecordsEqual(readSchema, expectedGenericRecords, readRecords); + assertRecordsEqual(readSchema, expectedGenericRecords, reader); } } } diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java index dc9f86d67827..085e4e85f2eb 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -30,6 +30,7 @@ import org.apache.iceberg.data.DataGenerators; import org.apache.iceberg.data.ReadFormatModelTests; import org.apache.iceberg.data.Record; +import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.spark.SparkUtil; import org.apache.iceberg.spark.data.vectorized.IcebergArrowColumnVector; @@ -86,6 +87,8 @@ protected void assertContainersReused(ColumnarBatch first, ColumnarBatch second) } private ValueVector arrowVector(ColumnarBatch batch, int col) { + assertThat(batch.column(col)).isInstanceOf(IcebergArrowColumnVector.class); + IcebergArrowColumnVector vector = (IcebergArrowColumnVector) batch.column(col); return vector.vectorAccessor().getVector(); } @@ -127,16 +130,21 @@ protected void assertEquals( @Override protected void assertRecordsEqual( - Schema schema, List expected, List actual) { - int rowId = 0; + Schema schema, List expected, CloseableIterable actual) { + int rowOffset = 0; + for (ColumnarBatch batch : actual) { + assertThat(rowOffset + batch.numRows()).isLessThanOrEqualTo(expected.size()); + for (int i = 0; i < batch.numRows(); i++) { TestHelpers.assertEquals( - schema, InternalRowConverter.convert(schema, expected.get(rowId)), batch.getRow(i)); - rowId++; + schema, + InternalRowConverter.convert(schema, expected.get(rowOffset + i)), + batch.getRow(i)); } + rowOffset += batch.numRows(); } - assertThat(rowId).isEqualTo(expected.size()); + assertThat(rowOffset).isEqualTo(expected.size()); } }