diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 552552279..5b95315c5 100644 --- a/python/python/lancedb/table.py +++ b/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): diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index cc64b1379..7cb0ce448 100644 --- a/python/python/tests/test_util.py +++ b/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]])