diff --git a/orc/src/main/java/org/apache/iceberg/orc/ORC.java b/orc/src/main/java/org/apache/iceberg/orc/ORC.java index 89cd1ad43..bfe00e788 100644 --- a/orc/src/main/java/org/apache/iceberg/orc/ORC.java +++ b/orc/src/main/java/org/apache/iceberg/orc/ORC.java @@ -672,6 +672,7 @@ public static class ReadBuilder { private Function> readerFunc; private Function> batchedReaderFunc; + private boolean supportsInitialDefaults = false; private int recordsPerBatch = VectorizedRowBatch.DEFAULT_SIZE; private ReadBuilder(InputFile file) { @@ -725,6 +726,19 @@ public ReadBuilder createReaderFunc(Function> r return this; } + /** + * Signals that the configured reader can fill {@code initial-default} values for fields omitted + * from an ORC file. Disabled by default so existing readers retain null-synthesizing projection + * behavior. + */ + public ReadBuilder supportsInitialDefaults() { + Preconditions.checkState( + this.readerFunc != null || this.batchedReaderFunc != null, + "A reader function must be configured before enabling initial defaults"); + this.supportsInitialDefaults = true; + return this; + } + public ReadBuilder filter(Expression newFilter) { this.filter = newFilter; return this; @@ -733,7 +747,7 @@ public ReadBuilder filter(Expression newFilter) { public ReadBuilder createBatchedReaderFunc( Function> batchReaderFunction) { Preconditions.checkArgument( - this.readerFunc == null, + this.readerFunc == null && !this.supportsInitialDefaults, "Batched reader function cannot be set since the non-batched version is already set"); this.batchedReaderFunc = batchReaderFunction; return this; @@ -759,6 +773,7 @@ public CloseableIterable build() { start, length, readerFunc, + supportsInitialDefaults, caseSensitive, filter, batchedReaderFunc, diff --git a/orc/src/main/java/org/apache/iceberg/orc/ORCSchemaUtil.java b/orc/src/main/java/org/apache/iceberg/orc/ORCSchemaUtil.java index fae1a76c3..e641c53dd 100644 --- a/orc/src/main/java/org/apache/iceberg/orc/ORCSchemaUtil.java +++ b/orc/src/main/java/org/apache/iceberg/orc/ORCSchemaUtil.java @@ -261,18 +261,41 @@ public static Schema convert(TypeDescription orcSchema) { */ public static TypeDescription buildOrcProjection( Schema schema, TypeDescription originalOrcSchema) { + return buildOrcProjection(schema, originalOrcSchema, false); + } + + /** + * Builds the ORC read projection, optionally omitting absent fields that declare an {@code + * initial-default} so a default-aware reader can fill them via {@code idToConstant}. + * + *

When {@code supportsInitialDefaults} is true and a field is absent from the file but + * declares {@code initialDefault()}, it is omitted instead of being synthesized as a null column. + */ + static TypeDescription buildOrcProjection( + Schema schema, TypeDescription originalOrcSchema, boolean supportsInitialDefaults) { final Map icebergToOrc = icebergToOrcMapping("root", originalOrcSchema); - return buildOrcProjection(Integer.MIN_VALUE, schema.asStruct(), true, icebergToOrc); + return buildOrcProjection( + Integer.MIN_VALUE, schema.asStruct(), true, supportsInitialDefaults, icebergToOrc); } private static TypeDescription buildOrcProjection( - Integer fieldId, Type type, boolean isRequired, Map mapping) { + Integer fieldId, + Type type, + boolean isRequired, + boolean supportsInitialDefaults, + Map mapping) { final TypeDescription orcType; switch (type.typeId()) { case STRUCT: orcType = TypeDescription.createStruct(); for (Types.NestedField nestedField : type.asStructType().fields()) { + // Omit so the reader fills via idToConstant instead of a synthetic null column. + if (supportsInitialDefaults + && mapping.get(nestedField.fieldId()) == null + && nestedField.initialDefault() != null) { + continue; + } // Using suffix _r to avoid potential underlying issues in ORC reader // with reused column names between ORC and Iceberg; // e.g. renaming column c -> d and adding new column d @@ -285,6 +308,7 @@ private static TypeDescription buildOrcProjection( nestedField.fieldId(), nestedField.type(), isRequired && nestedField.isRequired(), + supportsInitialDefaults, mapping); orcType.addField(name, childType); } @@ -296,16 +320,22 @@ private static TypeDescription buildOrcProjection( list.elementId(), list.elementType(), isRequired && list.isElementRequired(), + supportsInitialDefaults, mapping); orcType = TypeDescription.createList(elementType); break; case MAP: Types.MapType map = (Types.MapType) type; TypeDescription keyType = - buildOrcProjection(map.keyId(), map.keyType(), isRequired, mapping); + buildOrcProjection( + map.keyId(), map.keyType(), isRequired, supportsInitialDefaults, mapping); TypeDescription valueType = buildOrcProjection( - map.valueId(), map.valueType(), isRequired && map.isValueRequired(), mapping); + map.valueId(), + map.valueType(), + isRequired && map.isValueRequired(), + supportsInitialDefaults, + mapping); orcType = TypeDescription.createMap(keyType, valueType); break; default: diff --git a/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java b/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java index 58cf5d1f9..cd8148fdb 100644 --- a/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java +++ b/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java @@ -46,6 +46,7 @@ class OrcIterable extends CloseableGroup implements CloseableIterable { private final Long start; private final Long length; private final Function> readerFunction; + private final boolean supportsInitialDefaults; private final Expression filter; private final boolean caseSensitive; private final Function> batchReaderFunction; @@ -60,12 +61,14 @@ class OrcIterable extends CloseableGroup implements CloseableIterable { Long start, Long length, Function> readerFunction, + boolean supportsInitialDefaults, boolean caseSensitive, Expression filter, Function> batchReaderFunction, int recordsPerBatch) { this.schema = schema; this.readerFunction = readerFunction; + this.supportsInitialDefaults = supportsInitialDefaults; this.file = file; this.nameMapping = nameMapping; this.start = start; @@ -86,13 +89,14 @@ public CloseableIterator iterator() { TypeDescription fileSchema = orcFileReader.getSchema(); final TypeDescription readOrcSchema; if (ORCSchemaUtil.hasIds(fileSchema)) { - readOrcSchema = ORCSchemaUtil.buildOrcProjection(schema, fileSchema); + readOrcSchema = ORCSchemaUtil.buildOrcProjection(schema, fileSchema, supportsInitialDefaults); } else { if (nameMapping == null) { nameMapping = MappingUtil.create(schema); } TypeDescription typeWithIds = ORCSchemaUtil.applyNameMapping(fileSchema, nameMapping); - readOrcSchema = ORCSchemaUtil.buildOrcProjection(schema, typeWithIds); + readOrcSchema = + ORCSchemaUtil.buildOrcProjection(schema, typeWithIds, supportsInitialDefaults); } SearchArgument sarg = null; diff --git a/orc/src/main/java/org/apache/iceberg/orc/OrcSchemaWithTypeVisitor.java b/orc/src/main/java/org/apache/iceberg/orc/OrcSchemaWithTypeVisitor.java index fd37283a8..568de0219 100644 --- a/orc/src/main/java/org/apache/iceberg/orc/OrcSchemaWithTypeVisitor.java +++ b/orc/src/main/java/org/apache/iceberg/orc/OrcSchemaWithTypeVisitor.java @@ -36,7 +36,7 @@ public static T visit( Type iType, TypeDescription schema, OrcSchemaWithTypeVisitor visitor) { switch (schema.getCategory()) { case STRUCT: - return visitRecord(iType != null ? iType.asStructType() : null, schema, visitor); + return visitor.visitRecord(iType != null ? iType.asStructType() : null, schema, visitor); case UNION: throw new UnsupportedOperationException("Cannot handle " + schema); @@ -61,7 +61,11 @@ public static T visit( } } - private static T visitRecord( + /** + * Visits a struct. Overridden by Spark to inject {@code initial-default} values into {@code + * idToConstant} for fields omitted from the ORC projection. + */ + protected T visitRecord( Types.StructType struct, TypeDescription record, OrcSchemaWithTypeVisitor visitor) { List fields = record.getChildren(); List names = record.getFieldNames(); diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/OrcSchemaWithTypeVisitorSpark.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/OrcSchemaWithTypeVisitorSpark.java new file mode 100644 index 000000000..95e8e2894 --- /dev/null +++ b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/OrcSchemaWithTypeVisitorSpark.java @@ -0,0 +1,122 @@ +/* + * 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; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.orc.ORCSchemaUtil; +import org.apache.iceberg.orc.OrcSchemaWithTypeVisitor; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.spark.source.BaseDataReader; +import org.apache.iceberg.types.Types; +import org.apache.orc.TypeDescription; + +/** + * Spark ORC schema visitor that injects {@code initial-default} values into {@code idToConstant} + * for fields omitted from the per-file ORC projection. + * + *

{@link org.apache.iceberg.spark.data.SparkOrcReader} uses this inject. Vectorized ORC does + * not; {@code SparkBatchScan} keeps defaulted projections on the row reader. + */ +public abstract class OrcSchemaWithTypeVisitorSpark extends OrcSchemaWithTypeVisitor { + + private final Map idToConstant; + + public Map getIdToConstant() { + return idToConstant; + } + + protected OrcSchemaWithTypeVisitorSpark(Map idToConstant) { + this.idToConstant = Maps.newHashMap(); + this.idToConstant.putAll(idToConstant); + } + + @Override + protected T visitRecord( + Types.StructType struct, TypeDescription record, OrcSchemaWithTypeVisitor visitor) { + Preconditions.checkState( + icebergFieldIdsContainOrcFieldIdsInOrder(struct, record), + "Iceberg schema and ORC schema doesn't align, please call ORCSchemaUtil.buildOrcProjection" + + " to get an aligned ORC schema first!"); + List iFields = struct.fields(); + List fields = record.getChildren(); + List names = record.getFieldNames(); + List results = Lists.newArrayListWithExpectedSize(fields.size()); + + for (int i = 0, j = 0; i < iFields.size(); i++) { + Types.NestedField iField = iFields.get(i); + TypeDescription field = j < fields.size() ? fields.get(j) : null; + if (field == null || (iField.fieldId() != ORCSchemaUtil.fieldId(field))) { + // Cases that use idToConstant for an iField: + // 1. MetadataColumns.ROW_POSITION → RowPositionReader + // 2. Partition column → ConstantReader (already in idToConstant from PartitionUtil) + // 3. Field omitted because it declares initial-default → ConstantReader (inject here) + if (MetadataColumns.nonMetadataColumn(iField.name()) + && !idToConstant.containsKey(iField.fieldId()) + && iField.initialDefault() != null) { + idToConstant.put( + iField.fieldId(), + BaseDataReader.convertConstant(iField.type(), iField.initialDefault())); + } + } else { + results.add(visit(iField.type(), field, visitor)); + j++; + } + } + return visitor.record(struct, record, names, results); + } + + private static boolean icebergFieldIdsContainOrcFieldIdsInOrder( + Types.StructType struct, TypeDescription record) { + List icebergIDList = + struct.fields().stream().map(Types.NestedField::fieldId).collect(Collectors.toList()); + List orcIDList = + record.getChildren().stream().map(ORCSchemaUtil::fieldId).collect(Collectors.toList()); + + return containsInOrder(icebergIDList, orcIDList); + } + + /** + * Checks whether {@code list1} contains all integers from {@code list2} in the same relative + * order. {@code list1} may contain extra integers that {@code list2} does not. + */ + private static boolean containsInOrder(List list1, List list2) { + if (list1.size() < list2.size()) { + return false; + } + + for (int i = 0, j = 0; j < list2.size(); j++) { + if (i >= list1.size()) { + return false; + } + while (!list1.get(i).equals(list2.get(j))) { + i++; + if (i >= list1.size()) { + return false; + } + } + i++; + } + return true; + } +} diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcReader.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcReader.java index 3d5323353..f102213de 100644 --- a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcReader.java +++ b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/data/SparkOrcReader.java @@ -25,6 +25,7 @@ import org.apache.iceberg.orc.OrcValueReader; import org.apache.iceberg.orc.OrcValueReaders; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.spark.OrcSchemaWithTypeVisitorSpark; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; import org.apache.orc.TypeDescription; @@ -64,11 +65,10 @@ public void setBatchContext(long batchOffsetInFile) { reader.setBatchContext(batchOffsetInFile); } - private static class ReadBuilder extends OrcSchemaWithTypeVisitor> { - private final Map idToConstant; + private static class ReadBuilder extends OrcSchemaWithTypeVisitorSpark> { private ReadBuilder(Map idToConstant) { - this.idToConstant = idToConstant; + super(idToConstant); } @Override @@ -77,7 +77,7 @@ public OrcValueReader record( TypeDescription record, List names, List> fields) { - return SparkOrcValueReaders.struct(record, fields, expected, idToConstant); + return SparkOrcValueReaders.struct(record, fields, expected, getIdToConstant()); } @Override diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java index f3ddd50ee..1b70db6c1 100644 --- a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java +++ b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java @@ -22,9 +22,11 @@ import java.io.IOException; import java.math.BigDecimal; import java.nio.ByteBuffer; +import java.util.Collection; import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.avro.generic.GenericData; import org.apache.avro.util.Utf8; @@ -41,6 +43,7 @@ import org.apache.iceberg.io.InputFile; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types.NestedField; @@ -48,18 +51,23 @@ import org.apache.iceberg.util.ByteBuffers; import org.apache.iceberg.util.PartitionUtil; import org.apache.spark.rdd.InputFileBlockHolder; +import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.sql.catalyst.util.ArrayBasedMapData; +import org.apache.spark.sql.catalyst.util.ArrayData; +import org.apache.spark.sql.catalyst.util.GenericArrayData; import org.apache.spark.sql.types.Decimal; import org.apache.spark.unsafe.types.UTF8String; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import scala.collection.JavaConverters; /** * Base class of Spark readers. * * @param is the Java class returned by this reader whose objects contain one or more rows. */ -abstract class BaseDataReader implements Closeable { +public abstract class BaseDataReader implements Closeable { private static final Logger LOG = LoggerFactory.getLogger(BaseDataReader.class); private final Table table; @@ -160,12 +168,56 @@ protected InputFile getInputFile(String location) { } } - protected static Object convertConstant(Type type, Object value) { + /** + * Converts a constant (partition value or initial-default) to Spark's in-memory representation. + * + *

List/map/struct branches are unused until {@code NestedField.castDefault} allows non-null + * nested types. + */ + public static Object convertConstant(Type type, Object value) { if (value == null) { return null; } switch (type.typeId()) { + case STRUCT: + StructType structType = type.asStructType(); + + if (structType.fields().isEmpty()) { + return new GenericInternalRow(); + } + + InternalRow ret = new GenericInternalRow(structType.fields().size()); + for (int i = 0; i < structType.fields().size(); i++) { + NestedField field = structType.fields().get(i); + Type fieldType = field.type(); + if (value instanceof Map) { + ret.update(i, convertConstant(field.type(), ((Map) value).get(field.name()))); + } else { + ret.update( + i, + convertConstant( + fieldType, ((StructLike) value).get(i, fieldType.typeId().javaClass()))); + } + } + return ret; + case LIST: + List javaList = + ((Collection) value) + .stream() + .map(e -> convertConstant(type.asListType().elementType(), e)) + .collect(Collectors.toList()); + return ArrayData.toArrayData( + JavaConverters.collectionAsScalaIterableConverter(javaList).asScala().toSeq()); + case MAP: + List keyList = Lists.newArrayList(); + List valueList = Lists.newArrayList(); + for (Map.Entry entry : ((Map) value).entrySet()) { + keyList.add(convertConstant(type.asMapType().keyType(), entry.getKey())); + valueList.add(convertConstant(type.asMapType().valueType(), entry.getValue())); + } + return new ArrayBasedMapData( + new GenericArrayData(keyList.toArray()), new GenericArrayData(valueList.toArray())); case DECIMAL: return Decimal.apply((BigDecimal) value); case STRING: @@ -183,25 +235,6 @@ protected static Object convertConstant(Type type, Object value) { return ByteBuffers.toByteArray((ByteBuffer) value); case BINARY: return ByteBuffers.toByteArray((ByteBuffer) value); - case STRUCT: - StructType structType = (StructType) type; - - if (structType.fields().isEmpty()) { - return new GenericInternalRow(); - } - - List fields = structType.fields(); - Object[] values = new Object[fields.size()]; - StructLike struct = (StructLike) value; - - for (int index = 0; index < fields.size(); index++) { - NestedField field = fields.get(index); - Type fieldType = field.type(); - values[index] = - convertConstant(fieldType, struct.get(index, fieldType.typeId().javaClass())); - } - - return new GenericInternalRow(values); default: } return value; diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java index f206149da..041be30f6 100644 --- a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java +++ b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java @@ -159,6 +159,7 @@ private CloseableIterable newOrcIterable( .split(task.start(), task.length()) .createReaderFunc( readOrcSchema -> new SparkOrcReader(readSchema, readOrcSchema, idToConstant)) + .supportsInitialDefaults() .filter(task.residual()) .caseSensitive(caseSensitive); diff --git a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchScan.java b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchScan.java index 63489d505..678842478 100644 --- a/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchScan.java +++ b/spark/v3.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchScan.java @@ -38,6 +38,7 @@ import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.util.PropertyUtil; import org.apache.iceberg.util.TableScanUtil; import org.apache.iceberg.util.Tasks; @@ -193,7 +194,10 @@ public PartitionReaderFactory createReaderFactory() { boolean batchReadsEnabled = batchReadsEnabled(allParquetFileScanTasks, allOrcFileScanTasks); - boolean batchReadOrc = hasNoDeleteFiles && allOrcFileScanTasks; + boolean hasNoInitialDefaults = hasNoInitialDefaults(expectedSchema); + + // Defaulted projections stay on the row reader; batched ORC does not fill initial-defaults. + boolean batchReadOrc = hasNoDeleteFiles && allOrcFileScanTasks && hasNoInitialDefaults; boolean batchReadParquet = hasNoEqDeleteFiles && allParquetFileScanTasks && atLeastOneColumn && onlyPrimitives; @@ -205,6 +209,11 @@ public PartitionReaderFactory createReaderFactory() { return new ReaderFactory(batchSize); } + static boolean hasNoInitialDefaults(Schema schema) { + return TypeUtil.indexById(schema.asStruct()).values().stream() + .noneMatch(field -> field.initialDefault() != null); + } + private boolean batchReadsEnabled(boolean isParquetOnly, boolean isOrcOnly) { if (isParquetOnly) { return readConf.parquetVectorizationEnabled(); diff --git a/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkOrcReaderForFieldsWithDefaultValue.java b/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkOrcReaderForFieldsWithDefaultValue.java new file mode 100644 index 000000000..0f6de5051 --- /dev/null +++ b/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkOrcReaderForFieldsWithDefaultValue.java @@ -0,0 +1,186 @@ +/* + * 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.apache.iceberg.spark.data.TestHelpers.assertEquals; + +import java.io.File; +import java.io.IOException; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.iceberg.Files; +import org.apache.iceberg.Schema; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.orc.ORC; +import org.apache.iceberg.types.Types; +import org.apache.orc.OrcFile; +import org.apache.orc.TypeDescription; +import org.apache.orc.Writer; +import org.apache.orc.storage.ql.exec.vector.LongColumnVector; +import org.apache.orc.storage.ql.exec.vector.VectorizedRowBatch; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +/** + * Spark ORC reads of fields with {@code initial-default}. Defaulted projections use the row reader. + * Nested-typed column defaults are not covered: {@code castDefault} rejects them. + */ +public class TestSparkOrcReaderForFieldsWithDefaultValue { + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + @Test + public void testOrcScalarDefaultValues() throws IOException { + final int numRows = 10; + + final InternalRow expectedFirstRow = new GenericInternalRow(2); + expectedFirstRow.update(0, 0); + expectedFirstRow.update(1, UTF8String.fromString("foo")); + + TypeDescription orcSchema = TypeDescription.fromString("struct"); + + Schema readSchema = + new Schema( + Types.NestedField.required(1, "col1", Types.IntegerType.get()), + Types.NestedField.optional("col2") + .withId(2) + .ofType(Types.StringType.get()) + .withInitialDefault(Expressions.lit("foo")) + .build()); + + File orcFile = writeOrcWithIntColumn(orcSchema, numRows); + + try (CloseableIterable reader = + ORC.read(Files.localInput(orcFile)) + .project(readSchema) + .createReaderFunc(readOrcSchema -> new SparkOrcReader(readSchema, readOrcSchema)) + .supportsInitialDefaults() + .build()) { + InternalRow actualFirstRow = reader.iterator().next(); + assertEquals(readSchema, expectedFirstRow, actualFirstRow); + } + } + + @Test + public void testOrcNestedScalarDefaultValues() throws IOException { + // Parent struct `loc` is present in the file; child `country` is new with an initial-default. + // If the parent itself is absent, the synthetic null parent short-circuits child constants. + final int numRows = 10; + + final InternalRow expectedLoc = new GenericInternalRow(1); + expectedLoc.update(0, UTF8String.fromString("US")); + final InternalRow expectedFirstRow = new GenericInternalRow(2); + expectedFirstRow.update(0, 0L); + expectedFirstRow.update(1, expectedLoc); + + // Empty loc struct in the file: country is absent and will be filled from initial-default. + TypeDescription orcSchema = TypeDescription.fromString("struct>"); + + Schema readSchema = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional("loc") + .withId(2) + .ofType( + Types.StructType.of( + Types.NestedField.optional("country") + .withId(3) + .ofType(Types.StringType.get()) + .withInitialDefault(Expressions.lit("US")) + .build())) + .build()); + + File orcFile = writeOrcWithIdAndEmptyLoc(orcSchema, numRows); + + try (CloseableIterable reader = + ORC.read(Files.localInput(orcFile)) + .project(readSchema) + .createReaderFunc(readOrcSchema -> new SparkOrcReader(readSchema, readOrcSchema)) + .supportsInitialDefaults() + .build()) { + InternalRow actualFirstRow = reader.iterator().next(); + assertEquals(readSchema, expectedFirstRow, actualFirstRow); + } + } + + private File writeOrcWithIntColumn(TypeDescription orcSchema, int numRows) throws IOException { + Configuration conf = new Configuration(); + File orcFile = temp.newFile(); + Path orcFilePath = new Path(orcFile.getPath()); + + Writer writer = + OrcFile.createWriter( + orcFilePath, OrcFile.writerOptions(conf).setSchema(orcSchema).overwrite(true)); + + VectorizedRowBatch batch = orcSchema.createRowBatch(); + LongColumnVector firstCol = (LongColumnVector) batch.cols[0]; + for (int r = 0; r < numRows; ++r) { + int row = batch.size++; + firstCol.vector[row] = r; + if (batch.size == batch.getMaxSize()) { + writer.addRowBatch(batch); + batch.reset(); + } + } + if (batch.size != 0) { + writer.addRowBatch(batch); + batch.reset(); + } + writer.close(); + return orcFile; + } + + private File writeOrcWithIdAndEmptyLoc(TypeDescription orcSchema, int numRows) + throws IOException { + Configuration conf = new Configuration(); + File orcFile = temp.newFile(); + Path orcFilePath = new Path(orcFile.getPath()); + + Writer writer = + OrcFile.createWriter( + orcFilePath, OrcFile.writerOptions(conf).setSchema(orcSchema).overwrite(true)); + + VectorizedRowBatch batch = orcSchema.createRowBatch(); + LongColumnVector idCol = (LongColumnVector) batch.cols[0]; + org.apache.orc.storage.ql.exec.vector.StructColumnVector locCol = + (org.apache.orc.storage.ql.exec.vector.StructColumnVector) batch.cols[1]; + for (int r = 0; r < numRows; ++r) { + int row = batch.size++; + idCol.vector[row] = r; + // non-null empty loc struct + locCol.noNulls = true; + locCol.isNull[row] = false; + if (batch.size == batch.getMaxSize()) { + writer.addRowBatch(batch); + batch.reset(); + } + } + if (batch.size != 0) { + writer.addRowBatch(batch); + batch.reset(); + } + writer.close(); + return orcFile; + } +} diff --git a/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkBatchScanInitialDefaults.java b/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkBatchScanInitialDefaults.java new file mode 100644 index 000000000..7d0fc9807 --- /dev/null +++ b/spark/v3.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkBatchScanInitialDefaults.java @@ -0,0 +1,69 @@ +/* + * 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.source; + +import org.apache.iceberg.Schema; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.types.Types; +import org.junit.Assert; +import org.junit.Test; + +public class TestSparkBatchScanInitialDefaults { + + @Test + public void testDisablesVectorizationWhenTopLevelDefaultIsProjected() { + Schema schema = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional("country") + .withId(2) + .ofType(Types.StringType.get()) + .withInitialDefault(Expressions.lit("US")) + .build()); + + Assert.assertFalse(SparkBatchScan.hasNoInitialDefaults(schema)); + } + + @Test + public void testDisablesVectorizationWhenNestedDefaultIsProjected() { + Schema schema = + new Schema( + Types.NestedField.required( + 1, + "location", + Types.StructType.of( + Types.NestedField.optional("country") + .withId(2) + .ofType(Types.StringType.get()) + .withInitialDefault(Expressions.lit("US")) + .build()))); + + Assert.assertFalse(SparkBatchScan.hasNoInitialDefaults(schema)); + } + + @Test + public void testAllowsVectorizationWhenNoDefaultIsProjected() { + Schema schema = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "country", Types.StringType.get())); + + Assert.assertTrue(SparkBatchScan.hasNoInitialDefaults(schema)); + } +}