diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 8243c06cf..05004eafa 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -4154,17 +4154,58 @@ def _handle_bad_vector_column( raise ValueError( "`fill_value` must not be None if `on_bad_vectors` is 'fill'" ) - vec_arr = pc.if_else( - is_bad, - pa.scalar([fill_value] * dim, type=vec_arr.type), - vec_arr, - ) + vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value) else: raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}") return data.set_column(position, vector_column_name, vec_arr) +def _fill_bad_vector_values( + arr: Union[pa.Array, pa.ChunkedArray], + dim: int, + fill_value: float, +) -> pa.Array: + 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. + # 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) + 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 + ) + needs_fill = pc.is_null(sliced_child_nulls.values) + + values = sliced.values + if pa.types.is_floating(values.type): + values_for_nan_check = ( + values.cast(pa.float32()) if pa.types.is_float16(values.type) else values + ) + needs_fill = pc.or_kleene(needs_fill, pc.is_nan(values_for_nan_check)) + + fill_scalar = pa.scalar(fill_value).cast(values.type) + filled_values = pc.if_else(needs_fill, fill_scalar, values) + filled = pa.FixedSizeListArray.from_arrays(filled_values, dim) + return filled.cast(arr.type) + + def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray: if isinstance(arr, pa.ChunkedArray): values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks]) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 21cfb3df8..ac331838e 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1611,16 +1611,23 @@ def test_create_with_nans(mem_db: DBConnection): "fill_test", data=[ {"vector": [3.1, 4.1], "item": "foo", "price": 10.0}, + {"vector": [2.1, 4.1], "item": "foo", "price": 9.0}, {"vector": [np.nan], "item": "bar", "price": 20.0}, - {"vector": [np.nan, np.nan], "item": "bar", "price": 20.0}, + {"vector": [np.nan, 5.0], "item": "bar", "price": 21.0}, + {"vector": [5], "item": "bar", "price": 22.0}, ], on_bad_vectors="fill", fill_value=0.0, ) - assert len(table) == 3 + assert len(table) == 5 arrow_tbl = table.search().where("item == 'bar'").to_arrow() - v = arrow_tbl["vector"].to_pylist()[0] - assert np.allclose(v, np.array([0.0, 0.0])) + filled_vectors = { + row["price"]: row["vector"] + for row in arrow_tbl.select(["price", "vector"]).to_pylist() + } + assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0])) + assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0])) + assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) def test_add_with_nans(mem_db: DBConnection): @@ -1663,15 +1670,21 @@ def test_add_with_nans(mem_db: DBConnection): data=[ {"vector": [3.1, 4.1], "item": "foo", "price": 10.0}, {"vector": [np.nan], "item": "bar", "price": 20.0}, - {"vector": [np.nan, np.nan], "item": "bar", "price": 20.0}, + {"vector": [np.nan, 5.0], "item": "bar", "price": 21.0}, + {"vector": [5], "item": "bar", "price": 22.0}, ], on_bad_vectors="fill", fill_value=0.0, ) - assert len(table) == 3 + assert len(table) == 4 arrow_tbl = table.search().where("item == 'bar'").to_arrow() - v = arrow_tbl["vector"].to_pylist()[0] - assert np.allclose(v, np.array([0.0, 0.0])) + filled_vectors = { + row["price"]: row["vector"] + for row in arrow_tbl.select(["price", "vector"]).to_pylist() + } + assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0])) + assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0])) + assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection): @@ -1832,7 +1845,9 @@ def test_on_bad_vectors_fill_preserves_arrow_nested_vector_type(mem_db: DBConnec fill_value=0.0, ) - assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0], [0.0, 0.0]] + vector = table.to_arrow()["vector"] + assert vector.type == pa.list_(pa.float32()) + assert vector.to_pylist() == [[1.0, 2.0], [0.0, 3.0]] @pytest.mark.parametrize( diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index da33cc77f..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, @@ -287,7 +288,9 @@ def test_append_vector_columns(): @pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) def test_handle_bad_vectors_jagged(on_bad_vectors): - vector = pa.array([[1.0, 2.0], [3.0], [4.0, 5.0]]) + vector = pa.array( + [[1.0, 2.0], [3.0], [4.0, 5.0], [6.0, 7.0, 8.0], [None, 9.0], None] + ) schema = pa.schema({"vector": pa.list_(pa.float64())}) data = pa.table({"vector": vector}, schema=schema) @@ -313,15 +316,54 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ).read_all() if on_bad_vectors == "drop": - expected = pa.array([[1.0, 2.0], [4.0, 5.0]]) + expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]]) elif on_bad_vectors == "fill": - expected = pa.array([[1.0, 2.0], [42.0, 42.0], [4.0, 5.0]]) + expected = pa.array( + [ + [1.0, 2.0], + [3.0, 42.0], + [4.0, 5.0], + [6.0, 7.0], + [None, 9.0], + [42.0, 42.0], + ] + ) elif on_bad_vectors == "null": - expected = pa.array([[1.0, 2.0], None, [4.0, 5.0]]) + expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None]) 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]]) @@ -351,7 +393,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors): if on_bad_vectors == "drop": expected = pa.array([[3.0, 4.0]]) elif on_bad_vectors == "fill": - expected = pa.array([[42.0, 42.0], [3.0, 4.0]]) + expected = pa.array([[1.0, 42.0], [3.0, 4.0]]) elif on_bad_vectors == "null": expected = pa.array([None, [3.0, 4.0]])