mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 16:38:31 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 43465f06d5 | |||
| bc762926ab | |||
| 2426f275cd |
@@ -452,6 +452,210 @@ def _field_extension_name(field: pa.Field) -> Optional[str]:
|
||||
return extension_name
|
||||
|
||||
|
||||
_JSON_EXTENSION_NAMES = {"arrow.json", "lance.json"}
|
||||
_BLOB_EXTENSION_NAME = "lance.blob.v2"
|
||||
|
||||
|
||||
def _field_contains_write_extension(field: pa.Field) -> bool:
|
||||
extension_name = _field_extension_name(field)
|
||||
if (
|
||||
extension_name in _JSON_EXTENSION_NAMES
|
||||
or extension_name == _BLOB_EXTENSION_NAME
|
||||
):
|
||||
return True
|
||||
if pa.types.is_struct(field.type):
|
||||
return any(_field_contains_write_extension(child) for child in field.type)
|
||||
if (
|
||||
pa.types.is_list(field.type)
|
||||
or pa.types.is_large_list(field.type)
|
||||
or pa.types.is_fixed_size_list(field.type)
|
||||
):
|
||||
return _field_contains_write_extension(field.type.value_field)
|
||||
return False
|
||||
|
||||
|
||||
def _with_field_type(
|
||||
field: pa.Field,
|
||||
data_type: pa.DataType,
|
||||
*,
|
||||
name: Optional[str] = None,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> pa.Field:
|
||||
return pa.field(
|
||||
name or field.name,
|
||||
data_type,
|
||||
nullable=field.nullable,
|
||||
metadata=field.metadata if metadata is None else metadata,
|
||||
)
|
||||
|
||||
|
||||
def _with_list_value_field(
|
||||
data_type: pa.DataType, value_field: pa.Field
|
||||
) -> pa.DataType:
|
||||
if pa.types.is_list(data_type):
|
||||
return pa.list_(value_field)
|
||||
if pa.types.is_large_list(data_type):
|
||||
return pa.large_list(value_field)
|
||||
return pa.list_(value_field, data_type.list_size)
|
||||
|
||||
|
||||
def _extension_storage_field(field: pa.Field) -> pa.Field:
|
||||
"""Return a from-pylist-compatible field for nested write extensions."""
|
||||
extension_name = _field_extension_name(field)
|
||||
if extension_name in _JSON_EXTENSION_NAMES:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = b"arrow.json"
|
||||
return _with_field_type(field, pa.string(), metadata=metadata)
|
||||
if extension_name == _BLOB_EXTENSION_NAME:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = _BLOB_EXTENSION_NAME.encode()
|
||||
metadata[b"ARROW:extension:metadata"] = b""
|
||||
storage_type = getattr(field.type, "storage_type", field.type)
|
||||
return _with_field_type(field, storage_type, metadata=metadata)
|
||||
if pa.types.is_struct(field.type):
|
||||
children = [_extension_storage_field(child) for child in field.type]
|
||||
return _with_field_type(field, pa.struct(children))
|
||||
if _is_list_like(field.type):
|
||||
value_field = _extension_storage_field(field.type.value_field)
|
||||
return _with_field_type(field, _with_list_value_field(field.type, value_field))
|
||||
return field
|
||||
|
||||
|
||||
def _prepare_extension_field(
|
||||
field: pa.Field, target_field: pa.Field
|
||||
) -> Tuple[pa.Field, bool]:
|
||||
extension_name = _field_extension_name(target_field)
|
||||
if extension_name in _JSON_EXTENSION_NAMES:
|
||||
metadata = dict(field.metadata or {})
|
||||
metadata[b"ARROW:extension:name"] = b"arrow.json"
|
||||
return _with_field_type(field, pa.string(), metadata=metadata), True
|
||||
if extension_name == _BLOB_EXTENSION_NAME and pa.types.is_null(field.type):
|
||||
return _with_field_type(field, pa.large_binary()), True
|
||||
|
||||
if pa.types.is_struct(field.type) and pa.types.is_struct(target_field.type):
|
||||
target_children = {child.name: child for child in target_field.type}
|
||||
children = []
|
||||
changed = False
|
||||
for child in field.type:
|
||||
target_child = target_children.get(child.name)
|
||||
if target_child is None:
|
||||
children.append(child)
|
||||
continue
|
||||
prepared, child_changed = _prepare_extension_field(child, target_child)
|
||||
children.append(prepared)
|
||||
changed = changed or child_changed
|
||||
if changed:
|
||||
return _with_field_type(field, pa.struct(children)), True
|
||||
|
||||
if _is_list_like(field.type) and _is_list_like(target_field.type):
|
||||
target_value_field = target_field.type.value_field
|
||||
if _field_contains_write_extension(target_value_field):
|
||||
prepared = _extension_storage_field(target_value_field)
|
||||
data_type = _with_list_value_field(target_field.type, prepared)
|
||||
return _with_field_type(field, data_type), True
|
||||
|
||||
return field, False
|
||||
|
||||
|
||||
def _prepare_extension_value(
|
||||
value: Any, target_field: pa.Field, *, within_list: bool = False
|
||||
) -> Any:
|
||||
"""Shape raw nested blob values for PyArrow's struct construction."""
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
extension_name = _field_extension_name(target_field)
|
||||
if extension_name == _BLOB_EXTENSION_NAME and within_list:
|
||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||
return {"data": value}
|
||||
if isinstance(value, str):
|
||||
return {"uri": value}
|
||||
return value
|
||||
|
||||
if pa.types.is_struct(target_field.type) and isinstance(value, dict):
|
||||
target_children = {child.name: child for child in target_field.type}
|
||||
return {
|
||||
name: _prepare_extension_value(
|
||||
child_value, target_children[name], within_list=within_list
|
||||
)
|
||||
if name in target_children
|
||||
else child_value
|
||||
for name, child_value in value.items()
|
||||
}
|
||||
|
||||
if _is_list_like(target_field.type) and isinstance(value, (list, tuple)):
|
||||
return [
|
||||
_prepare_extension_value(
|
||||
item, target_field.type.value_field, within_list=True
|
||||
)
|
||||
for item in value
|
||||
]
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _prepare_extension_list(data: DATA, target_schema: pa.Schema) -> DATA:
|
||||
"""Give inferred list columns the logical type required by extensions."""
|
||||
if not isinstance(data, list) or not data or not isinstance(data[0], dict):
|
||||
return data
|
||||
|
||||
target_fields = {field.name: field for field in target_schema}
|
||||
if not any(
|
||||
_field_contains_write_extension(field) for field in target_fields.values()
|
||||
):
|
||||
return data
|
||||
|
||||
inferred = pa.Table.from_pylist(data)
|
||||
fields = []
|
||||
changed = False
|
||||
for field in inferred.schema:
|
||||
target_field = target_fields.get(field.name)
|
||||
if target_field is None:
|
||||
fields.append(field)
|
||||
continue
|
||||
prepared, field_changed = _prepare_extension_field(field, target_field)
|
||||
fields.append(prepared)
|
||||
changed = changed or field_changed
|
||||
|
||||
if not changed:
|
||||
return inferred
|
||||
|
||||
insert_schema = pa.schema(fields, metadata=inferred.schema.metadata)
|
||||
prepared_data = [
|
||||
{
|
||||
name: _prepare_extension_value(value, target_fields[name])
|
||||
if name in target_fields
|
||||
else value
|
||||
for name, value in row.items()
|
||||
}
|
||||
for row in data
|
||||
]
|
||||
return pa.Table.from_pylist(prepared_data, schema=insert_schema)
|
||||
|
||||
|
||||
def _is_blob_source_field(field: pa.Field) -> bool:
|
||||
if _field_extension_name(field) == _BLOB_EXTENSION_NAME:
|
||||
return True
|
||||
|
||||
predicates = (
|
||||
"is_binary",
|
||||
"is_large_binary",
|
||||
"is_binary_view",
|
||||
"is_string",
|
||||
"is_large_string",
|
||||
"is_string_view",
|
||||
)
|
||||
if any(
|
||||
predicate(field.type)
|
||||
for name in predicates
|
||||
if (predicate := getattr(pa.types, name, None)) is not None
|
||||
):
|
||||
return True
|
||||
return pa.types.is_struct(field.type) and any(
|
||||
child.name in {"data", "uri"} for child in field.type
|
||||
)
|
||||
|
||||
|
||||
def _align_field_types(
|
||||
fields: List[pa.Field],
|
||||
target_fields: List[pa.Field],
|
||||
@@ -464,13 +668,21 @@ def _align_field_types(
|
||||
target_field = next((f for f in target_fields if f.name == field.name), None)
|
||||
if target_field is None:
|
||||
raise ValueError(f"Field '{field.name}' not found in target schema")
|
||||
target_extension_name = _field_extension_name(target_field)
|
||||
# Preserve accepted blob carriers so Lance can construct the declared
|
||||
# blob struct after optional Python preprocessing.
|
||||
if target_extension_name == _BLOB_EXTENSION_NAME and _is_blob_source_field(
|
||||
field
|
||||
):
|
||||
new_fields.append(field)
|
||||
continue
|
||||
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
|
||||
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
|
||||
# input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
if (
|
||||
_field_extension_name(field) == "arrow.json"
|
||||
and _field_extension_name(target_field) == "lance.json"
|
||||
and target_extension_name in _JSON_EXTENSION_NAMES
|
||||
):
|
||||
new_fields.append(field)
|
||||
continue
|
||||
@@ -5442,6 +5654,9 @@ class AsyncTable:
|
||||
if fill_value is None:
|
||||
fill_value = 0.0
|
||||
|
||||
if mode != "overwrite":
|
||||
data = _prepare_extension_list(data, schema)
|
||||
|
||||
# _santitize_data is an old code path, but we will use it until the
|
||||
# new code path is ready.
|
||||
if mode == "overwrite":
|
||||
|
||||
@@ -203,6 +203,80 @@ def test_fetch_blobs_preserves_null_and_empty_values():
|
||||
assert blobs[3].as_py() == b"present"
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column():
|
||||
table = _blob_table("all_null_add", [{"id": 1, "image": None}])
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_blob_column_with_sanitizer():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("all_null_sanitized_add", schema=schema)
|
||||
|
||||
table.add([{"id": 1, "image": None}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert len(blobs) == 1
|
||||
assert blobs[0].as_py() is None
|
||||
|
||||
|
||||
def test_add_all_null_list_to_nested_blob_column():
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
info_field = pa.field("info", pa.struct([blob_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[_blob_array("image", [b"seed"])], fields=[blob_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = db.create_table("nested_null_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "info": {"image": None}}])
|
||||
table.add([{"id": 2, "info": {"image": None}}], on_bad_vectors="fill")
|
||||
|
||||
hits = table.search().where("id > 0").to_arrow()
|
||||
blobs = table.fetch_blobs("info.image", hits)
|
||||
assert len(blobs) == 2
|
||||
assert all(blob.as_py() is None for blob in blobs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("large_list", [False, True], ids=["list", "large_list"])
|
||||
def test_add_list_of_dicts_to_blob_list_column(large_list):
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
blob_values = _blob_array("image", [b"seed"])
|
||||
if large_list:
|
||||
items_field = pa.field("items", pa.large_list(blob_field))
|
||||
items = pa.LargeListArray.from_arrays(
|
||||
pa.array([0, 1], type=pa.int64()), blob_values
|
||||
)
|
||||
else:
|
||||
items_field = pa.field("items", pa.list_(blob_field))
|
||||
items = pa.ListArray.from_arrays(pa.array([0, 1], type=pa.int32()), blob_values)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), items],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), items_field]),
|
||||
)
|
||||
table = db.create_table(f"blob_{large_list}_list_add", data=seed)
|
||||
|
||||
table.add([{"id": 1, "items": [None]}])
|
||||
table.add(
|
||||
[{"id": 2, "items": [b"a", None]}],
|
||||
on_bad_vectors="fill",
|
||||
)
|
||||
|
||||
ids = table.search().select(["id"]).to_arrow()["id"].to_pylist()
|
||||
assert sorted(ids) == [0, 1, 2]
|
||||
assert pa.types.is_large_list(table.schema.field("items").type) is large_list
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls():
|
||||
table = _blob_table(
|
||||
"range_alignment",
|
||||
@@ -679,6 +753,7 @@ def test_add_external_uri_string_round_trips_with_flag(tmp_path):
|
||||
table = db.create_table("external_string", schema=schema)
|
||||
table.add(
|
||||
[{"id": 1, "image": blob_path.as_uri()}],
|
||||
on_bad_vectors="fill",
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -770,6 +770,55 @@ async def test_add_async(mem_db_async: AsyncConnection):
|
||||
assert await table.count_rows() == 3
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("values", "expected"),
|
||||
[
|
||||
([None], [None]),
|
||||
([None, '{"k": 1}'], [None, '{"k":1}']),
|
||||
(['{"k": 2}'], ['{"k":2}']),
|
||||
],
|
||||
)
|
||||
async def test_add_list_of_dicts_to_json_column(
|
||||
mem_db_async: AsyncConnection, values, expected
|
||||
):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("value", pa.json_())])
|
||||
table = await mem_db_async.create_table("json_list_add", schema=schema)
|
||||
|
||||
await table.add([{"id": idx, "value": value} for idx, value in enumerate(values)])
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert [row["value"] for row in rows] == expected
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_list_of_dicts_to_nested_json_column(
|
||||
mem_db_async: AsyncConnection,
|
||||
):
|
||||
json_field = pa.field("value", pa.json_())
|
||||
info_field = pa.field("info", pa.struct([json_field]))
|
||||
info = pa.StructArray.from_arrays(
|
||||
[pa.array(['{"seed": 0}'], type=pa.json_())], fields=[json_field]
|
||||
)
|
||||
seed = pa.Table.from_arrays(
|
||||
[pa.array([0], type=pa.int64()), info],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), info_field]),
|
||||
)
|
||||
table = await mem_db_async.create_table("nested_json_list_add", data=seed)
|
||||
|
||||
await table.add([{"id": 1, "info": {"value": '{"k": 1}'}}])
|
||||
await table.add([{"id": 2, "info": {"value": '{"k": 2}'}}], on_bad_vectors="fill")
|
||||
|
||||
rows = (await table.to_arrow()).sort_by("id").to_pylist()
|
||||
assert rows == [
|
||||
{"id": 0, "info": {"value": '{"seed":0}'}},
|
||||
{"id": 1, "info": {"value": '{"k":1}'}},
|
||||
{"id": 2, "info": {"value": '{"k":2}'}},
|
||||
]
|
||||
|
||||
|
||||
def test_add_overwrite_infers_vector_schema(mem_db: DBConnection):
|
||||
"""Overwrite should infer vector columns the same way create_table does.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user