mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-16 00:32:26 +00:00
fix(python): convert objects for JSON fields (#4211)
## Summary Closes: #4060 - Convert Python dict/list objects to JSON strings during ingestion. - Support JSON fields nested inside structs and lists. - Cover `add` and `merge_insert` JSON ingestion paths. - Add regression coverage for nested struct and list JSON fields. ## Testing - `git diff --check` - Python compilation passed. - Focused pytest was attempted but could not complete because the native extension build stalled during `uv` bootstrap.
This commit is contained in:
Rudra Prasad Bhuyan
@@ -4,6 +4,7 @@
|
||||
|
||||
import ctypes
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
@@ -17,6 +18,7 @@ from typing import List
|
||||
from unittest.mock import patch
|
||||
|
||||
import lancedb
|
||||
from lancedb import table as table_module
|
||||
from lancedb.dependencies import _PANDAS_AVAILABLE
|
||||
from lancedb.index import BTree, FTS, HnswFlat, HnswPq, HnswSq, IvfPq
|
||||
import numpy as np
|
||||
@@ -3027,6 +3029,236 @@ async def test_add_sanitization_encodes_json(mem_db_async: AsyncConnection):
|
||||
assert rows == [{"id": "c", "j": '{"k":3}'}]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
def test_add_python_objects_to_json_column(mem_db: DBConnection):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("payload", pa.json_()),
|
||||
pa.field("metadata", pa.struct([("label", pa.string())])),
|
||||
pa.field("tags", pa.list_(pa.string())),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table("json_python_objects_add", schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"payload": {"foo": "bar", "count": 2},
|
||||
"metadata": {"label": "dict"},
|
||||
"tags": ["a", "b"],
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"payload": {
|
||||
"name": "alice",
|
||||
"tags": ["x", "y"],
|
||||
"nested": {"enabled": True},
|
||||
},
|
||||
"metadata": {"label": "nested"},
|
||||
"tags": ["c"],
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"payload": ["x", "y", "z"],
|
||||
"metadata": {"label": "list"},
|
||||
"tags": ["d", "e"],
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"payload": '{"foo": "bar"}',
|
||||
"metadata": {"label": "string"},
|
||||
"tags": ["f"],
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"payload": None,
|
||||
"metadata": {"label": "null"},
|
||||
"tags": ["g"],
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
|
||||
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
|
||||
assert json.loads(rows[2]["payload"]) == {
|
||||
"name": "alice",
|
||||
"tags": ["x", "y"],
|
||||
"nested": {"enabled": True},
|
||||
}
|
||||
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
|
||||
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
|
||||
assert rows[5]["payload"] is None
|
||||
assert rows[1]["metadata"] == {"label": "dict"}
|
||||
assert rows[1]["tags"] == ["a", "b"]
|
||||
|
||||
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
|
||||
assert [row["id"] for row in matched] == [1]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
def test_merge_insert_python_objects_to_json_column(mem_db: DBConnection):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("payload", pa.json_()),
|
||||
pa.field("metadata", pa.struct([("label", pa.string())])),
|
||||
pa.field("tags", pa.list_(pa.string())),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table("json_python_objects_merge", schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"payload": {"old": True},
|
||||
"metadata": {"label": "old"},
|
||||
"tags": ["old"],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
table.merge_insert(
|
||||
"id"
|
||||
).when_matched_update_all().when_not_matched_insert_all().execute(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"payload": {"foo": "bar", "count": 2},
|
||||
"metadata": {"label": "dict"},
|
||||
"tags": ["a", "b"],
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"payload": {
|
||||
"name": "alice",
|
||||
"tags": ["x", "y"],
|
||||
"nested": {"enabled": True},
|
||||
},
|
||||
"metadata": {"label": "nested"},
|
||||
"tags": ["c"],
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"payload": ["x", "y", "z"],
|
||||
"metadata": {"label": "list"},
|
||||
"tags": ["d", "e"],
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"payload": '{"foo": "bar"}',
|
||||
"metadata": {"label": "string"},
|
||||
"tags": ["f"],
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"payload": None,
|
||||
"metadata": {"label": "null"},
|
||||
"tags": ["g"],
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
rows = {row["id"]: row for row in table.to_arrow().to_pylist()}
|
||||
assert json.loads(rows[1]["payload"]) == {"foo": "bar", "count": 2}
|
||||
assert json.loads(rows[2]["payload"]) == {
|
||||
"name": "alice",
|
||||
"tags": ["x", "y"],
|
||||
"nested": {"enabled": True},
|
||||
}
|
||||
assert json.loads(rows[3]["payload"]) == ["x", "y", "z"]
|
||||
assert json.loads(rows[4]["payload"]) == {"foo": "bar"}
|
||||
assert rows[5]["payload"] is None
|
||||
assert rows[1]["metadata"] == {"label": "dict"}
|
||||
assert rows[1]["tags"] == ["a", "b"]
|
||||
|
||||
matched = table.search().where("json_extract(payload, '$.count') = '2'").to_list()
|
||||
assert [row["id"] for row in matched] == [1]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
def test_add_python_objects_to_nested_json_fields(mem_db: DBConnection):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"info",
|
||||
pa.struct([pa.field("payload", pa.json_())]),
|
||||
),
|
||||
pa.field("documents", pa.list_(pa.field("item", pa.json_()))),
|
||||
]
|
||||
)
|
||||
table = mem_db.create_table("nested_json_python_objects_add", schema=schema)
|
||||
|
||||
table.add(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"info": {"payload": {"kind": "struct", "value": 1}},
|
||||
"documents": [{"kind": "list", "value": 2}],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
row = table.to_arrow().to_pylist()[0]
|
||||
assert json.loads(row["info"]["payload"]) == {"kind": "struct", "value": 1}
|
||||
assert json.loads(row["documents"][0]) == {"kind": "list", "value": 2}
|
||||
assert [
|
||||
row["id"]
|
||||
for row in table.search()
|
||||
.where("json_extract(info.payload, '$.value') = '1'")
|
||||
.to_list()
|
||||
] == [1]
|
||||
assert [
|
||||
row["id"]
|
||||
for row in table.search()
|
||||
.where("json_extract(documents[1], '$.value') = '2'")
|
||||
.to_list()
|
||||
] == [1]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
def test_json_serialization_plan_skips_non_json_branches():
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("vector", pa.list_(pa.float32(), 3)),
|
||||
pa.field("payload", pa.json_()),
|
||||
pa.field(
|
||||
"metadata",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("label", pa.string()),
|
||||
pa.field("payload", pa.json_()),
|
||||
]
|
||||
),
|
||||
),
|
||||
pa.field(
|
||||
"documents",
|
||||
pa.list_(
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("title", pa.string()),
|
||||
pa.field("payload", pa.json_()),
|
||||
]
|
||||
)
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
plans = table_module._json_serialization_plans(schema)
|
||||
|
||||
assert set(plans) == {"payload", "metadata", "documents"}
|
||||
metadata_children = plans["metadata"].children
|
||||
assert metadata_children is not None
|
||||
assert set(metadata_children) == {"payload"}
|
||||
documents_item = plans["documents"].item
|
||||
assert documents_item is not None
|
||||
assert documents_item.children is not None
|
||||
assert set(documents_item.children) == {"payload"}
|
||||
|
||||
|
||||
@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type")
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_all_null_json_batch(mem_db_async: AsyncConnection):
|
||||
|
||||
Reference in New Issue
Block a user