perf(python): keep bad vector filling Arrow-native

python/python/lancedb/table.py

@@ -4125,18 +4125,32 @@ def _fill_bad_vector_values(
     dim: int,
     fill_value: float,
 ) -> pa.Array:
-    if isinstance(arr, pa.ChunkedArray):
-        arr = arr.combine_chunks()
+    if not isinstance(arr, pa.ChunkedArray):
+        arr = pa.chunked_array([arr])
+    arr = arr.combine_chunks()
 
     # A fixed-size slice truncates long vectors and pads short vectors with nulls.
-    # Build a mask for only those padded positions so existing null values remain
-    # unchanged. Null vectors have a filled length of zero and are padded entirely.
+    # Slice an array marking the original child nulls in parallel so padding nulls
+    # can be distinguished from null values that were already present.
     sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True)
-    vector_lengths = pc.fill_null(pc.list_value_length(arr), 0).to_numpy(
-        zero_copy_only=False
+    child_nulls = pc.is_null(arr.values)
+    parent_nulls = pc.is_null(arr)
+    if pa.types.is_list(arr.type):
+        original_child_nulls = pa.ListArray.from_arrays(
+            arr.offsets, child_nulls, mask=parent_nulls
+        )
+    elif pa.types.is_large_list(arr.type):
+        original_child_nulls = pa.LargeListArray.from_arrays(
+            arr.offsets, child_nulls, mask=parent_nulls
+        )
+    else:
+        original_child_nulls = pa.FixedSizeListArray.from_arrays(
+            child_nulls, arr.type.list_size, mask=parent_nulls
+        )
+    sliced_child_nulls = pc.list_slice(
+        original_child_nulls, 0, dim, return_fixed_size_list=True
     )
-    value_positions = np.arange(dim, dtype=vector_lengths.dtype)
-    needs_fill = pa.array((value_positions >= vector_lengths[:, None]).reshape(-1))
+    needs_fill = pc.is_null(sliced_child_nulls.values)
 
     values = sliced.values
     if pa.types.is_floating(values.type):

python/python/tests/test_util.py

@@ -13,6 +13,7 @@ from lancedb.embeddings.registry import EmbeddingFunctionRegistry
 from lancedb.table import (
     _append_vector_columns,
     _cast_to_target_schema,
+    _fill_bad_vector_values,
     _handle_bad_vectors,
     _into_pyarrow_reader,
     _infer_target_schema,
@@ -333,6 +334,36 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
     assert output["vector"].combine_chunks() == expected
 
 
+@pytest.mark.parametrize(
+    ("vector_type", "vectors", "expected"),
+    [
+        (
+            pa.list_(pa.float64()),
+            [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
+            [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
+        ),
+        (
+            pa.large_list(pa.float64()),
+            [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
+            [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
+        ),
+        (
+            pa.list_(pa.float64(), 2),
+            [[1.0, float("nan")], None, [None, 3.0]],
+            [[1.0, 42.0], [42.0, 42.0], [None, 3.0]],
+        ),
+    ],
+)
+def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
+    arr = pa.array([[0.0, 0.0], *vectors, [9.0, 9.0]], type=vector_type)
+    arr = arr.slice(1, len(vectors))
+
+    actual = _fill_bad_vector_values(arr, dim=2, fill_value=42.0)
+
+    assert actual.type == vector_type
+    assert actual.to_pylist() == expected
+
+
 @pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
 def test_handle_bad_vectors_nan(on_bad_vectors):
     vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])