feat: push down nested struct field projections

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<ColumnVector> buildSparkOrderedVectors(
       VectorSchemaRoot root, LanceInputPartition inputPartition) {
     StructType schema = inputPartition.getSchema();
 
-    Map<String, FieldVector> actualFields = new HashMap<>();
+    Map<String, FieldVector> actualFields = buildActualFieldMap(root);
     List<FieldVector> 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<String> blobColumnNames = fragmentScanner.getBlobColumnNames();
@@ -157,16 +157,16 @@ private List<ColumnVector> 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<ColumnVector> buildSparkOrderedVectors(
     return fieldVectors;
   }
 
+  private Map<String, FieldVector> buildActualFieldMap(VectorSchemaRoot root) {
+    Map<String, FieldVector> actualFields = new HashMap<>();
+    List<FieldVector> 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<String> columnPath,
+      Map<String, FieldVector> 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<String, FieldVector> actualFields) {
+    String childPrefix = columnPath + ".";
+    return actualFields.keySet().stream().anyMatch(name -> name.startsWith(childPrefix));
+  }
+
+  private ColumnVector buildChildColumnVector(
+      StructField field,
+      List<String> columnPath,
+      Map<String, FieldVector> 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<String> appendPath(List<String> path, String fieldName) {
+    List<String> 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<String, FieldVector> actualFields) {
+    List<String> fieldNames = new ArrayList<>(actualFields.keySet());
+    int limit = Math.min(fieldNames.size(), 8);
+    List<String> 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

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<String> blobColumnNames = getBlobColumnNames(scanSchema);
       boolean hasBlobColumns = !blobColumnNames.isEmpty();
 
-      List<String> projectedColumns = getColumnNames(scanSchema);
+      List<String> projectedColumns =
+          getColumnNames(scanSchema, inputPartition.getProjectedColumns());
       if (projectedColumns.isEmpty() && scanSchema.isEmpty()) {
         scanOptions.withRowId(true);
       }
@@ -266,25 +266,15 @@ private static Set<String> getBlobColumnNames(StructType schema) {
     return blobColumns;
   }
 
-  private static List<String> getColumnNames(StructType schema) {
+  private static List<String> getColumnNames(StructType schema, List<String> projectedDataColumns) {
     java.util.Set<String> schemaFields = new java.util.HashSet<>();
     for (StructField field : schema.fields()) {
       schemaFields.add(field.name());
     }
 
     List<String> 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);

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");
+  }
+}

lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceInputPartition.java

@@ -22,13 +22,16 @@
 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;
 
 public class LanceInputPartition implements HasPartitionKey {
   private static final long serialVersionUID = 4723894723984723985L;
 
   private final StructType schema;
+  private final List<String> 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<String> projectedColumns,
       int partitionId,
       LanceSplit lanceSplit,
       LanceSparkReadOptions readOptions,
@@ -73,6 +77,10 @@ public LanceInputPartition(
       Map<String, String> 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<String> getProjectedColumns() {
+    return projectedColumns;
+  }
+
   public int getPartitionId() {
     return partitionId;
   }