fix: struct writer for lance v2.1+

Author: qidian99

lance-spark-base_2.12/src/main/scala/com/lancedb/lance/spark/arrow/LanceArrowFieldWriter.scala

@@ -42,10 +42,15 @@ abstract private[arrow] class LanceArrowFieldWriter {
 
   def write(input: SpecializedGetters, ordinal: Int): Unit = {
     if (input.isNullAt(ordinal)) {
-      setNull()
+      writeNull()
     } else {
       setValue(input, ordinal)
+      count += 1
     }
+  }
+
+  private[arrow] def writeNull(): Unit = {
+    setNull()
     count += 1
   }
 

lance-spark-base_2.12/src/main/scala/com/lancedb/lance/spark/arrow/LanceArrowWriter.scala

@@ -388,8 +388,18 @@ private[arrow] class MapWriter(
 private[arrow] class StructWriter(
     val valueVector: StructVector,
     val children: Array[LanceArrowFieldWriter]) extends LanceArrowFieldWriter {
-  override def setNull(): Unit = {}
+  override def setNull(): Unit = {
+    // mark the parent struct as null for this row
+    valueVector.setNull(count)
+    var i = 0
+    while (i < children.length) {
+      children(i).writeNull()
+      i += 1
+    }
+  }
   override def setValue(input: SpecializedGetters, ordinal: Int): Unit = {
+    // mark the parent struct as defined (not null) before writing children
+    valueVector.setIndexDefined(count)
     val struct = input.getStruct(ordinal, children.length)
     var i = 0
     while (i < children.length) {

lance-spark-base_2.12/src/test/java/com/lancedb/lance/spark/write/LanceArrowWriterTest.java

@@ -15,9 +15,12 @@
 
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.IntVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -33,6 +36,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class LanceArrowWriterTest {
   @Test
@@ -107,4 +111,54 @@ public void test() throws Exception {
       arrowWriter.close();
     }
   }
+
+  @Test
+  public void propagatesStructNullsToChildren() {
+    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      Field childField =
+          new Field(
+              "child", FieldType.nullable(new ArrowType.Int(32, true)), Collections.emptyList());
+      Field structField =
+          new Field(
+              "struct_col",
+              FieldType.nullable(ArrowType.Struct.INSTANCE),
+              Collections.singletonList(childField));
+      Schema arrowSchema = new Schema(Collections.singletonList(structField));
+
+      StructType childType =
+          new StructType(
+              new StructField[] {
+                DataTypes.createStructField("child", DataTypes.IntegerType, true)
+              });
+      StructType sparkSchema =
+          new StructType(
+              new StructField[] {DataTypes.createStructField("struct_col", childType, true)});
+
+      try (VectorSchemaRoot root = VectorSchemaRoot.create(arrowSchema, allocator)) {
+        com.lancedb.lance.spark.arrow.LanceArrowWriter structWriter =
+            com.lancedb.lance.spark.arrow.LanceArrowWriter$.MODULE$.create(root, sparkSchema);
+
+        InternalRow[] rows =
+            new InternalRow[] {
+              new GenericInternalRow(new Object[] {new GenericInternalRow(new Object[] {1})}),
+              new GenericInternalRow(new Object[] {null}),
+              new GenericInternalRow(new Object[] {new GenericInternalRow(new Object[] {3})})
+            };
+
+        for (InternalRow row : rows) {
+          structWriter.write(row);
+        }
+        structWriter.finish();
+
+        StructVector structVector = (StructVector) root.getVector("struct_col");
+        IntVector childVector = (IntVector) structVector.getChild("child");
+
+        assertEquals(rows.length, structVector.getValueCount());
+        assertEquals(rows.length, childVector.getValueCount());
+        assertTrue(structVector.isNull(1));
+        assertEquals(1, childVector.get(0));
+        assertEquals(3, childVector.get(2));
+      }
+    }
+  }
 }

lance-spark-base_2.12/src/test/scala/com/lancedb/lance/spark/arrow/StructWriterSuite.scala

@@ -0,0 +1,64 @@
+/*
+ * 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 com.lancedb.lance.spark.arrow
+
+import org.apache.arrow.memory.RootAllocator
+import org.apache.arrow.vector.{IntVector, VectorSchemaRoot}
+import org.apache.arrow.vector.complex.StructVector
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.types.{DataTypes, StructType}
+import org.apache.spark.sql.util.LanceArrowUtils
+import org.scalatest.funsuite.AnyFunSuite
+
+import java.time.ZoneId
+
+class StructWriterSuite extends AnyFunSuite {
+
+  test("struct null rows advance child writers") {
+    val childType = new StructType().add("child", DataTypes.IntegerType, nullable = true)
+    val schema = new StructType().add("struct_col", childType, nullable = true)
+
+    val allocator = new RootAllocator(Long.MaxValue)
+    try {
+      val arrowSchema = LanceArrowUtils.toArrowSchema(
+        schema,
+        ZoneId.systemDefault().getId,
+        errorOnDuplicatedFieldNames = true)
+      val root = VectorSchemaRoot.create(arrowSchema, allocator)
+      try {
+        val writer = LanceArrowWriter.create(root, schema)
+        val rows = Seq(
+          new GenericInternalRow(Array[Any](new GenericInternalRow(Array[Any](1)))),
+          new GenericInternalRow(Array[Any](null)),
+          new GenericInternalRow(Array[Any](new GenericInternalRow(Array[Any](3)))))
+
+        rows.foreach(writer.write)
+        writer.finish()
+
+        val structVector = root.getVector("struct_col").asInstanceOf[StructVector]
+        val childVector = structVector.getChild("child").asInstanceOf[IntVector]
+
+        assert(structVector.getValueCount === rows.length)
+        assert(childVector.getValueCount === rows.length)
+        assert(structVector.isNull(1))
+        assert(childVector.get(0) === 1)
+        assert(childVector.get(2) === 3)
+      } finally {
+        root.close()
+      }
+    } finally {
+      allocator.close()
+    }
+  }
+}