fix: refuse graph schemas a definition cannot carry, and take any Table.add input on insert

A node or edge type whose schema would read back as another type from its JSON form (timestamp units and time zones, durations, times, dictionaries) is refused before a request carries it. Inserting into a property graph registers DataFrame converters first, as Table.add does, and the async graph handles are in the API reference.
This commit is contained in:
Jack Ye
2026-09-29 09:13:42 -05:00
parent 973a793d18
commit 5f4d448a34
5 changed files with 180 additions and 8 deletions
+4 -1
View File
@@ -33,7 +33,7 @@ import pyarrow as pa
from ._lancedb import graph_schema_from_json, graph_schema_to_json
from .background_loop import LOOP
from .scannable import to_scannable
from .scannable import _register_optional_converters, to_scannable
if TYPE_CHECKING:
from .common import DATA
@@ -409,6 +409,9 @@ class AsyncPropertyGraph(_AsyncGraph):
the edge type names; an edge whose endpoint is not a node of the graph
refuses the whole insert.
"""
# A DataFrame library imported after lancedb registers its converter
# here, as it does for Table.add.
_register_optional_converters()
return _property_description(
await self._inner.insert_into_property_graph(
self._name, label, to_scannable(data), self._namespace_path
+88
View File
@@ -8,7 +8,9 @@ import json
import multiprocessing as mp
import pickle
import re
import subprocess
import sys
import textwrap
import threading
import time
from unittest.mock import MagicMock, patch
@@ -3079,6 +3081,92 @@ def test_a_property_graph_carries_schemas_and_takes_arrow_rows():
assert drop[:2] == ("POST", "/v1/property_graph/people/drop")
def test_a_graph_refuses_a_schema_its_definition_cannot_carry():
# A type the definition's JSON form would read back as another type is
# refused before any request carries it, in node and edge schemas.
calls = []
person = NodeType(
"Person",
key="person_id",
schema=pa.schema(
[
pa.field("person_id", pa.int64(), nullable=False),
("at", pa.timestamp("ns", tz="UTC")),
]
),
)
plain = NodeType(
"Person",
key="person_id",
schema=pa.schema([pa.field("person_id", pa.int64(), nullable=False)]),
)
knows = EdgeType(
"KNOWS",
source=("Person", "src_id"),
destination=("Person", "dst_id"),
schema=pa.schema([("took", pa.duration("ns"))]),
)
with mock_lancedb_connection(_graph_handler(calls, lambda path, body: {})) as db:
with pytest.raises(ValueError, match="property 'at'"):
db.create_property_graph("people", [person], [])
with pytest.raises(ValueError, match="property 'took'"):
db.create_property_graph("people", [plain], [knows])
assert calls == []
def test_a_property_graph_inserts_a_dataframe_imported_after_lancedb():
# A DataFrame library imported after lancedb has its converter registered
# at the insert boundary, as Table.add does. A fresh interpreter makes the
# import order the test's own.
pytest.importorskip("polars")
script = textwrap.dedent(
"""\
import asyncio
import json
import lancedb
import polars as pl
from lancedb.graph import AsyncPropertyGraph
from lancedb.scannable import Scannable
class Inner:
async def insert_into_property_graph(self, name, label, data, namespace):
names = data.schema.names if isinstance(data, Scannable) else None
if names != ["person_id"]:
raise SystemExit(f"unexpected rows: {data!r}")
return json.dumps(
{
"name": name,
"namespace": namespace,
"nodes": [],
"edges": [],
"commit": "c1",
"committed_at": "2026-09-29T00:00:00Z",
"vertex_count": 1,
"edge_count": 0,
}
)
described = asyncio.run(
AsyncPropertyGraph(Inner(), "people", []).insert(
"Person", pl.DataFrame({"person_id": [1]})
)
)
if described.vertex_count != 1:
raise SystemExit(f"unexpected description: {described}")
"""
)
result = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
def test_a_graph_takes_only_its_own_kind_of_definition():
person = NodeType(
"Person",
+3 -2
View File
@@ -29,10 +29,11 @@ impl GraphKind {
}
}
/// A pyarrow schema as the JSON Arrow schema of a node or edge type.
/// A pyarrow schema as the JSON Arrow schema of a node or edge type, refused
/// when that form would not read back as the same schema.
#[pyfunction]
pub fn graph_schema_to_json(schema: PyArrowType<Schema>) -> PyResult<String> {
let json = lance_namespace::schema::arrow_schema_to_json(&schema.0)
let json = lancedb::graph::schema_to_json(&schema.0)
.map_err(|error| PyValueError::new_err(error.to_string()))?;
serde_json::to_string(&json).map_err(|error| PyValueError::new_err(error.to_string()))
}