This is an automated email from the ASF dual-hosted git repository.
pvary 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 22ed06b979 Data, Spark 4.2: Add format model TCK coverage for
vectorized reads (#17610)
22ed06b979 is described below
commit 22ed06b9792c7161523e5ffa5dde4b199e0c6629
Author: Joy Haldar <[email protected]>
AuthorDate: Wed Nov 4 20:26:35 2026 +0530
Data, Spark 4.2: Add format model TCK coverage for vectorized reads (#17610)
---
.../apache/iceberg/data/BaseFormatModelTests.java | 10 +-
.../org/apache/iceberg/data/DataGenerator.java | 2 +-
.../org/apache/iceberg/data/DataGenerators.java | 4 +-
.../apache/iceberg/data/ReadFormatModelTests.java | 286 ++++++++++++---------
.../spark/data/TestSparkVectorizedFormatModel.java | 207 +++++++++++++++
5 files changed, 382 insertions(+), 127 deletions(-)
diff --git
a/data/src/test/java/org/apache/iceberg/data/BaseFormatModelTests.java
b/data/src/test/java/org/apache/iceberg/data/BaseFormatModelTests.java
index e45ebbbedc..496f77a321 100644
--- a/data/src/test/java/org/apache/iceberg/data/BaseFormatModelTests.java
+++ b/data/src/test/java/org/apache/iceberg/data/BaseFormatModelTests.java
@@ -63,7 +63,7 @@ public abstract class BaseFormatModelTests<T> extends
ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testDataWriterEngineWriteGenericRead(FileFormat fileFormat,
DataGenerator dataGenerator)
throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
List<Record> genericRecords = project(dataGenerator.generateRecords(),
schema);
List<T> engineRecords = convertToEngineRecords(genericRecords, schema);
writeEngineRecords(fileFormat, schema, engineRecords,
engineSchema(schema));
@@ -75,7 +75,7 @@ public abstract class BaseFormatModelTests<T> extends
ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testDataWriterEngineWriteWithoutEngineSchema(
FileFormat fileFormat, DataGenerator dataGenerator) throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
List<Record> genericRecords = project(dataGenerator.generateRecords(),
schema);
List<T> engineRecords = convertToEngineRecords(genericRecords, schema);
writeEngineRecords(fileFormat, schema, engineRecords);
@@ -87,7 +87,7 @@ public abstract class BaseFormatModelTests<T> extends
ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testDataWriterEngineWriteEngineRead(FileFormat fileFormat,
DataGenerator dataGenerator)
throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
List<Record> genericRecords = project(dataGenerator.generateRecords(),
schema);
List<T> engineRecords = convertToEngineRecords(genericRecords, schema);
writeEngineRecords(fileFormat, schema, engineRecords);
@@ -99,7 +99,7 @@ public abstract class BaseFormatModelTests<T> extends
ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testEqualityDeleteWriterEngineWriteGenericRead(
FileFormat fileFormat, DataGenerator dataGenerator) throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
FileWriterBuilder<EqualityDeleteWriter<T>, Object> writerBuilder =
FormatModelRegistry.equalityDeleteWriteBuilder(fileFormat,
engineType(), encryptedFile);
@@ -135,7 +135,7 @@ public abstract class BaseFormatModelTests<T> extends
ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testEqualityDeleteWriterEngineWriteWithoutEngineSchema(
FileFormat fileFormat, DataGenerator dataGenerator) throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
FileWriterBuilder<EqualityDeleteWriter<T>, Object> writerBuilder =
FormatModelRegistry.equalityDeleteWriteBuilder(fileFormat,
engineType(), encryptedFile);
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 ad068fbc68..7c969f44e2 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 java.util.List;
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 cf2ab00d7c..bdfbe3acf8 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 @@ import org.apache.iceberg.util.DateTimeUtil;
* 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 b8596c632b..971012e274 100644
--- a/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java
+++ b/data/src/test/java/org/apache/iceberg/data/ReadFormatModelTests.java
@@ -24,7 +24,6 @@ import static org.apache.iceberg.TestBase.SCHEMA;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assumptions.assumeThat;
-import static org.junit.jupiter.api.Assumptions.assumeFalse;
import java.io.File;
import java.io.IOException;
@@ -105,6 +104,8 @@ import org.junit.jupiter.params.provider.FieldSource;
*/
public abstract class ReadFormatModelTests<T> {
+ @TempDir private File tableDir;
+
protected abstract Class<T> engineType();
protected abstract Object engineSchema(Schema schema);
@@ -119,19 +120,49 @@ public abstract class ReadFormatModelTests<T> {
return false;
}
- @TempDir private File tableDir;
+ protected void assertRecordsEqual(
+ Schema schema, List<Record> expected, CloseableIterable<T> 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());
+ }
/**
* 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<Type.TypeID> unsupportedTypeIds() {
return Set.of();
}
- private Schema filterUnsupported(Schema schema) {
- Collection<Type.TypeID> unsupportedTypeIds = unsupportedTypeIds();
+ /**
+ * Types unsupported by this engine for the given file format. Defaults to
the format-independent
+ * set.
+ */
+ protected Collection<Type.TypeID> unsupportedTypeIds(FileFormat fileFormat) {
+ return unsupportedTypeIds();
+ }
+
+ /** Returns true if the schema contains any of the given types, ignoring the
root struct. */
+ protected static boolean containsType(Schema schema, Collection<Type.TypeID>
typeIds) {
+ Types.StructType root = schema.asStruct();
+ return TypeUtil.find(schema, type -> type != root &&
typeIds.contains(type.typeId())) != null;
+ }
+
+ private Schema filterUnsupported(FileFormat fileFormat, Schema schema) {
+ Collection<Type.TypeID> unsupportedTypeIds =
unsupportedTypeIds(fileFormat);
if (unsupportedTypeIds.isEmpty()) {
return schema;
}
@@ -145,8 +176,35 @@ public abstract class ReadFormatModelTests<T> {
return kept.size() == schema.columns().size() ? schema : new Schema(kept);
}
- protected Schema supportedSchema(DataGenerator generator) {
- Schema filtered = filterUnsupported(generator.schema());
+ protected void assumeSupportedTypes(FileFormat fileFormat, Schema schema) {
+ assumeThat(containsType(schema, unsupportedTypeIds(fileFormat)))
+ .as("Schema contains a type unsupported by this engine")
+ .isFalse();
+ }
+
+ /**
+ * Returns the number of rows in the read result, batch engines override to
sum across batches.
+ */
+ protected int countRows(List<T> records) {
+ return records.size();
+ }
+
+ private ReadBuilder<T, ?> readBuilder(FileFormat fileFormat, InputFile
inputFile) {
+ ReadBuilder<T, ?> builder =
+ FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile);
+ return supportsBatchReads() ? builder.recordsPerBatch(3) : builder;
+ }
+
+ 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(FileFormat fileFormat, DataGenerator
generator) {
+ Schema filtered = filterUnsupported(fileFormat, generator.schema());
assumeThat(filtered.columns())
.as("All columns of %s are unsupported by this engine", generator)
.isNotEmpty();
@@ -254,22 +312,17 @@ public abstract class ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testDataWriterGenericWriteEngineRead(FileFormat fileFormat,
DataGenerator dataGenerator)
throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
List<Record> genericRecords = project(dataGenerator.generateRecords(),
schema);
writeGenericRecords(fileFormat, schema, genericRecords);
// Read back and verify
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
- try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .build()) {
- readRecords = ImmutableList.copyOf(reader);
- }
- assertEquals(schema, convertToEngineRecords(genericRecords, schema),
readRecords);
+ try (CloseableIterable<T> reader = readBuilder(fileFormat,
inputFile).project(schema).build()) {
+ assertRecordsEqual(schema, genericRecords, reader);
+ }
}
/** Write with Generic Record, read with engine type T */
@@ -277,7 +330,7 @@ public abstract class ReadFormatModelTests<T> {
@FieldSource("FORMAT_AND_GENERATOR")
void testEqualityDeleteWriterGenericWriteEngineRead(
FileFormat fileFormat, DataGenerator dataGenerator) throws IOException {
- Schema schema = supportedSchema(dataGenerator);
+ Schema schema = supportedSchema(fileFormat, dataGenerator);
FileWriterBuilder<EqualityDeleteWriter<Record>, Object> writerBuilder =
FormatModelRegistry.equalityDeleteWriteBuilder(fileFormat,
Record.class, encryptedFile);
@@ -303,15 +356,10 @@ public abstract class ReadFormatModelTests<T> {
// Read back and verify
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
- try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .build()) {
- readRecords = ImmutableList.copyOf(reader);
- }
- assertEquals(schema, convertToEngineRecords(genericRecords, schema),
readRecords);
+ try (CloseableIterable<T> reader = readBuilder(fileFormat,
inputFile).project(schema).build()) {
+ assertRecordsEqual(schema, genericRecords, reader);
+ }
}
/** Write position deletes, read with Generic Record */
@@ -370,19 +418,13 @@ public abstract class ReadFormatModelTests<T> {
writeGenericRecords(fileFormat, fullSchema, genericRecords);
List<Record> projectedGenericRecords = project(genericRecords,
projectedSchema);
- List<T> expectedEngineRecords =
- convertToEngineRecords(projectedGenericRecords, projectedSchema);
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
+
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(projectedSchema)
- .build()) {
- readRecords = ImmutableList.copyOf(reader);
+ readBuilder(fileFormat, inputFile).project(projectedSchema).build()) {
+ assertRecordsEqual(projectedSchema, projectedGenericRecords, reader);
}
-
- assertEquals(projectedSchema, expectedEngineRecords, readRecords);
}
@ParameterizedTest
@@ -408,27 +450,21 @@ public abstract class ReadFormatModelTests<T> {
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
List<T> readRecords;
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .filter(lessThanFilter)
- .build()) {
+ readBuilder(fileFormat,
inputFile).project(schema).filter(lessThanFilter).build()) {
readRecords = ImmutableList.copyOf(reader);
}
- assertThat(readRecords).isEmpty();
+ assertThat(countRows(readRecords)).isZero();
// Filter: id >= count - 1, so only the last record matches across all row
groups
Expression greaterThanFilter = Expressions.greaterThanOrEqual("id", count
- 1);
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .filter(greaterThanFilter)
- .build()) {
+ readBuilder(fileFormat,
inputFile).project(schema).filter(greaterThanFilter).build()) {
readRecords = ImmutableList.copyOf(reader);
}
-
assertThat(readRecords).hasSizeGreaterThan(0).hasSizeLessThan(genericRecords.size());
+
assertThat(countRows(readRecords)).isGreaterThan(0).isLessThan(genericRecords.size());
}
/**
@@ -458,7 +494,7 @@ public abstract class ReadFormatModelTests<T> {
// caseSensitive=false: upper-cased column name must be resolved correctly.
List<T> readRecords;
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
+ readBuilder(fileFormat, inputFile)
.project(schema)
.filter(upperCaseFilter)
.caseSensitive(false)
@@ -466,13 +502,13 @@ public abstract class ReadFormatModelTests<T> {
readRecords = ImmutableList.copyOf(reader);
}
- assertThat(readRecords).isNotEmpty();
+ assertThat(countRows(readRecords)).isGreaterThan(0);
// caseSensitive=true: upper-cased column name cannot be resolved, so must
throw.
assertThatThrownBy(
() -> {
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(),
inputFile)
+ readBuilder(fileFormat, inputFile)
.project(schema)
.filter(upperCaseFilter)
.caseSensitive(true)
@@ -495,7 +531,6 @@ public abstract class ReadFormatModelTests<T> {
@ParameterizedTest
@FieldSource("FILE_FORMATS")
void testReaderBuilderSplit(FileFormat fileFormat) throws IOException {
-
assumeSupports(fileFormat, FEATURE_SPLIT);
DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
@@ -514,44 +549,37 @@ public abstract class ReadFormatModelTests<T> {
List<T> readRecords;
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
+ readBuilder(fileFormat, inputFile)
.project(schema)
.split(firstSplitStart, firstSplitLength)
.build()) {
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<T> emptyReadRecords;
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .split(fileLength, 0)
- .build()) {
+ readBuilder(fileFormat, inputFile).project(schema).split(fileLength,
0).build()) {
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<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .split(0, fileLength)
- .build()) {
+ readBuilder(fileFormat, inputFile).project(schema).split(0,
fileLength).build()) {
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 +592,25 @@ public abstract class ReadFormatModelTests<T> {
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
// Without reuseContainers: every record must be a distinct object instance
- List<T> noReuseRecords;
- try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .build()) {
- noReuseRecords = ImmutableList.copyOf(reader);
- }
+ try (CloseableIterable<T> reader = readBuilder(fileFormat,
inputFile).project(schema).build()) {
+ List<T> noReuseRecords = ImmutableList.copyOf(reader);
- 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));
+ 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<T> reuseRecords;
+ // With reuseContainers: verify that the underlying containers are reused
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(schema)
- .reuseContainers()
- .build()) {
- reuseRecords = ImmutableList.copyOf(reader);
- }
+ readBuilder(fileFormat,
inputFile).project(schema).reuseContainers().build()) {
+ List<T> reuseRecords = ImmutableList.copyOf(reader);
- reuseRecords.forEach(r -> assertThat(r).isSameAs(reuseRecords.get(0)));
+ assertThat(reuseRecords).hasSizeGreaterThan(1);
+ reuseRecords.forEach(r -> assertContainersReused(reuseRecords.get(0),
r));
+ }
}
@ParameterizedTest
@@ -709,6 +731,7 @@ public abstract class ReadFormatModelTests<T> {
.withDoc("Used to test nested field defaults")
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
List<Record> genericRecords = RandomGenericData.generate(writeSchema, 10,
1L);
writeGenericRecords(fileFormat, writeSchema, genericRecords);
@@ -779,6 +802,7 @@ public abstract class ReadFormatModelTests<T> {
.withDoc("Used to test nested map value field defaults")
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
List<Record> genericRecords = RandomGenericData.generate(writeSchema, 10,
1L);
writeGenericRecords(fileFormat, writeSchema, genericRecords);
@@ -864,6 +888,7 @@ public abstract class ReadFormatModelTests<T> {
.withDoc("Used to test nested field defaults")
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
List<Record> genericRecords = RandomGenericData.generate(writeSchema, 10,
1L);
writeGenericRecords(fileFormat, writeSchema, genericRecords);
@@ -938,6 +963,8 @@ public abstract class ReadFormatModelTests<T> {
Types.NestedField.optional("nested").withId(2).ofType(writeNested).build();
Schema writeSchema = new Schema(idField, nestedField);
+ assumeSupportedTypes(fileFormat, writeSchema);
+
Record presentNested = GenericRecord.create(writeNested);
presentNested.setField("a", "x");
presentNested.setField("b", "y");
@@ -1001,6 +1028,7 @@ public abstract class ReadFormatModelTests<T> {
Types.NestedField.required(5, "inner",
Types.StringType.get())))
.build();
Schema writeSchema = new Schema(idField, nestedField);
+ assumeSupportedTypes(fileFormat, writeSchema);
Types.StructType nestedType = nestedField.type().asStructType();
// alternate present and null structs, with a different number of list
elements per row so
@@ -1076,6 +1104,7 @@ public abstract class ReadFormatModelTests<T> {
3, "tags", Types.ListType.ofRequired(4,
Types.StringType.get()))))
.build();
Schema writeSchema = new Schema(idField, nestedField);
+ assumeSupportedTypes(fileFormat, writeSchema);
Types.StructType nestedType = nestedField.type().asStructType();
List<Record> genericRecords = Lists.newArrayList();
@@ -1113,6 +1142,8 @@ public abstract class ReadFormatModelTests<T> {
Types.StructType.of(Types.NestedField.required(3, "inner",
Types.StringType.get())));
Schema writeSchema = new Schema(idField, nestedField);
+ assumeSupportedTypes(fileFormat, writeSchema);
+
Record present = GenericRecord.create(writeSchema);
present.setField("id", 1L);
Record presentNested =
GenericRecord.create(nestedField.type().asStructType());
@@ -1170,6 +1201,7 @@ public abstract class ReadFormatModelTests<T> {
.ofType(Types.ListType.ofOptional(3, writeElement))
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
Record present = GenericRecord.create(writeElement);
present.setField("inner", "a");
Record row = GenericRecord.create(writeSchema);
@@ -1225,6 +1257,7 @@ public abstract class ReadFormatModelTests<T> {
.ofType(Types.ListType.ofOptional(3, writeElement))
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
Record nullList = GenericRecord.create(writeSchema);
nullList.setField("id", 1L);
nullList.setField("arr", null);
@@ -1279,6 +1312,8 @@ public abstract class ReadFormatModelTests<T> {
.ofType(Types.MapType.ofOptional(3, 4, Types.StringType.get(),
writeValue))
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
+
Record present = GenericRecord.create(writeValue);
present.setField("inner", "a");
Map<String, Record> writtenMap = Maps.newLinkedHashMap();
@@ -1340,6 +1375,8 @@ public abstract class ReadFormatModelTests<T> {
.ofType(Types.MapType.ofOptional(3, 4, Types.StringType.get(),
writeValue))
.build());
+ assumeSupportedTypes(fileFormat, writeSchema);
+
Record nullMap = GenericRecord.create(writeSchema);
nullMap.setField("id", 1L);
nullMap.setField("map", null);
@@ -1411,7 +1448,7 @@ public abstract class ReadFormatModelTests<T> {
assumeSupports(fileFormat, FEATURE_READER_DEFAULT);
Schema writeSchema = DataGenerators.PrimitiveDefaults.WRITE_SCHEMA;
- Schema readSchema =
filterUnsupported(DataGenerators.PrimitiveDefaults.READ_SCHEMA);
+ Schema readSchema = filterUnsupported(fileFormat,
DataGenerators.PrimitiveDefaults.READ_SCHEMA);
List<Record> sourceRecords = RandomGenericData.generate(writeSchema, 10,
1L);
writeGenericRecords(fileFormat, writeSchema, sourceRecords);
@@ -1440,7 +1477,7 @@ public abstract class ReadFormatModelTests<T> {
void testPrimitiveDefaultValuesNotApplied(FileFormat fileFormat) throws
IOException {
assumeSupports(fileFormat, FEATURE_READER_DEFAULT);
- Schema readSchema =
filterUnsupported(DataGenerators.PrimitiveDefaults.READ_SCHEMA);
+ Schema readSchema = filterUnsupported(fileFormat,
DataGenerators.PrimitiveDefaults.READ_SCHEMA);
List<Record> sourceRecords = RandomGenericData.generate(readSchema, 10,
1L);
writeGenericRecords(fileFormat, readSchema, sourceRecords);
@@ -1451,7 +1488,9 @@ public abstract class ReadFormatModelTests<T> {
@ParameterizedTest
@FieldSource("FILE_FORMATS")
void testReaderBuilderRecordsPerBatchNotSupported(FileFormat fileFormat)
throws IOException {
- assumeFalse(supportsBatchReads(), engineType().getSimpleName() + "
supports batch reads");
+ assumeThat(supportsBatchReads())
+ .as(engineType().getSimpleName() + " supports batch reads")
+ .isFalse();
DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
Schema schema = dataGenerator.schema();
@@ -1467,6 +1506,35 @@ public abstract class ReadFormatModelTests<T> {
.isInstanceOf(UnsupportedOperationException.class);
}
+ @ParameterizedTest
+ @FieldSource("FILE_FORMATS")
+ void testReaderBuilderRecordsPerBatch(FileFormat fileFormat) throws
IOException {
+ assumeThat(supportsBatchReads())
+ .as(engineType().getSimpleName() + " does not support batch reads")
+ .isTrue();
+
+ DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
+ Schema schema = dataGenerator.schema();
+ List<Record> genericRecords = dataGenerator.generateRecords(1000);
+ writeGenericRecords(fileFormat, schema, genericRecords);
+
+ InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
+ List<T> readRecords;
+ try (CloseableIterable<T> 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
@FieldSource("FILE_FORMATS")
void testDataWriterMetricsCollection(FileFormat fileFormat) throws
IOException {
@@ -1842,8 +1910,7 @@ public abstract class ReadFormatModelTests<T> {
@ParameterizedTest
@FieldSource("FILE_FORMATS")
- void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws
IOException {
-
+ protected void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws
IOException {
DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
Schema schema = dataGenerator.schema();
List<Record> genericRecords = dataGenerator.generateRecords();
@@ -1862,8 +1929,7 @@ public abstract class ReadFormatModelTests<T> {
@ParameterizedTest
@FieldSource("FILE_FORMATS")
- void testReadMetadataColumnRowLineage(FileFormat fileFormat) throws
IOException {
-
+ protected void testReadMetadataColumnRowLineage(FileFormat fileFormat)
throws IOException {
DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
Schema schema = dataGenerator.schema();
List<Record> genericRecords = dataGenerator.generateRecords();
@@ -1895,8 +1961,8 @@ public abstract class ReadFormatModelTests<T> {
@ParameterizedTest
@FieldSource("FILE_FORMATS")
- void testReadMetadataColumnRowLineageExistingValues(FileFormat fileFormat)
throws IOException {
-
+ protected void testReadMetadataColumnRowLineageExistingValues(FileFormat
fileFormat)
+ throws IOException {
DataGenerator dataGenerator = new DataGenerators.DefaultSchema();
Schema dataSchema = dataGenerator.schema();
@@ -2367,13 +2433,11 @@ public abstract class ReadFormatModelTests<T> {
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
List<T> readRecords;
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(emptySchema)
- .build()) {
+ readBuilder(fileFormat, inputFile).project(emptySchema).build()) {
readRecords = ImmutableList.copyOf(reader);
}
- assertThat(readRecords).hasSameSizeAs(genericRecords);
+ assertThat(countRows(readRecords)).isEqualTo(genericRecords.size());
}
@ParameterizedTest
@@ -2392,16 +2456,14 @@ public abstract class ReadFormatModelTests<T> {
NameMapping nameMapping = MappingUtil.create(icebergSchema);
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
+
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
+ readBuilder(fileFormat, 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 +2518,7 @@ public abstract class ReadFormatModelTests<T> {
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);
@@ -2870,23 +2932,15 @@ public abstract class ReadFormatModelTests<T> {
throws IOException {
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
-
- ReadBuilder<T, ?> readerBuilder =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(projectionSchema);
+ ReadBuilder<T, ?> readerBuilder = readBuilder(fileFormat,
inputFile).project(projectionSchema);
if (idToConstant != null) {
readerBuilder.idToConstant(convertConstantsToEngine(projectionSchema,
idToConstant));
}
try (CloseableIterable<T> 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 +2981,10 @@ public abstract class ReadFormatModelTests<T> {
throws IOException {
List<Record> expectedGenericRecords =
sourceRecords.stream().map(converter).toList();
InputFile inputFile = encryptedFile.encryptingOutputFile().toInputFile();
- List<T> readRecords;
+
try (CloseableIterable<T> reader =
- FormatModelRegistry.readBuilder(fileFormat, engineType(), inputFile)
- .project(readSchema)
- .build()) {
- readRecords = ImmutableList.copyOf(reader);
+ readBuilder(fileFormat, inputFile).project(readSchema).build()) {
+ 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 0000000000..88b4d7e214
--- /dev/null
+++
b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkVectorizedFormatModel.java
@@ -0,0 +1,207 @@
+/*
+ * 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 static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.io.IOException;
+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.DataGenerator;
+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.relocated.com.google.common.collect.ImmutableSet;
+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.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.FieldSource;
+
+class TestSparkVectorizedFormatModel extends
ReadFormatModelTests<ColumnarBatch> {
+
+ private static final FileFormat[] FILE_FORMATS = {FileFormat.ORC,
FileFormat.PARQUET};
+
+ private static final Set<Type.TypeID> SPARK_UNSUPPORTED_TYPE_IDS =
+ Set.of(Type.TypeID.TIME, Type.TypeID.TIMESTAMP_NANO);
+
+ // The Arrow-based Parquet reader cannot read nested or geospatial columns;
+ // the vectorized ORC reader supports nested types.
+ private static final Set<Type.TypeID> ARROW_UNSUPPORTED_TYPE_IDS =
+ Set.of(
+ Type.TypeID.STRUCT,
+ Type.TypeID.LIST,
+ Type.TypeID.MAP,
+ Type.TypeID.GEOMETRY,
+ Type.TypeID.GEOGRAPHY);
+
+ private static final Set<Type.TypeID> PARQUET_UNSUPPORTED_TYPE_IDS =
+ ImmutableSet.<Type.TypeID>builder()
+ .addAll(SPARK_UNSUPPORTED_TYPE_IDS)
+ .addAll(ARROW_UNSUPPORTED_TYPE_IDS)
+ .build();
+
+ private static final List<Arguments> FORMAT_AND_GENERATOR =
+ Arrays.stream(FILE_FORMATS)
+ .flatMap(
+ format ->
+ Arrays.stream(DataGenerators.ALL)
+ .filter(generator -> supportsVectorized(format,
generator))
+ .map(generator -> Arguments.of(format, generator)))
+ .toList();
+
+ private static boolean supportsVectorized(FileFormat format, DataGenerator
generator) {
+ return supportsGenerator(format, generator)
+ && (format != FileFormat.PARQUET
+ || !containsType(generator.schema(), ARROW_UNSUPPORTED_TYPE_IDS));
+ }
+
+ @Override
+ protected Collection<Type.TypeID> unsupportedTypeIds(FileFormat fileFormat) {
+ return fileFormat == FileFormat.PARQUET
+ ? PARQUET_UNSUPPORTED_TYPE_IDS
+ : SPARK_UNSUPPORTED_TYPE_IDS;
+ }
+
+ @Override
+ protected boolean supportsBatchReads() {
+ return true;
+ }
+
+ @Override
+ protected int countRows(List<ColumnarBatch> 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));
+ }
+ }
+
+ @Override
+ @ParameterizedTest
+ @FieldSource("FILE_FORMATS")
+ protected void testReadMetadataColumnIsDeleted(FileFormat fileFormat) throws
IOException {
+ assumeThat(fileFormat)
+ .as("Vectorized Parquet reader only populates _deleted when a delete
filter is set")
+ .isNotEqualTo(FileFormat.PARQUET);
+ super.testReadMetadataColumnIsDeleted(fileFormat);
+ }
+
+ @Override
+ @ParameterizedTest
+ @FieldSource("FILE_FORMATS")
+ protected void testReadMetadataColumnRowLineage(FileFormat fileFormat)
throws IOException {
+ assumeRowLineageSupported(fileFormat);
+ super.testReadMetadataColumnRowLineage(fileFormat);
+ }
+
+ @Override
+ @ParameterizedTest
+ @FieldSource("FILE_FORMATS")
+ protected void testReadMetadataColumnRowLineageExistingValues(FileFormat
fileFormat)
+ throws IOException {
+ assumeRowLineageSupported(fileFormat);
+ super.testReadMetadataColumnRowLineageExistingValues(fileFormat);
+ }
+
+ // TODO: remove when https://github.com/apache/iceberg/issues/18543 is fixed
+ private static void assumeRowLineageSupported(FileFormat fileFormat) {
+ assumeThat(fileFormat)
+ .as("Vectorized ORC reader does not apply row lineage")
+ .isNotEqualTo(FileFormat.ORC);
+ }
+
+ 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<ColumnarBatch> 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<ColumnarBatch> expected, List<ColumnarBatch> actual)
{
+ throw new UnsupportedOperationException("Use assertRecordsEqual for batch
reads");
+ }
+
+ @Override
+ protected void assertRecordsEqual(
+ Schema schema, List<Record> expected, CloseableIterable<ColumnarBatch>
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());
+ }
+}
