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 cf2ab00d7c02..bdfbe3acf86a 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 b8596c632b49..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,12 +119,30 @@ protected boolean supportsBatchReads() { return false; } + 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(); @@ -145,6 +163,29 @@ 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(); + } + + /** + * 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 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()) @@ -261,15 +302,17 @@ 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); + + ReadBuilder builder = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile).project(schema); + + if (supportsBatchReads()) { + builder.recordsPerBatch(2); } - assertEquals(schema, convertToEngineRecords(genericRecords, schema), readRecords); + try (CloseableIterable reader = builder.build()) { + assertRecordsEqual(schema, genericRecords, reader); + } } /** Write with Generic Record, read with engine type T */ @@ -303,15 +346,17 @@ 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); + + ReadBuilder builder = + FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile).project(schema); + + if (supportsBatchReads()) { + builder.recordsPerBatch(2); } - assertEquals(schema, convertToEngineRecords(genericRecords, schema), readRecords); + try (CloseableIterable reader = builder.build()) { + assertRecordsEqual(schema, genericRecords, reader); + } } /** Write position deletes, read with Generic Record */ @@ -370,19 +415,15 @@ 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; + try (CloseableIterable reader = FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile) .project(projectedSchema) .build()) { - readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(projectedSchema, projectedGenericRecords, reader); } - - assertEquals(projectedSchema, expectedEngineRecords, readRecords); } @ParameterizedTest @@ -428,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()); } /** @@ -495,7 +536,6 @@ void testReaderBuilderCaseSensitive(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReaderBuilderSplit(FileFormat fileFormat) throws IOException { - assumeSupports(fileFormat, FEATURE_SPLIT); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -521,7 +561,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; @@ -533,7 +573,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 = @@ -544,14 +584,13 @@ 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 */ @ParameterizedTest @FieldSource("FILE_FORMATS") void testReaderBuilderReuseContainers(FileFormat fileFormat) throws IOException { - assumeSupports(fileFormat, FEATURE_REUSE_CONTAINERS); DataGenerator dataGenerator = new DataGenerators.DefaultSchema(); @@ -564,31 +603,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 @@ -709,6 +757,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); @@ -779,6 +828,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); @@ -864,6 +914,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); @@ -938,6 +989,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"); @@ -1001,6 +1054,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 @@ -1076,6 +1130,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(); @@ -1113,6 +1168,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()); @@ -1170,6 +1227,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); @@ -1225,6 +1283,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); @@ -1279,6 +1338,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(); @@ -1340,6 +1401,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); @@ -1450,21 +1513,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 @@ -1843,7 +1921,9 @@ void testReadMetadataColumnRowPosition(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { - + 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(); @@ -1863,7 +1943,10 @@ void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws IOException { @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException { - + // 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(); @@ -1896,7 +1979,10 @@ void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws IOException @ParameterizedTest @FieldSource("FILE_FORMATS") void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat) throws IOException { - + // 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(); @@ -2373,7 +2459,7 @@ void testSchemaEvolutionEmptyProjection(FileFormat fileFormat) throws IOExceptio readRecords = ImmutableList.copyOf(reader); } - assertThat(readRecords).hasSameSizeAs(genericRecords); + assertThat(countRows(readRecords)).isEqualTo(genericRecords.size()); } @ParameterizedTest @@ -2392,16 +2478,14 @@ 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, reader); } - - assertEquals(icebergSchema, convertToEngineRecords(genericRecords, icebergSchema), readRecords); } protected void readAndAssertGenericRecords( @@ -2456,7 +2540,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) { Schema schema = generator.schema(); return supportsTypes(fileFormat, schema, VARIANT_TYPE_IDS, FEATURE_VARIANT) && supportsTypes(fileFormat, schema, GEOSPATIAL_TYPE_IDS, FEATURE_GEOSPATIAL); @@ -2881,12 +2965,8 @@ private void readAndAssertMetadataColumn( } try (CloseableIterable reader = readerBuilder.build()) { - readRecords = ImmutableList.copyOf(reader); + assertRecordsEqual(projectionSchema, expectedRecords, reader); } - - assertThat(readRecords).hasSize(expectedRecords.size()); - assertEquals( - projectionSchema, convertToEngineRecords(expectedRecords, projectionSchema), readRecords); } private static Record copy(Record source, Schema sourceSchema, Schema targetSchema) { @@ -2927,16 +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, reader); } - - assertThat(readRecords).hasSize(expectedGenericRecords.size()); - assertEquals( - readSchema, convertToEngineRecords(expectedGenericRecords, readSchema), readRecords); } } 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 new file mode 100644 index 000000000000..085e4e85f2eb --- /dev/null +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java @@ -0,0 +1,150 @@ +/* + * 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.arrow.vector.ValueVector; +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.io.CloseableIterable; +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; + +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, + Type.TypeID.GEOMETRY, + Type.TypeID.GEOGRAPHY); + + @Override + protected Collection unsupportedTypeIds() { + return UNSUPPORTED_TYPE_IDS; + } + + @Override + protected boolean supportsBatchReads() { + return true; + } + + @Override + 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) { + assertThat(batch.column(col)).isInstanceOf(IcebergArrowColumnVector.class); + + 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); + } + + @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, 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(rowOffset + i)), + batch.getRow(i)); + } + rowOffset += batch.numRows(); + } + + assertThat(rowOffset).isEqualTo(expected.size()); + } +}