diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java index 57c700540..163919a6f 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentColumnarBatchScanner.java @@ -16,6 +16,7 @@ import org.lance.ipc.ScanStats; import org.lance.spark.LanceConstant; import org.lance.spark.read.LanceInputPartition; +import org.lance.spark.utils.FieldPathUtils; import org.lance.spark.vectorized.BlobStructAccessor; import org.lance.spark.vectorized.LanceArrowColumnVector; @@ -25,6 +26,7 @@ import org.apache.arrow.vector.complex.StructVector; import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.spark.sql.execution.vectorized.ConstantColumnVector; +import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; @@ -35,6 +37,7 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -119,11 +122,8 @@ private List buildSparkOrderedVectors( VectorSchemaRoot root, LanceInputPartition inputPartition) { StructType schema = inputPartition.getSchema(); - Map actualFields = new HashMap<>(); + Map actualFields = buildActualFieldMap(root); List rootVectors = root.getFieldVectors(); - for (int i = 0; i < rootVectors.size(); i++) { - actualFields.put(rootVectors.get(i).getField().getName(), rootVectors.get(i)); - } // Extract row addresses for blob reference support Set blobColumnNames = fragmentScanner.getBlobColumnNames(); @@ -157,16 +157,16 @@ private List buildSparkOrderedVectors( fieldVectors.add(sizeVector); } } else { - FieldVector vector = actualFields.get(fieldName); - if (vector == null) { - throw new IllegalStateException( - "Lance scan did not return expected field '" + fieldName + "'"); - } - LanceArrowColumnVector colVec = new LanceArrowColumnVector(vector, false, field); + ColumnVector colVec = + buildColumnVector( + field, Collections.singletonList(fieldName), actualFields, root.getRowCount()); // Set blob reference context so getBinary() produces blob references - if (rowAddresses != null && blobColumnNames.contains(fieldName)) { - BlobStructAccessor blobAccessor = colVec.getBlobStructAccessor(); + if (rowAddresses != null + && blobColumnNames.contains(fieldName) + && colVec instanceof LanceArrowColumnVector) { + BlobStructAccessor blobAccessor = + ((LanceArrowColumnVector) colVec).getBlobStructAccessor(); if (blobAccessor != null) { blobAccessor.setBlobReferenceContext( fragmentScanner.getDatasetUri(), fieldName, rowAddresses); @@ -179,6 +179,89 @@ private List buildSparkOrderedVectors( return fieldVectors; } + private Map buildActualFieldMap(VectorSchemaRoot root) { + Map actualFields = new HashMap<>(); + List rootVectors = root.getFieldVectors(); + for (int i = 0; i < rootVectors.size(); i++) { + actualFields.put(rootVectors.get(i).getField().getName(), rootVectors.get(i)); + } + return actualFields; + } + + private ColumnVector buildColumnVector( + StructField field, + List columnPath, + Map actualFields, + int rowCount) { + String canonicalPath = FieldPathUtils.canonicalPath(columnPath); + FieldVector exactVector = actualFields.get(canonicalPath); + if (exactVector != null) { + return new LanceArrowColumnVector(exactVector, false, field); + } + + if (field.dataType() instanceof StructType + && hasProjectedChildren(canonicalPath, actualFields)) { + StructType structType = (StructType) field.dataType(); + ColumnVector[] childVectors = new ColumnVector[structType.fields().length]; + for (int i = 0; i < structType.fields().length; i++) { + StructField childField = structType.fields()[i]; + childVectors[i] = + buildChildColumnVector( + childField, appendPath(columnPath, childField.name()), actualFields, rowCount); + } + return new ProjectedStructColumnVector(structType, childVectors); + } + + throw new IllegalStateException( + "Cannot materialize projected column '" + + canonicalPath + + "' with Spark type " + + field.dataType().catalogString() + + " from Lance output fields " + + summarizeActualFields(actualFields)); + } + + private boolean hasProjectedChildren(String columnPath, Map actualFields) { + String childPrefix = columnPath + "."; + return actualFields.keySet().stream().anyMatch(name -> name.startsWith(childPrefix)); + } + + private ColumnVector buildChildColumnVector( + StructField field, + List columnPath, + Map actualFields, + int rowCount) { + String canonicalPath = FieldPathUtils.canonicalPath(columnPath); + if (actualFields.containsKey(canonicalPath) + || hasProjectedChildren(canonicalPath, actualFields)) { + return buildColumnVector(field, columnPath, actualFields, rowCount); + } + return buildNullColumnVector(field.dataType(), rowCount); + } + + private List appendPath(List path, String fieldName) { + List childPath = new ArrayList<>(path.size() + 1); + childPath.addAll(path); + childPath.add(fieldName); + return childPath; + } + + private ColumnVector buildNullColumnVector(DataType dataType, int rowCount) { + ConstantColumnVector nullVector = new ConstantColumnVector(rowCount, dataType); + nullVector.setNull(); + return nullVector; + } + + private String summarizeActualFields(Map actualFields) { + List fieldNames = new ArrayList<>(actualFields.keySet()); + int limit = Math.min(fieldNames.size(), 8); + List preview = fieldNames.subList(0, limit); + if (fieldNames.size() <= limit) { + return preview.toString(); + } + return preview + " ... (" + fieldNames.size() + " fields total)"; + } + /** * Extracts row addresses from the {@code _rowaddr} column appended by the native scanner. Row * addresses are needed to construct blob references that allow the write side to fetch actual diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index a235a1171..aff277b71 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -30,7 +30,6 @@ import org.apache.spark.sql.types.StructType; import java.io.IOException; -import java.util.Arrays; import java.util.HashSet; import java.util.List; import java.util.Optional; @@ -109,7 +108,8 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in Set blobColumnNames = getBlobColumnNames(scanSchema); boolean hasBlobColumns = !blobColumnNames.isEmpty(); - List projectedColumns = getColumnNames(scanSchema); + List projectedColumns = + getColumnNames(scanSchema, inputPartition.getProjectedColumns()); if (projectedColumns.isEmpty() && scanSchema.isEmpty()) { scanOptions.withRowId(true); } @@ -266,25 +266,15 @@ private static Set getBlobColumnNames(StructType schema) { return blobColumns; } - private static List getColumnNames(StructType schema) { + private static List getColumnNames(StructType schema, List projectedDataColumns) { java.util.Set schemaFields = new java.util.HashSet<>(); for (StructField field : schema.fields()) { schemaFields.add(field.name()); } List columns = - Arrays.stream(schema.fields()) - .map(StructField::name) - .filter( - name -> - !name.equals(LanceConstant.FRAGMENT_ID) - && !name.equals(LanceConstant.ROW_ID) - && !name.equals(LanceConstant.ROW_ADDRESS) - && !name.equals(LanceConstant.ROW_CREATED_AT_VERSION) - && !name.equals(LanceConstant.ROW_LAST_UPDATED_AT_VERSION) - && !name.equals(LanceConstant.SCORE) - && !name.endsWith(LanceConstant.BLOB_POSITION_SUFFIX) - && !name.endsWith(LanceConstant.BLOB_SIZE_SUFFIX)) + projectedDataColumns.stream() + .filter(name -> !name.equals(LanceConstant.SCORE)) .collect(Collectors.toList()); if (schemaFields.contains(LanceConstant.ROW_LAST_UPDATED_AT_VERSION)) { columns.add(LanceConstant.ROW_LAST_UPDATED_AT_VERSION); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ProjectedStructColumnVector.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ProjectedStructColumnVector.java new file mode 100644 index 000000000..b8d05c0db --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/ProjectedStructColumnVector.java @@ -0,0 +1,133 @@ +/* + * Licensed 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.lance.spark.internal; + +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; +import org.apache.spark.unsafe.types.UTF8String; + +class ProjectedStructColumnVector extends ColumnVector { + private final ColumnVector[] childColumns; + private boolean closed; + + ProjectedStructColumnVector(StructType dataType, ColumnVector[] childColumns) { + super(dataType); + this.childColumns = childColumns; + } + + @Override + public void close() { + if (closed) { + return; + } + closed = true; + for (ColumnVector childColumn : childColumns) { + if (childColumn != null) { + childColumn.close(); + } + } + } + + @Override + public boolean hasNull() { + return false; + } + + @Override + public int numNulls() { + return 0; + } + + @Override + public boolean isNullAt(int rowId) { + // This vector is only used to provide struct child access for nested projection pushdown. + // Projection planning only reconstructs non-nullable parent structs. Nullable structs + // are kept as exact top-level projections so Spark still sees the real parent validity + // bitmap from Arrow. + return false; + } + + @Override + public boolean getBoolean(int rowId) { + throw unsupported(); + } + + @Override + public byte getByte(int rowId) { + throw unsupported(); + } + + @Override + public short getShort(int rowId) { + throw unsupported(); + } + + @Override + public int getInt(int rowId) { + throw unsupported(); + } + + @Override + public long getLong(int rowId) { + throw unsupported(); + } + + @Override + public float getFloat(int rowId) { + throw unsupported(); + } + + @Override + public double getDouble(int rowId) { + throw unsupported(); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + throw unsupported(); + } + + @Override + public UTF8String getUTF8String(int rowId) { + throw unsupported(); + } + + @Override + public byte[] getBinary(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarArray getArray(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarMap getMap(int rowId) { + throw unsupported(); + } + + @Override + public ColumnVector getChild(int ordinal) { + return childColumns[ordinal]; + } + + private UnsupportedOperationException unsupported() { + return new UnsupportedOperationException( + "ProjectedStructColumnVector only supports nested child access"); + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceInputPartition.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceInputPartition.java index db32c762c..8f2e8c79d 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceInputPartition.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceInputPartition.java @@ -22,6 +22,8 @@ import org.apache.spark.sql.connector.read.HasPartitionKey; import org.apache.spark.sql.types.StructType; +import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -29,6 +31,7 @@ public class LanceInputPartition implements HasPartitionKey { private static final long serialVersionUID = 4723894723984723985L; private final StructType schema; + private final List projectedColumns; private final int partitionId; private final LanceSplit lanceSplit; private final LanceSparkReadOptions readOptions; @@ -59,6 +62,7 @@ public class LanceInputPartition implements HasPartitionKey { public LanceInputPartition( StructType schema, + List projectedColumns, int partitionId, LanceSplit lanceSplit, LanceSparkReadOptions readOptions, @@ -73,6 +77,10 @@ public LanceInputPartition( Map namespaceProperties, InternalRow partitionKeyRow) { this.schema = schema; + this.projectedColumns = + projectedColumns == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(projectedColumns)); this.partitionId = partitionId; this.lanceSplit = lanceSplit; this.readOptions = readOptions; @@ -92,6 +100,10 @@ public StructType getSchema() { return schema; } + public List getProjectedColumns() { + return projectedColumns; + } + public int getPartitionId() { return partitionId; } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java index b8a5d1a2a..9dcbd4ef9 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScan.java @@ -48,6 +48,7 @@ import scala.collection.immutable.Map; import java.io.Serializable; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -68,6 +69,7 @@ public class LanceScan private static final Logger LOG = LoggerFactory.getLogger(LanceScan.class); private final StructType schema; + private final List projectedColumns; private final LanceSparkReadOptions readOptions; private final Optional whereConditions; private final Optional limit; @@ -136,6 +138,7 @@ public LanceScan( Optional pushedAggregation, Predicate[] pushedPredicates, LanceStatistics statistics, + List projectedColumns, java.util.Map> zonemapStats, Set survivingFragmentIds, List precomputedSplits, @@ -147,6 +150,10 @@ public LanceScan( java.util.Map namespaceProperties) { this.schema = schema; this.readOptions = readOptions; + this.projectedColumns = + projectedColumns == null + ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(projectedColumns)); this.whereConditions = whereConditions; this.limit = limit; this.offset = offset; @@ -222,6 +229,7 @@ public InputPartition[] planInputPartitions() { } return new LanceInputPartition( schema, + projectedColumns, i, split, readOptions, @@ -473,6 +481,7 @@ public boolean equals(Object o) { } LanceScan that = (LanceScan) o; return Objects.equals(schema, that.schema) + && Objects.equals(projectedColumns, that.projectedColumns) && Objects.equals(readOptions, that.readOptions) && Objects.equals(whereConditions, that.whereConditions) && Objects.equals(limit, that.limit) @@ -486,7 +495,13 @@ && aggregationEquals(pushedAggregation, that.pushedAggregation) public int hashCode() { int result = Objects.hash( - schema, readOptions, whereConditions, limit, offset, topNSortOrders.toString()); + schema, + projectedColumns, + readOptions, + whereConditions, + limit, + offset, + topNSortOrders.toString()); result = 31 * result + Arrays.hashCode(sortedByHash(pushedPredicates)); result = 31 * result + aggregationHashCode(pushedAggregation); return result; diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java index 9d858d665..ddd136d77 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java @@ -86,6 +86,7 @@ public class LanceScanBuilder private final Set blobV2Columns; private StructType schema; + private List projectedColumns; private Predicate[] pushedPredicates = new Predicate[0]; @@ -133,6 +134,8 @@ public LanceScanBuilder( this.fullSchema = BlobUtils.applyBlobV2DescriptorSchema(schema); this.blobV2Columns = BlobUtils.blobV2ColumnNames(this.fullSchema); this.schema = this.fullSchema; + this.projectedColumns = + ReadSchemaNestedColumnProjection.buildProjectedColumns(this.fullSchema, this.fullSchema); this.readOptions = readOptions; this.initialStorageOptions = initialStorageOptions; this.namespaceImpl = namespaceImpl; @@ -281,6 +284,7 @@ public Scan build() { pushedAggregation, pushedPredicates, statistics, + projectedColumns, zonemapStats, survivingFragmentIds, scanPlan.getSplits(), @@ -343,6 +347,8 @@ private Scan buildNamespaceFtsScan() { @Override public void pruneColumns(StructType requiredSchema) { this.schema = ReadSchemaNestedStructWidening.widenRequiredSchema(requiredSchema, fullSchema); + this.projectedColumns = + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema); } @Override diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/ReadSchemaNestedColumnProjection.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/ReadSchemaNestedColumnProjection.java new file mode 100644 index 000000000..cbf125914 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/ReadSchemaNestedColumnProjection.java @@ -0,0 +1,130 @@ +/* + * Licensed 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.lance.spark.read; + +import org.lance.spark.LanceConstant; +import org.lance.spark.utils.FieldPathUtils; + +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public final class ReadSchemaNestedColumnProjection { + private ReadSchemaNestedColumnProjection() {} + + public static List buildProjectedColumns( + StructType requiredSchema, StructType fullSchema) { + List projectedColumns = new ArrayList<>(); + if (requiredSchema == null) { + return projectedColumns; + } + + for (StructField requiredField : requiredSchema.fields()) { + if (isScannerSpecialField(requiredField.name())) { + continue; + } + appendProjectedColumns( + requiredField, + findField(fullSchema, requiredField.name()), + newPath(requiredField.name()), + projectedColumns); + } + return projectedColumns; + } + + private static void appendProjectedColumns( + StructField requiredField, + StructField fullField, + List columnPath, + List projectedColumns) { + if (fullField == null) { + throw new IllegalArgumentException( + "Required projection column '" + + FieldPathUtils.canonicalPath(columnPath) + + "' with type " + + requiredField.dataType().catalogString() + + " does not exist in the full schema"); + } + + if (shouldProjectNestedChildren(requiredField, fullField)) { + StructType requiredStruct = (StructType) requiredField.dataType(); + StructType fullStruct = (StructType) fullField.dataType(); + for (StructField childField : requiredStruct.fields()) { + appendProjectedColumns( + childField, + findField(fullStruct, childField.name()), + appendPath(columnPath, childField.name()), + projectedColumns); + } + return; + } + + projectedColumns.add(FieldPathUtils.canonicalPath(columnPath)); + } + + private static List newPath(String fieldName) { + List path = new ArrayList<>(); + path.add(fieldName); + return path; + } + + private static List appendPath(List path, String fieldName) { + List childPath = new ArrayList<>(path.size() + 1); + childPath.addAll(path); + childPath.add(fieldName); + return childPath; + } + + private static boolean shouldProjectNestedChildren( + StructField requiredField, StructField fullField) { + if (fullField == null) { + return false; + } + + DataType requiredType = requiredField.dataType(); + DataType fullType = fullField.dataType(); + return requiredType instanceof StructType + && fullType instanceof StructType + // Reconstructing a nullable parent struct from projected child vectors loses the + // parent validity bitmap. Keep nullable structs as top-level projections so Spark + // still observes exact parent null semantics. + && !fullField.nullable() + && !requiredType.equals(fullType); + } + + private static StructField findField(StructType schema, String fieldName) { + if (schema == null) { + return null; + } + return Arrays.stream(schema.fields()) + .filter(field -> field.name().equals(fieldName)) + .findFirst() + .orElse(null); + } + + private static boolean isScannerSpecialField(String fieldName) { + return fieldName.equals(LanceConstant.FRAGMENT_ID) + || fieldName.equals(LanceConstant.ROW_ID) + || fieldName.equals(LanceConstant.ROW_ADDRESS) + || fieldName.equals(LanceConstant.ROW_CREATED_AT_VERSION) + || fieldName.equals(LanceConstant.ROW_LAST_UPDATED_AT_VERSION) + || fieldName.equals(LanceConstant.SCORE) + || fieldName.endsWith(LanceConstant.BLOB_POSITION_SUFFIX) + || fieldName.endsWith(LanceConstant.BLOB_SIZE_SUFFIX); + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/TestUtils.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/TestUtils.java index 02c8ebe83..b5344e593 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/TestUtils.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/TestUtils.java @@ -16,6 +16,7 @@ import org.lance.namespace.LanceNamespace; import org.lance.spark.read.LanceInputPartition; import org.lance.spark.read.LanceSplit; +import org.lance.spark.read.ReadSchemaNestedColumnProjection; import org.lance.spark.utils.Optional; import org.apache.arrow.memory.BufferAllocator; @@ -82,6 +83,7 @@ public static class TestTable1Config { inputPartition = new LanceInputPartition( schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns(schema, schema), 0 /* partitionId */, new LanceSplit(Arrays.asList(0, 1)), readOptions, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java index e0230f994..488230bd1 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceFragmentScannerTest.java @@ -17,6 +17,7 @@ import org.lance.spark.LanceConstant; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.read.LanceInputPartition; +import org.lance.spark.read.ReadSchemaNestedColumnProjection; import org.lance.spark.utils.BlobUtils; import org.lance.spark.utils.Optional; @@ -44,10 +45,16 @@ public class LanceFragmentScannerTest { private List callGetColumnNames(StructType schema) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException { Method method = - LanceFragmentScanner.class.getDeclaredMethod("getColumnNames", StructType.class); + LanceFragmentScanner.class.getDeclaredMethod( + "getColumnNames", StructType.class, List.class); method.setAccessible(true); @SuppressWarnings("unchecked") - List result = (List) method.invoke(null, schema); + List result = + (List) + method.invoke( + null, + schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns(schema, schema)); return result; } @@ -249,6 +256,8 @@ public void testCreateSkipsNamespaceRebuildWhenExecutorCredentialRefreshDisabled LanceInputPartition partition = new LanceInputPartition( new StructType(), + ReadSchemaNestedColumnProjection.buildProjectedColumns( + new StructType(), new StructType()), 0, null, readOptions, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ProjectedStructColumnVectorTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ProjectedStructColumnVectorTest.java new file mode 100644 index 000000000..3a28c7a8a --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/ProjectedStructColumnVectorTest.java @@ -0,0 +1,156 @@ +/* + * Licensed 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.lance.spark.internal; + +import org.apache.spark.sql.execution.vectorized.ConstantColumnVector; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; + +public class ProjectedStructColumnVectorTest { + @Test + public void shouldNotInferParentNullnessFromProjectedChildren() { + StructType structType = + new StructType() + .add("token_total", DataTypes.LongType, true) + .add("token_prompt", DataTypes.LongType, true); + ConstantColumnVector tokenTotal = new ConstantColumnVector(2, DataTypes.LongType); + tokenTotal.setNull(); + ConstantColumnVector tokenPrompt = new ConstantColumnVector(2, DataTypes.LongType); + tokenPrompt.setNull(); + + try (ProjectedStructColumnVector vector = + new ProjectedStructColumnVector(structType, new ColumnVector[] {tokenTotal, tokenPrompt})) { + assertFalse(vector.isNullAt(0)); + assertFalse(vector.hasNull()); + assertSame(tokenTotal, vector.getChild(0)); + assertSame(tokenPrompt, vector.getChild(1)); + } + } + + @Test + public void shouldCloseChildrenOnlyOnce() { + StructType structType = new StructType().add("token_total", DataTypes.LongType, true); + TrackingColumnVector child = new TrackingColumnVector(DataTypes.LongType); + + ProjectedStructColumnVector vector = + new ProjectedStructColumnVector(structType, new ColumnVector[] {child}); + vector.close(); + vector.close(); + + assertSame(child, vector.getChild(0)); + org.junit.jupiter.api.Assertions.assertEquals(1, child.closeCalls); + } + + private static final class TrackingColumnVector extends ColumnVector { + private int closeCalls; + + private TrackingColumnVector(org.apache.spark.sql.types.DataType dataType) { + super(dataType); + } + + @Override + public void close() { + closeCalls++; + } + + @Override + public boolean hasNull() { + throw unsupported(); + } + + @Override + public int numNulls() { + throw unsupported(); + } + + @Override + public boolean isNullAt(int rowId) { + throw unsupported(); + } + + @Override + public boolean getBoolean(int rowId) { + throw unsupported(); + } + + @Override + public byte getByte(int rowId) { + throw unsupported(); + } + + @Override + public short getShort(int rowId) { + throw unsupported(); + } + + @Override + public int getInt(int rowId) { + throw unsupported(); + } + + @Override + public long getLong(int rowId) { + throw unsupported(); + } + + @Override + public float getFloat(int rowId) { + throw unsupported(); + } + + @Override + public double getDouble(int rowId) { + throw unsupported(); + } + + @Override + public org.apache.spark.sql.types.Decimal getDecimal(int rowId, int precision, int scale) { + throw unsupported(); + } + + @Override + public org.apache.spark.unsafe.types.UTF8String getUTF8String(int rowId) { + throw unsupported(); + } + + @Override + public byte[] getBinary(int rowId) { + throw unsupported(); + } + + @Override + public org.apache.spark.sql.vectorized.ColumnarArray getArray(int rowId) { + throw unsupported(); + } + + @Override + public org.apache.spark.sql.vectorized.ColumnarMap getMap(int rowId) { + throw unsupported(); + } + + @Override + public ColumnVector getChild(int ordinal) { + throw unsupported(); + } + + private UnsupportedOperationException unsupported() { + return new UnsupportedOperationException("TrackingColumnVector is only used for close()"); + } + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/BaseSparkConnectorReadTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/BaseSparkConnectorReadTest.java index dadf649fe..20d5af5c2 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/BaseSparkConnectorReadTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/BaseSparkConnectorReadTest.java @@ -400,6 +400,114 @@ public void testFilterWithSingleQuoteInStringValue() { assertEquals("O'Brien", rows.get(0).getString(1)); } + @Test + public void nestedStructAggregationShouldReturnConsistentResultWithoutNestedSchemaPruning() { + StructType usageMetricsSchema = + new StructType() + .add("token_prompt", DataTypes.LongType, true) + .add("token_total", DataTypes.LongType, true) + .add("token_completion", DataTypes.LongType, true); + StructType schema = + new StructType() + .add("id", DataTypes.IntegerType, false) + .add("business_domain", DataTypes.StringType, false) + .add("usage_metrics", usageMetricsSchema, true); + + List testData = + Arrays.asList( + RowFactory.create(1, "public", RowFactory.create(10L, 10L, 0L)), + RowFactory.create(2, "public", RowFactory.create(20L, 20L, 0L)), + RowFactory.create(3, "public", RowFactory.create(30L, 30L, 0L)), + RowFactory.create(4, "private", RowFactory.create(40L, 40L, 0L))); + + Dataset df = spark.createDataFrame(testData, schema); + String datasetPath = tempDir.toString() + "/nested_struct_aggregation_test"; + df.write().format(LanceDataSource.name).save(datasetPath); + + Dataset lanceData = spark.read().format(LanceDataSource.name).load(datasetPath); + lanceData.createOrReplaceTempView("nested_struct_aggregation_test"); + + String sql = + "SELECT COUNT(1), SUM(usage_metrics.token_total) " + + "FROM nested_struct_aggregation_test " + + "WHERE business_domain = 'public'"; + + String originalNestedSchemaPruning = + spark.conf().get("spark.sql.optimizer.nestedSchemaPruning.enabled"); + try { + spark.conf().set("spark.sql.optimizer.nestedSchemaPruning.enabled", "false"); + Row withoutPruning = spark.sql(sql).collectAsList().get(0); + + spark.conf().set("spark.sql.optimizer.nestedSchemaPruning.enabled", "true"); + Row withPruning = spark.sql(sql).collectAsList().get(0); + + assertEquals(3L, withoutPruning.getLong(0)); + assertEquals(60L, withoutPruning.getLong(1)); + assertEquals(withoutPruning.getLong(0), withPruning.getLong(0)); + assertEquals(withoutPruning.getLong(1), withPruning.getLong(1)); + } finally { + spark + .conf() + .set("spark.sql.optimizer.nestedSchemaPruning.enabled", originalNestedSchemaPruning); + } + } + + @Test + public void nestedStructNullnessShouldRemainCorrectWhenProjectingNullableSubfield() { + StructType usageMetricsSchema = + new StructType() + .add("token_prompt", DataTypes.LongType, true) + .add("token_total", DataTypes.LongType, true); + StructType schema = + new StructType() + .add("id", DataTypes.IntegerType, false) + .add("usage_metrics", usageMetricsSchema, true); + + List testData = + Arrays.asList( + RowFactory.create(1, null), + RowFactory.create(2, RowFactory.create(7L, null)), + RowFactory.create(3, RowFactory.create(9L, 10L))); + + Dataset df = spark.createDataFrame(testData, schema); + String datasetPath = tempDir.toString() + "/nested_struct_nullness_projection_test"; + df.write().format(LanceDataSource.name).save(datasetPath); + + String originalNestedSchemaPruning = + spark.conf().get("spark.sql.optimizer.nestedSchemaPruning.enabled"); + try { + spark.conf().set("spark.sql.optimizer.nestedSchemaPruning.enabled", "true"); + + List result = + spark + .read() + .format(LanceDataSource.name) + .load(datasetPath) + .selectExpr( + "id", "usage_metrics IS NULL AS metrics_is_null", "usage_metrics.token_total") + .orderBy("id") + .collectAsList(); + + assertEquals(3, result.size()); + + assertEquals(1, result.get(0).getInt(0)); + assertTrue(result.get(0).getBoolean(1)); + assertTrue(result.get(0).isNullAt(2)); + + assertEquals(2, result.get(1).getInt(0)); + assertFalse(result.get(1).getBoolean(1)); + assertTrue(result.get(1).isNullAt(2)); + + assertEquals(3, result.get(2).getInt(0)); + assertFalse(result.get(2).getBoolean(1)); + assertEquals(10L, result.get(2).getLong(2)); + } finally { + spark + .conf() + .set("spark.sql.optimizer.nestedSchemaPruning.enabled", originalNestedSchemaPruning); + } + } + @Test public void testArrayMaxOnNestedStructField() { // Create a schema with nested struct containing an array: diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceColumnarPartitionReaderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceColumnarPartitionReaderTest.java index 03a47774d..55794bae5 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceColumnarPartitionReaderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceColumnarPartitionReaderTest.java @@ -36,6 +36,8 @@ public void test() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0 /* partitionId */, split, TestUtils.TestTable1Config.readOptions, @@ -79,6 +81,8 @@ public void testOffsetAndLimit() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0 /* partitionId */, split, TestUtils.TestTable1Config.readOptions, @@ -124,6 +128,8 @@ public void testTopN() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0 /* partitionId */, split, TestUtils.TestTable1Config.readOptions, @@ -168,6 +174,8 @@ public void testCloseAfterMidIterationCreateFailure() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0 /* partitionId */, split, TestUtils.TestTable1Config.readOptions, @@ -209,6 +217,8 @@ public void testCloseIsIdempotent() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0 /* partitionId */, split, TestUtils.TestTable1Config.readOptions, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java index 23135af59..7f48363b4 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceCountStarPartitionReaderTest.java @@ -40,6 +40,8 @@ public void testCloseReleasesArrowMemory() throws Exception { LanceInputPartition partition = new LanceInputPartition( TestUtils.TestTable1Config.schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns( + TestUtils.TestTable1Config.schema, TestUtils.TestTable1Config.schema), 0, new LanceSplit(Arrays.asList(0, 1)), TestUtils.TestTable1Config.readOptions, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceDatasetReadTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceDatasetReadTest.java index 2a8c6670c..3e047c29f 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceDatasetReadTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceDatasetReadTest.java @@ -116,6 +116,7 @@ public void validateFragment(List> expectedValues, int fragment, St fragment, new LanceInputPartition( schema, + ReadSchemaNestedColumnProjection.buildProjectedColumns(schema, schema), 0 /* partitionId */, new LanceSplit(Arrays.asList(fragment)), TestUtils.TestTable1Config.readOptions, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanTest.java index 78291ff96..8473180e7 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanTest.java @@ -210,6 +210,7 @@ public void testOutputPartitioningWithPartitionInfo() { org.lance.spark.utils.Optional.empty(), new Predicate[0], null, + ReadSchemaNestedColumnProjection.buildProjectedColumns(TEST_SCHEMA, TEST_SCHEMA), Collections.emptyMap(), null, plan.getSplits(), @@ -295,6 +296,7 @@ public void testOutputPartitioningWithBucketInfo() { org.lance.spark.utils.Optional.empty(), new Predicate[0], null, + ReadSchemaNestedColumnProjection.buildProjectedColumns(TEST_SCHEMA, TEST_SCHEMA), Collections.emptyMap(), null, bucketPlan.getSplits(), @@ -361,4 +363,60 @@ public void testNotEqualWithDifferentSchema() { assertNotEquals(scan1, scan2, "Scans with different schemas should not be equal"); } + + @Test + public void testNotEqualWithDifferentProjectedColumns() { + StructType schema = + new StructType() + .add( + "usage_metrics", + new StructType() + .add("token_total", DataTypes.LongType) + .add("token_prompt", DataTypes.LongType)); + LanceScan scan1 = + new LanceScan( + schema, + TestUtils.TestTable1Config.readOptions, + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + new Predicate[0], + null, + Collections.singletonList("usage_metrics.token_total"), + Collections.emptyMap(), + null, + null, + Collections.emptyMap(), + null, + null, + Collections.emptyMap(), + null, + Collections.emptyMap()); + LanceScan scan2 = + new LanceScan( + schema, + TestUtils.TestTable1Config.readOptions, + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + org.lance.spark.utils.Optional.empty(), + new Predicate[0], + null, + Collections.singletonList("usage_metrics.token_prompt"), + Collections.emptyMap(), + null, + null, + Collections.emptyMap(), + null, + null, + Collections.emptyMap(), + null, + Collections.emptyMap()); + + assertNotEquals( + scan1, scan2, "Scans with different nested projected columns should not be equal"); + } } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/ReadSchemaNestedColumnProjectionTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/ReadSchemaNestedColumnProjectionTest.java new file mode 100644 index 000000000..8681cdc34 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/ReadSchemaNestedColumnProjectionTest.java @@ -0,0 +1,203 @@ +/* + * Licensed 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.lance.spark.read; + +import org.lance.spark.LanceConstant; + +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +public class ReadSchemaNestedColumnProjectionTest { + @Test + public void shouldReturnEmptyProjectionForNullRequiredSchema() { + StructType fullSchema = new StructType().add("record_id", DataTypes.IntegerType); + + assertEquals( + Collections.emptyList(), + ReadSchemaNestedColumnProjection.buildProjectedColumns(null, fullSchema)); + } + + @Test + public void shouldProjectNestedStructSubsetAsLeafPaths() { + StructType fullSchema = + new StructType() + .add("record_id", DataTypes.IntegerType) + .add( + "usage_metrics", + new StructType() + .add("auxiliary_count", DataTypes.LongType) + .add("token_total", DataTypes.LongType), + false) + .add("business_domain", DataTypes.StringType); + StructType requiredSchema = + new StructType() + .add("usage_metrics", new StructType().add("token_total", DataTypes.LongType), false) + .add("business_domain", DataTypes.StringType); + + assertEquals( + Arrays.asList("usage_metrics.token_total", "business_domain"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + } + + @Test + public void shouldProjectDeeplyNestedStructSubsetAsLeafPaths() { + StructType fullSchema = + new StructType() + .add( + "profile", + new StructType() + .add( + "usage_metrics", + new StructType() + .add("token_prompt", DataTypes.LongType) + .add("token_total", DataTypes.LongType), + false) + .add("region", DataTypes.StringType), + false); + StructType requiredSchema = + new StructType() + .add( + "profile", + new StructType() + .add( + "usage_metrics", + new StructType().add("token_total", DataTypes.LongType), + false), + false); + + assertEquals( + Collections.singletonList("profile.usage_metrics.token_total"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + } + + @Test + public void shouldQuoteNestedFieldPathParts() { + StructType fullSchema = + new StructType() + .add( + "usage.metrics", + new StructType() + .add("token.total", DataTypes.LongType) + .add("token_prompt", DataTypes.LongType), + false); + StructType requiredSchema = + new StructType() + .add("usage.metrics", new StructType().add("token.total", DataTypes.LongType), false); + + assertEquals( + Collections.singletonList("`usage.metrics`.`token.total`"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + } + + @Test + public void shouldKeepFullStructProjectionAsTopLevelColumn() { + StructType fullSchema = + new StructType() + .add("record_id", DataTypes.IntegerType) + .add( + "usage_metrics", + new StructType() + .add("auxiliary_count", DataTypes.LongType) + .add("token_total", DataTypes.LongType)); + + assertEquals( + Arrays.asList("record_id", "usage_metrics"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(fullSchema, fullSchema)); + } + + @Test + public void shouldKeepNullableStructSubsetAsTopLevelProjection() { + StructType fullSchema = + new StructType() + .add("record_id", DataTypes.IntegerType) + .add( + "usage_metrics", + new StructType() + .add("auxiliary_count", DataTypes.LongType) + .add("token_total", DataTypes.LongType), + true); + StructType requiredSchema = + new StructType() + .add("usage_metrics", new StructType().add("token_total", DataTypes.LongType), true); + + assertEquals( + Collections.singletonList("usage_metrics"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + } + + @Test + public void shouldLeaveArrayOfStructAsTopLevelProjection() { + // Nested child pruning currently only expands plain StructType projections. Array/Map + // containers remain top-level so the scanner behavior stays explicit and predictable. + StructType fullSchema = + new StructType() + .add( + "event_groups", + DataTypes.createArrayType( + new StructType() + .add("reserved_count", DataTypes.LongType) + .add("total_events", DataTypes.LongType))); + StructType requiredSchema = + new StructType() + .add( + "event_groups", + DataTypes.createArrayType( + new StructType().add("total_events", DataTypes.LongType))); + + assertEquals( + Collections.singletonList("event_groups"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + } + + @Test + public void shouldSkipScannerSpecialFields() { + StructType schema = + new StructType() + .add("record_id", DataTypes.IntegerType) + .add(LanceConstant.FRAGMENT_ID, DataTypes.IntegerType) + .add(LanceConstant.ROW_ID, DataTypes.LongType) + .add(LanceConstant.ROW_ADDRESS, DataTypes.LongType) + .add(LanceConstant.ROW_CREATED_AT_VERSION, DataTypes.LongType) + .add(LanceConstant.ROW_LAST_UPDATED_AT_VERSION, DataTypes.LongType) + .add("payload" + LanceConstant.BLOB_POSITION_SUFFIX, DataTypes.LongType) + .add("payload" + LanceConstant.BLOB_SIZE_SUFFIX, DataTypes.LongType); + + assertEquals( + Collections.singletonList("record_id"), + ReadSchemaNestedColumnProjection.buildProjectedColumns(schema, schema)); + } + + @Test + public void shouldFailWhenRequiredFieldIsMissingFromFullSchema() { + StructType fullSchema = new StructType().add("record_id", DataTypes.IntegerType); + StructType requiredSchema = new StructType().add("missing_field", DataTypes.StringType); + + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + ReadSchemaNestedColumnProjection.buildProjectedColumns(requiredSchema, fullSchema)); + + assertEquals( + "Required projection column 'missing_field' with type string does not exist in the full schema", + error.getMessage()); + } +}