mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-12 08:12:28 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3d562a7c28 |
+127
-34
@@ -4035,7 +4035,10 @@ def _handle_bad_vectors(
|
|||||||
for vector_column in vector_columns:
|
for vector_column in vector_columns:
|
||||||
dim = vector_column["expected_dim"]
|
dim = vector_column["expected_dim"]
|
||||||
if target_schema is not None and dim is None:
|
if target_schema is not None and dim is None:
|
||||||
dim = _infer_vector_dim(batch[vector_column["name"]])
|
dim = _infer_vector_column_dim(
|
||||||
|
batch[vector_column["name"]],
|
||||||
|
vector_column["is_multivector"],
|
||||||
|
)
|
||||||
pending_dims.append(vector_column)
|
pending_dims.append(vector_column)
|
||||||
batch = _handle_bad_vector_column(
|
batch = _handle_bad_vector_column(
|
||||||
batch,
|
batch,
|
||||||
@@ -4044,11 +4047,13 @@ def _handle_bad_vectors(
|
|||||||
fill_value=fill_value,
|
fill_value=fill_value,
|
||||||
expected_dim=dim,
|
expected_dim=dim,
|
||||||
expected_value_type=vector_column["expected_value_type"],
|
expected_value_type=vector_column["expected_value_type"],
|
||||||
|
is_multivector=vector_column["is_multivector"],
|
||||||
)
|
)
|
||||||
for vector_column in pending_dims:
|
for vector_column in pending_dims:
|
||||||
if vector_column["expected_dim"] is None:
|
if vector_column["expected_dim"] is None:
|
||||||
vector_column["expected_dim"] = _infer_vector_dim(
|
vector_column["expected_dim"] = _infer_vector_column_dim(
|
||||||
batch[vector_column["name"]]
|
batch[vector_column["name"]],
|
||||||
|
vector_column["is_multivector"],
|
||||||
)
|
)
|
||||||
if batch.schema.equals(output_schema, check_metadata=True):
|
if batch.schema.equals(output_schema, check_metadata=True):
|
||||||
yield batch
|
yield batch
|
||||||
@@ -4074,22 +4079,30 @@ def _find_vector_columns(
|
|||||||
if target_schema is None:
|
if target_schema is None:
|
||||||
vector_columns = []
|
vector_columns = []
|
||||||
for field in reader_schema:
|
for field in reader_schema:
|
||||||
named_vector_col = (
|
is_multivector = _is_multivector_type(field.type)
|
||||||
_is_list_like(field.type)
|
is_fixed_multivector = is_multivector and pa.types.is_fixed_size_list(
|
||||||
and pa.types.is_floating(field.type.value_type)
|
field.type.value_type
|
||||||
and field.name == VECTOR_COLUMN_NAME
|
|
||||||
)
|
)
|
||||||
|
named_vector_col = (
|
||||||
|
_is_float_vector_type(field.type) or is_multivector
|
||||||
|
) and field.name == VECTOR_COLUMN_NAME
|
||||||
likely_vector_col = (
|
likely_vector_col = (
|
||||||
pa.types.is_fixed_size_list(field.type)
|
pa.types.is_fixed_size_list(field.type)
|
||||||
and pa.types.is_floating(field.type.value_type)
|
and pa.types.is_floating(field.type.value_type)
|
||||||
and (field.type.list_size >= 10)
|
and (field.type.list_size >= 10)
|
||||||
)
|
)
|
||||||
if named_vector_col or likely_vector_col:
|
if named_vector_col or likely_vector_col or is_fixed_multivector:
|
||||||
|
vector_type = field.type.value_type if is_multivector else field.type
|
||||||
vector_columns.append(
|
vector_columns.append(
|
||||||
{
|
{
|
||||||
"name": field.name,
|
"name": field.name,
|
||||||
"expected_dim": None,
|
"expected_dim": (
|
||||||
"expected_value_type": None,
|
vector_type.list_size
|
||||||
|
if pa.types.is_fixed_size_list(vector_type)
|
||||||
|
else None
|
||||||
|
),
|
||||||
|
"expected_value_type": vector_type.value_type,
|
||||||
|
"is_multivector": is_multivector,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
return vector_columns
|
return vector_columns
|
||||||
@@ -4103,9 +4116,8 @@ def _find_vector_columns(
|
|||||||
for field in target_schema:
|
for field in target_schema:
|
||||||
if field.name not in reader_column_names:
|
if field.name not in reader_column_names:
|
||||||
continue
|
continue
|
||||||
if not _is_list_like(field.type) or not pa.types.is_floating(
|
is_multivector = _is_multivector_type(field.type)
|
||||||
field.type.value_type
|
if not _is_float_vector_type(field.type) and not is_multivector:
|
||||||
):
|
|
||||||
continue
|
continue
|
||||||
|
|
||||||
reader_field = reader_schema.field(field.name)
|
reader_field = reader_schema.field(field.name)
|
||||||
@@ -4120,16 +4132,18 @@ def _find_vector_columns(
|
|||||||
and reader_field.type.list_size >= 10
|
and reader_field.type.list_size >= 10
|
||||||
)
|
)
|
||||||
|
|
||||||
if named_vector_col or typed_fixed_vector_col:
|
if named_vector_col or typed_fixed_vector_col or is_multivector:
|
||||||
|
vector_type = field.type.value_type if is_multivector else field.type
|
||||||
vector_columns.append(
|
vector_columns.append(
|
||||||
{
|
{
|
||||||
"name": field.name,
|
"name": field.name,
|
||||||
"expected_dim": (
|
"expected_dim": (
|
||||||
field.type.list_size
|
vector_type.list_size
|
||||||
if pa.types.is_fixed_size_list(field.type)
|
if pa.types.is_fixed_size_list(vector_type)
|
||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
"expected_value_type": field.type.value_type,
|
"expected_value_type": vector_type.value_type,
|
||||||
|
"is_multivector": is_multivector,
|
||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -4180,6 +4194,7 @@ def _handle_bad_vector_column(
|
|||||||
fill_value: float = 0.0,
|
fill_value: float = 0.0,
|
||||||
expected_dim: Optional[int] = None,
|
expected_dim: Optional[int] = None,
|
||||||
expected_value_type: Optional[pa.DataType] = None,
|
expected_value_type: Optional[pa.DataType] = None,
|
||||||
|
is_multivector: bool = False,
|
||||||
) -> pa.RecordBatch:
|
) -> pa.RecordBatch:
|
||||||
"""
|
"""
|
||||||
Ensure that the vector column exists and has type fixed_size_list(float)
|
Ensure that the vector column exists and has type fixed_size_list(float)
|
||||||
@@ -4200,6 +4215,8 @@ def _handle_bad_vector_column(
|
|||||||
vec_arr = data[vector_column_name]
|
vec_arr = data[vector_column_name]
|
||||||
if not _is_list_like(vec_arr.type):
|
if not _is_list_like(vec_arr.type):
|
||||||
return data
|
return data
|
||||||
|
if is_multivector and not _is_multivector_type(vec_arr.type):
|
||||||
|
return data
|
||||||
|
|
||||||
if (
|
if (
|
||||||
expected_dim is not None
|
expected_dim is not None
|
||||||
@@ -4216,12 +4233,21 @@ def _handle_bad_vector_column(
|
|||||||
vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type))
|
vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type))
|
||||||
data = data.set_column(position, vector_column_name, vec_arr)
|
data = data.set_column(position, vector_column_name, vec_arr)
|
||||||
|
|
||||||
if pa.types.is_floating(vec_arr.type.value_type):
|
if is_multivector or pa.types.is_floating(vec_arr.type.value_type):
|
||||||
has_nan = has_nan_values(vec_arr)
|
has_nan = has_nan_values(vec_arr)
|
||||||
else:
|
else:
|
||||||
has_nan = pa.array([False] * len(vec_arr))
|
has_nan = pa.array([False] * len(vec_arr))
|
||||||
|
|
||||||
if expected_dim is not None:
|
if is_multivector:
|
||||||
|
dim = (
|
||||||
|
expected_dim
|
||||||
|
if expected_dim is not None
|
||||||
|
else _infer_vector_column_dim(vec_arr, True)
|
||||||
|
)
|
||||||
|
if dim is None:
|
||||||
|
return data
|
||||||
|
has_wrong_dim = _multivector_has_wrong_dim(vec_arr, dim)
|
||||||
|
elif expected_dim is not None:
|
||||||
dim = expected_dim
|
dim = expected_dim
|
||||||
elif pa.types.is_fixed_size_list(vec_arr.type):
|
elif pa.types.is_fixed_size_list(vec_arr.type):
|
||||||
dim = vec_arr.type.list_size
|
dim = vec_arr.type.list_size
|
||||||
@@ -4230,15 +4256,16 @@ def _handle_bad_vector_column(
|
|||||||
if dim is None:
|
if dim is None:
|
||||||
return data
|
return data
|
||||||
|
|
||||||
is_null = pc.is_null(vec_arr)
|
if not is_multivector:
|
||||||
# pc.list_value_length returns null for null list entries, so
|
is_null = pc.is_null(vec_arr)
|
||||||
# pc.not_equal(null, dim) also returns null. Use or_kleene so that
|
# pc.list_value_length returns null for null list entries, so
|
||||||
# True OR null = True (Kleene three-valued logic), ensuring null vectors
|
# pc.not_equal(null, dim) also returns null. Use or_kleene so that
|
||||||
# are counted as wrong-dim.
|
# True OR null = True (Kleene three-valued logic), ensuring null vectors
|
||||||
has_wrong_dim = pc.or_kleene(
|
# are counted as wrong-dim.
|
||||||
is_null,
|
has_wrong_dim = pc.or_kleene(
|
||||||
pc.not_equal(pc.list_value_length(vec_arr), dim),
|
is_null,
|
||||||
)
|
pc.not_equal(pc.list_value_length(vec_arr), dim),
|
||||||
|
)
|
||||||
|
|
||||||
has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
|
has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
|
||||||
|
|
||||||
@@ -4273,7 +4300,10 @@ def _handle_bad_vector_column(
|
|||||||
raise ValueError(
|
raise ValueError(
|
||||||
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
||||||
)
|
)
|
||||||
vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
|
if is_multivector:
|
||||||
|
vec_arr = _fill_bad_multivector_values(vec_arr, dim, fill_value)
|
||||||
|
else:
|
||||||
|
vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
||||||
|
|
||||||
@@ -4325,17 +4355,60 @@ def _fill_bad_vector_values(
|
|||||||
return filled.cast(arr.type)
|
return filled.cast(arr.type)
|
||||||
|
|
||||||
|
|
||||||
|
def _fill_bad_multivector_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()
|
||||||
|
|
||||||
|
filled_vectors = _fill_bad_vector_values(arr.values, dim, fill_value)
|
||||||
|
parent_nulls = pc.is_null(arr)
|
||||||
|
if pa.types.is_large_list(arr.type):
|
||||||
|
filled = pa.LargeListArray.from_arrays(
|
||||||
|
arr.offsets, filled_vectors, mask=parent_nulls
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
filled = pa.ListArray.from_arrays(
|
||||||
|
arr.offsets, filled_vectors, mask=parent_nulls
|
||||||
|
)
|
||||||
|
return filled.cast(arr.type)
|
||||||
|
|
||||||
|
|
||||||
|
def _multivector_has_wrong_dim(
|
||||||
|
arr: Union[pa.Array, pa.ChunkedArray], dim: int
|
||||||
|
) -> pa.BooleanArray:
|
||||||
|
if isinstance(arr, pa.ChunkedArray):
|
||||||
|
results = [_multivector_has_wrong_dim(chunk, dim) for chunk in arr.chunks]
|
||||||
|
return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_())
|
||||||
|
|
||||||
|
vectors = arr.flatten()
|
||||||
|
vector_is_wrong = pc.or_kleene(
|
||||||
|
pc.is_null(vectors),
|
||||||
|
pc.not_equal(pc.list_value_length(vectors), dim),
|
||||||
|
)
|
||||||
|
parent_indices = pc.list_parent_indices(arr)
|
||||||
|
wrong_parent_indices = pc.unique(pc.filter(parent_indices, vector_is_wrong))
|
||||||
|
indices = pa.array(range(len(arr)), type=pa.uint32())
|
||||||
|
return pc.or_(pc.is_null(arr), pc.is_in(indices, wrong_parent_indices))
|
||||||
|
|
||||||
|
|
||||||
def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
|
def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
|
||||||
if isinstance(arr, pa.ChunkedArray):
|
if isinstance(arr, pa.ChunkedArray):
|
||||||
values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
|
results = [has_nan_values(chunk) for chunk in arr.chunks]
|
||||||
else:
|
return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_())
|
||||||
values = arr.flatten()
|
|
||||||
if pa.types.is_float16(values.type):
|
values = arr.flatten()
|
||||||
|
if _is_list_like(values.type):
|
||||||
|
values_has_nan = has_nan_values(values)
|
||||||
|
elif pa.types.is_float16(values.type):
|
||||||
# is_nan isn't yet implemented for f16, so we cast to f32
|
# is_nan isn't yet implemented for f16, so we cast to f32
|
||||||
# https://github.com/apache/arrow/issues/45083
|
# https://github.com/apache/arrow/issues/45083
|
||||||
values_has_nan = pc.is_nan(values.cast(pa.float32()))
|
values_has_nan = pc.is_nan(values.cast(pa.float32()))
|
||||||
else:
|
elif pa.types.is_floating(values.type):
|
||||||
values_has_nan = pc.is_nan(values)
|
values_has_nan = pc.is_nan(values)
|
||||||
|
else:
|
||||||
|
return pa.array([False] * len(arr))
|
||||||
values_indices = pc.list_parent_indices(arr)
|
values_indices = pc.list_parent_indices(arr)
|
||||||
has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
|
has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
|
||||||
indices = pa.array(range(len(arr)), type=pa.uint32())
|
indices = pa.array(range(len(arr)), type=pa.uint32())
|
||||||
@@ -4350,6 +4423,16 @@ def _is_list_like(data_type: pa.DataType) -> bool:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_float_vector_type(data_type: pa.DataType) -> bool:
|
||||||
|
return _is_list_like(data_type) and pa.types.is_floating(data_type.value_type)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_multivector_type(data_type: pa.DataType) -> bool:
|
||||||
|
return (
|
||||||
|
pa.types.is_list(data_type) or pa.types.is_large_list(data_type)
|
||||||
|
) and _is_float_vector_type(data_type.value_type)
|
||||||
|
|
||||||
|
|
||||||
def _merge_metadata(*metadata_dicts: Optional[dict]) -> dict:
|
def _merge_metadata(*metadata_dicts: Optional[dict]) -> dict:
|
||||||
merged = {}
|
merged = {}
|
||||||
for metadata in metadata_dicts:
|
for metadata in metadata_dicts:
|
||||||
@@ -4441,6 +4524,16 @@ def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]:
|
|||||||
return pc.mode(lengths)[0].as_py()["mode"]
|
return pc.mode(lengths)[0].as_py()["mode"]
|
||||||
|
|
||||||
|
|
||||||
|
def _infer_vector_column_dim(
|
||||||
|
arr: Union[pa.Array, pa.ChunkedArray], is_multivector: bool
|
||||||
|
) -> Optional[int]:
|
||||||
|
if not is_multivector:
|
||||||
|
return _infer_vector_dim(arr)
|
||||||
|
if isinstance(arr, pa.ChunkedArray):
|
||||||
|
arr = arr.combine_chunks()
|
||||||
|
return _infer_vector_dim(arr.flatten())
|
||||||
|
|
||||||
|
|
||||||
def _validate_schema(schema: pa.Schema):
|
def _validate_schema(schema: pa.Schema):
|
||||||
"""
|
"""
|
||||||
Make sure the metadata is valid utf8
|
Make sure the metadata is valid utf8
|
||||||
|
|||||||
@@ -1710,6 +1710,19 @@ def test_create_with_nans(mem_db: DBConnection):
|
|||||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_with_nans_in_multivectors(mem_db: DBConnection):
|
||||||
|
multivector_type = pa.list_(pa.list_(pa.float32(), 128))
|
||||||
|
schema = pa.schema(
|
||||||
|
[pa.field("filename", pa.string()), pa.field("vector", multivector_type)]
|
||||||
|
)
|
||||||
|
vector = [0.1] * 128
|
||||||
|
vector[-1] = np.nan
|
||||||
|
data = [{"filename": "img1.jpg", "vector": [vector]}]
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Vector column 'vector' has NaNs"):
|
||||||
|
mem_db.create_table("nan_multivector", data=data, schema=schema)
|
||||||
|
|
||||||
|
|
||||||
def test_add_with_nans(mem_db: DBConnection):
|
def test_add_with_nans(mem_db: DBConnection):
|
||||||
schema = pa.schema(
|
schema = pa.schema(
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -400,6 +400,83 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
|
|||||||
assert output["vector"].combine_chunks() == expected
|
assert output["vector"].combine_chunks() == expected
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||||
|
def test_handle_bad_multivectors_nan(on_bad_vectors):
|
||||||
|
multivector_type = pa.list_(pa.list_(pa.float32(), 2))
|
||||||
|
vectors = pa.array(
|
||||||
|
[
|
||||||
|
[[1.0, float("nan")], [2.0, 3.0]],
|
||||||
|
[[4.0, 5.0]],
|
||||||
|
],
|
||||||
|
type=multivector_type,
|
||||||
|
)
|
||||||
|
data = pa.table({"vector": vectors})
|
||||||
|
|
||||||
|
if on_bad_vectors == "error":
|
||||||
|
with pytest.raises(ValueError, match="Vector column 'vector' has NaNs"):
|
||||||
|
_handle_bad_vectors(data.to_reader()).read_all()
|
||||||
|
return
|
||||||
|
|
||||||
|
output = _handle_bad_vectors(
|
||||||
|
data.to_reader(),
|
||||||
|
on_bad_vectors=on_bad_vectors,
|
||||||
|
fill_value=42.0,
|
||||||
|
).read_all()
|
||||||
|
|
||||||
|
if on_bad_vectors == "drop":
|
||||||
|
expected = pa.array([[[4.0, 5.0]]], type=multivector_type)
|
||||||
|
elif on_bad_vectors == "fill":
|
||||||
|
expected = pa.array(
|
||||||
|
[[[1.0, 42.0], [2.0, 3.0]], [[4.0, 5.0]]],
|
||||||
|
type=multivector_type,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
expected = pa.array([None, [[4.0, 5.0]]], type=multivector_type)
|
||||||
|
|
||||||
|
assert output["vector"].combine_chunks() == expected
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||||
|
def test_handle_bad_variable_multivectors(on_bad_vectors):
|
||||||
|
target_type = pa.list_(pa.list_(pa.float32(), 2))
|
||||||
|
vectors = pa.array(
|
||||||
|
[
|
||||||
|
[[1.0, float("nan")], [2.0, 3.0]],
|
||||||
|
[[4.0]],
|
||||||
|
[[5.0, 6.0]],
|
||||||
|
]
|
||||||
|
)
|
||||||
|
data = pa.table({"vector": vectors})
|
||||||
|
|
||||||
|
if on_bad_vectors == "error":
|
||||||
|
with pytest.raises(ValueError, match="variable length vectors"):
|
||||||
|
_handle_bad_vectors(
|
||||||
|
data.to_reader(),
|
||||||
|
target_schema=pa.schema({"vector": target_type}),
|
||||||
|
).read_all()
|
||||||
|
return
|
||||||
|
|
||||||
|
output = _handle_bad_vectors(
|
||||||
|
data.to_reader(),
|
||||||
|
on_bad_vectors=on_bad_vectors,
|
||||||
|
fill_value=42.0,
|
||||||
|
target_schema=pa.schema({"vector": target_type}),
|
||||||
|
).read_all()
|
||||||
|
|
||||||
|
if on_bad_vectors == "drop":
|
||||||
|
expected = [[[5.0, 6.0]]]
|
||||||
|
elif on_bad_vectors == "fill":
|
||||||
|
expected = [
|
||||||
|
[[1.0, 42.0], [2.0, 3.0]],
|
||||||
|
[[4.0, 42.0]],
|
||||||
|
[[5.0, 6.0]],
|
||||||
|
]
|
||||||
|
else:
|
||||||
|
expected = [None, None, [[5.0, 6.0]]]
|
||||||
|
|
||||||
|
assert output["vector"].combine_chunks().to_pylist() == expected
|
||||||
|
|
||||||
|
|
||||||
def test_handle_bad_vectors_noop():
|
def test_handle_bad_vectors_noop():
|
||||||
# ChunkedArray should be preserved as-is
|
# ChunkedArray should be preserved as-is
|
||||||
vector = pa.chunked_array(
|
vector = pa.chunked_array(
|
||||||
|
|||||||
@@ -258,6 +258,7 @@ mod tests {
|
|||||||
FixedSizeListArray, Float32Array, Int32Array, LargeStringArray, ListArray, RecordBatch,
|
FixedSizeListArray, Float32Array, Int32Array, LargeStringArray, ListArray, RecordBatch,
|
||||||
RecordBatchIterator, record_batch,
|
RecordBatchIterator, record_batch,
|
||||||
};
|
};
|
||||||
|
use arrow_buffer::OffsetBuffer;
|
||||||
use arrow_schema::{ArrowError, DataType, Field, Schema};
|
use arrow_schema::{ArrowError, DataType, Field, Schema};
|
||||||
use futures::TryStreamExt;
|
use futures::TryStreamExt;
|
||||||
use lance::dataset::{WriteMode, WriteParams};
|
use lance::dataset::{WriteMode, WriteParams};
|
||||||
@@ -885,6 +886,63 @@ mod tests {
|
|||||||
assert_eq!(row_count, 1);
|
assert_eq!(row_count, 1);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_add_rejects_nan_multivectors() {
|
||||||
|
let vector_type =
|
||||||
|
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
|
||||||
|
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||||
|
"embedding",
|
||||||
|
DataType::List(Arc::new(Field::new("item", vector_type.clone(), true))),
|
||||||
|
false,
|
||||||
|
)]));
|
||||||
|
|
||||||
|
let db = connect("memory://").execute().await.unwrap();
|
||||||
|
let table = db
|
||||||
|
.create_empty_table("nan_multivector_test", schema.clone())
|
||||||
|
.execute()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let vectors = FixedSizeListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||||
|
4,
|
||||||
|
Arc::new(Float32Array::from(vec![
|
||||||
|
0.1,
|
||||||
|
0.2,
|
||||||
|
0.3,
|
||||||
|
0.4,
|
||||||
|
0.5,
|
||||||
|
f32::NAN,
|
||||||
|
0.7,
|
||||||
|
0.8,
|
||||||
|
])),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let multivectors = ListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", vector_type, true)),
|
||||||
|
OffsetBuffer::from_lengths([2]),
|
||||||
|
Arc::new(vectors),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let batch = RecordBatch::try_new(schema, vec![Arc::new(multivectors)]).unwrap();
|
||||||
|
|
||||||
|
let err = table.add(batch.clone()).execute().await.unwrap_err();
|
||||||
|
assert!(
|
||||||
|
err.to_string().contains("NaN"),
|
||||||
|
"Expected error mentioning NaN values, but got: {err:?}"
|
||||||
|
);
|
||||||
|
|
||||||
|
table
|
||||||
|
.add(batch)
|
||||||
|
.on_nan_vectors(NaNVectorBehavior::Keep)
|
||||||
|
.execute()
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(table.count_rows(None).await.unwrap(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_add_subschema() {
|
async fn test_add_subschema() {
|
||||||
let data = record_batch!(("id", Int64, [4, 5]), ("text", Utf8, ["foo", "bar"])).unwrap();
|
let data = record_batch!(("id", Int64, [4, 5]), ("text", Utf8, ["foo", "bar"])).unwrap();
|
||||||
|
|||||||
@@ -5,7 +5,7 @@
|
|||||||
|
|
||||||
use std::sync::{Arc, LazyLock};
|
use std::sync::{Arc, LazyLock};
|
||||||
|
|
||||||
use arrow_array::{Array, FixedSizeListArray};
|
use arrow_array::{Array, FixedSizeListArray, ListArray};
|
||||||
use arrow_schema::{DataType, Field, FieldRef};
|
use arrow_schema::{DataType, Field, FieldRef};
|
||||||
use datafusion_common::config::ConfigOptions;
|
use datafusion_common::config::ConfigOptions;
|
||||||
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
|
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
|
||||||
@@ -19,16 +19,20 @@ use crate::{Error, Result};
|
|||||||
static REJECT_NAN_UDF: LazyLock<Arc<datafusion_expr::ScalarUDF>> =
|
static REJECT_NAN_UDF: LazyLock<Arc<datafusion_expr::ScalarUDF>> =
|
||||||
LazyLock::new(|| Arc::new(datafusion_expr::ScalarUDF::from(RejectNanUdf::new())));
|
LazyLock::new(|| Arc::new(datafusion_expr::ScalarUDF::from(RejectNanUdf::new())));
|
||||||
|
|
||||||
/// Returns true if the field is a vector column: FixedSizeList<Float16/32/64>.
|
/// Returns true if the field is a vector or multivector column.
|
||||||
fn is_vector_field(field: &Field) -> bool {
|
fn is_vector_field(field: &Field) -> bool {
|
||||||
if let DataType::FixedSizeList(child, _) = field.data_type() {
|
fn is_vector_data_type(data_type: &DataType) -> bool {
|
||||||
matches!(
|
match data_type {
|
||||||
child.data_type(),
|
DataType::FixedSizeList(child, _) => matches!(
|
||||||
DataType::Float16 | DataType::Float32 | DataType::Float64
|
child.data_type(),
|
||||||
)
|
DataType::Float16 | DataType::Float32 | DataType::Float64
|
||||||
} else {
|
),
|
||||||
false
|
DataType::List(child) => is_vector_data_type(child.data_type()),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
is_vector_data_type(field.data_type())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Wraps the input plan with a projection that checks vector columns for NaN values.
|
/// Wraps the input plan with a projection that checks vector columns for NaN values.
|
||||||
@@ -69,8 +73,8 @@ pub fn reject_nan_vectors(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Execu
|
|||||||
Ok(Arc::new(projection))
|
Ok(Arc::new(projection))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A scalar UDF that passes through FixedSizeList arrays unchanged, but errors
|
/// A scalar UDF that passes through vector arrays unchanged, but errors if any
|
||||||
/// if any float values in the list are NaN.
|
/// float values in the vector or multivector are NaN.
|
||||||
#[derive(Debug, Hash, PartialEq, Eq)]
|
#[derive(Debug, Hash, PartialEq, Eq)]
|
||||||
struct RejectNanUdf {
|
struct RejectNanUdf {
|
||||||
signature: Signature,
|
signature: Signature,
|
||||||
@@ -113,44 +117,54 @@ impl ScalarUDFImpl for RejectNanUdf {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
|
fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
|
||||||
let fsl = array
|
fn contains_nan(array: &dyn Array) -> datafusion_common::Result<bool> {
|
||||||
.as_any()
|
match array.data_type() {
|
||||||
.downcast_ref::<FixedSizeListArray>()
|
DataType::Float16 => Ok(array
|
||||||
.ok_or_else(|| {
|
|
||||||
datafusion_common::DataFusionError::Internal(
|
|
||||||
"reject_nan expected FixedSizeList".to_string(),
|
|
||||||
)
|
|
||||||
})?;
|
|
||||||
|
|
||||||
// Only inspect elements that are both in a valid parent row and non-null
|
|
||||||
// themselves. Values backing null parent rows or null child elements may
|
|
||||||
// contain garbage (including NaN) per the Arrow spec.
|
|
||||||
let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| {
|
|
||||||
let row = fsl.value(i);
|
|
||||||
match row.data_type() {
|
|
||||||
DataType::Float16 => row
|
|
||||||
.as_any()
|
.as_any()
|
||||||
.downcast_ref::<arrow_array::Float16Array>()
|
.downcast_ref::<arrow_array::Float16Array>()
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|v| v.is_some_and(|v| v.is_nan())),
|
.any(|v| v.is_some_and(|v| v.is_nan()))),
|
||||||
DataType::Float32 => row
|
DataType::Float32 => Ok(array
|
||||||
.as_any()
|
.as_any()
|
||||||
.downcast_ref::<arrow_array::Float32Array>()
|
.downcast_ref::<arrow_array::Float32Array>()
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|v| v.is_some_and(|v| v.is_nan())),
|
.any(|v| v.is_some_and(|v| v.is_nan()))),
|
||||||
DataType::Float64 => row
|
DataType::Float64 => Ok(array
|
||||||
.as_any()
|
.as_any()
|
||||||
.downcast_ref::<arrow_array::Float64Array>()
|
.downcast_ref::<arrow_array::Float64Array>()
|
||||||
.unwrap()
|
.unwrap()
|
||||||
.iter()
|
.iter()
|
||||||
.any(|v| v.is_some_and(|v| v.is_nan())),
|
.any(|v| v.is_some_and(|v| v.is_nan()))),
|
||||||
_ => false,
|
DataType::FixedSizeList(_, _) => {
|
||||||
|
let lists = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
|
||||||
|
for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) {
|
||||||
|
if contains_nan(lists.value(i).as_ref())? {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
DataType::List(_) => {
|
||||||
|
let lists = array.as_any().downcast_ref::<ListArray>().unwrap();
|
||||||
|
for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) {
|
||||||
|
if contains_nan(lists.value(i).as_ref())? {
|
||||||
|
return Ok(true);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(false)
|
||||||
|
}
|
||||||
|
data_type => Err(datafusion_common::DataFusionError::Internal(format!(
|
||||||
|
"reject_nan expected a vector or multivector, got {data_type}"
|
||||||
|
))),
|
||||||
}
|
}
|
||||||
});
|
}
|
||||||
|
|
||||||
if has_nan {
|
// Only inspect elements that are both in a valid parent row and non-null
|
||||||
|
// themselves. Values backing null parent rows or null child elements may
|
||||||
|
// contain garbage (including NaN) per the Arrow spec.
|
||||||
|
if contains_nan(array)? {
|
||||||
return Err(datafusion_common::DataFusionError::ArrowError(
|
return Err(datafusion_common::DataFusionError::ArrowError(
|
||||||
Box::new(arrow_schema::ArrowError::ComputeError(
|
Box::new(arrow_schema::ArrowError::ComputeError(
|
||||||
"Vector column contains NaN values".to_string(),
|
"Vector column contains NaN values".to_string(),
|
||||||
@@ -166,6 +180,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use arrow_array::Float32Array;
|
use arrow_array::Float32Array;
|
||||||
|
use arrow_buffer::OffsetBuffer;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_passes_clean_vectors() {
|
fn test_passes_clean_vectors() {
|
||||||
@@ -191,6 +206,46 @@ mod tests {
|
|||||||
assert!(check_no_nans(&fsl).is_err());
|
assert!(check_no_nans(&fsl).is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_rejects_nan_multivectors() {
|
||||||
|
let vectors = FixedSizeListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||||
|
2,
|
||||||
|
Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, f32::NAN])),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let multivectors = ListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", vectors.data_type().clone(), true)),
|
||||||
|
OffsetBuffer::from_lengths([2]),
|
||||||
|
Arc::new(vectors),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(check_no_nans(&multivectors).is_err());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_skips_null_multivector_rows() {
|
||||||
|
let vectors = FixedSizeListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||||
|
2,
|
||||||
|
Arc::new(Float32Array::from(vec![f32::NAN, f32::NAN, 1.0, 2.0])),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let multivectors = ListArray::try_new(
|
||||||
|
Arc::new(Field::new("item", vectors.data_type().clone(), true)),
|
||||||
|
OffsetBuffer::from_lengths([1, 1]),
|
||||||
|
Arc::new(vectors),
|
||||||
|
Some(vec![false, true].into()),
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(check_no_nans(&multivectors).is_ok());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn test_skips_null_rows() {
|
fn test_skips_null_rows() {
|
||||||
// Values backing null rows may contain NaN per the Arrow spec.
|
// Values backing null rows may contain NaN per the Arrow spec.
|
||||||
@@ -254,6 +309,15 @@ mod tests {
|
|||||||
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), 4),
|
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), 4),
|
||||||
false,
|
false,
|
||||||
)));
|
)));
|
||||||
|
assert!(is_vector_field(&Field::new(
|
||||||
|
"v",
|
||||||
|
DataType::List(Arc::new(Field::new(
|
||||||
|
"item",
|
||||||
|
DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4,),
|
||||||
|
true,
|
||||||
|
))),
|
||||||
|
false,
|
||||||
|
)));
|
||||||
assert!(!is_vector_field(&Field::new("id", DataType::Int32, false)));
|
assert!(!is_vector_field(&Field::new("id", DataType::Int32, false)));
|
||||||
assert!(!is_vector_field(&Field::new(
|
assert!(!is_vector_field(&Field::new(
|
||||||
"v",
|
"v",
|
||||||
|
|||||||
Reference in New Issue
Block a user