mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 05:26:14 +00:00
fix(python): preserve Arrow map scalars in updates (#4448)
Arrow map scalars returned by a table query cannot currently be passed back to `update`: the generic Arrow scalar conversion turns them into a list of tuples, which the SQL literal encoder cannot handle. Empty map scalars become an untyped SQL list and also fail to update a map column. Add a dedicated `MapScalar` conversion that renders its keys and values as a SQL map while preserving the existing null conversion. The regression tests exercise real synchronous and asynchronous table updates, quoted keys, nullable values, empty typed maps, and nested maps. Validation: the complete documented unit-test scope passed locally (1,515 passed, 37 skipped; unchanged baseline: 1,507 passed, 37 skipped). Full doctests passed (71 passed, 6 skipped), and repository-wide Python formatting/lint, Rust formatting/clippy, and Python 3.13 type checking passed. The locally built native wheel also passed all eight regression cases in a separate Python 3.13 environment. Full Linux/S3 integration and all platform matrices remain for CI. The SQL map constructor's existing backend constraints still apply: duplicate keys are rejected, and a nested map containing both nonempty and empty inner maps can fail SQL type unification. This change does not claim arbitrary Arrow map roundtrip support. AI assistance: implementation and local verification used Codex with gpt-6.1-sol agents; the diff and evidence were reviewed before submission. --------- Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
@@ -452,6 +452,15 @@ def _(value: pa.Scalar):
|
||||
return value_to_sql(value.as_py())
|
||||
|
||||
|
||||
@value_to_sql.register(pa.MapScalar)
|
||||
def _(value: pa.MapScalar):
|
||||
if not value.is_valid:
|
||||
return value_to_sql(None)
|
||||
keys = value_to_sql(list(value.values.field(0)))
|
||||
values = value_to_sql(list(value.values.field(1)))
|
||||
return f"map({keys}, {values})"
|
||||
|
||||
|
||||
@value_to_sql.register(np.ndarray)
|
||||
def _(value: np.ndarray):
|
||||
return value_to_sql(value.tolist())
|
||||
|
||||
@@ -2724,6 +2724,57 @@ def test_update_with_arrow_scalar(mem_db: DBConnection):
|
||||
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("async_api", [False, True])
|
||||
@pytest.mark.parametrize("entries", [[("o'brien", 4), ("nil", None)], [], None])
|
||||
async def test_update_map_arrow_scalar(tmp_path, async_api, entries):
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
data = pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
"attrs": pa.array([[("original", 1)], entries], type=map_type),
|
||||
}
|
||||
)
|
||||
if async_api:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
table = await db.create_table("maps", data)
|
||||
source = await table.query().where("id = 2").to_arrow()
|
||||
result = await table.update({"attrs": source["attrs"][0]}, where="id = 1")
|
||||
actual = await table.query().to_arrow()
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("maps", data)
|
||||
source = table.search().where("id = 2").to_arrow()
|
||||
result = table.update(values={"attrs": source["attrs"][0]}, where="id = 1")
|
||||
actual = table.to_arrow()
|
||||
assert result.rows_updated == 1
|
||||
assert actual.schema.field("attrs").type == map_type
|
||||
assert actual.sort_by("id")["attrs"].to_pylist() == [entries, entries]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"map_type, entries",
|
||||
[
|
||||
(pa.map_(pa.int64(), pa.string()), []),
|
||||
(
|
||||
pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())),
|
||||
[("outer", [("inner", 3)]), ("null", None)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_update_typed_map_arrow_scalar(mem_db: DBConnection, map_type, entries):
|
||||
table = mem_db.create_table(
|
||||
"maps",
|
||||
pa.table({"id": [1, 2], "attrs": pa.array([None, entries], type=map_type)}),
|
||||
)
|
||||
scalar = table.search().where("id = 2").to_arrow()["attrs"][0]
|
||||
result = table.update(values={"attrs": scalar}, where="id = 1")
|
||||
actual = table.to_arrow().sort_by("id")["attrs"]
|
||||
assert result.rows_updated == 1
|
||||
assert actual.type == map_type
|
||||
assert actual.to_pylist() == [entries, entries]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_arrow_scalar", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"decimal_type, value",
|
||||
|
||||
Reference in New Issue
Block a user