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:
LeonSGP
2026-10-24 04:22:30 +08:00
committed by GitHub
co-authored by Xuanwo
parent e74a7c15b5
commit 3321860e9f
2 changed files with 60 additions and 0 deletions
+9
View File
@@ -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())
+51
View File
@@ -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",