diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 9393c57ae..3251f009e 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -478,6 +478,17 @@ Table hold your actual data as a collection of records / rows. ::: lancedb.materialized_view.AsyncMaterializedView +## Property Graphs (Asynchronous) + +`AsyncConnection` returns these handles from its `create_*` and `open_*` +property graph methods. + +::: lancedb.graph.AsyncPropertyGraph + +::: lancedb.graph.AsyncVirtualPropertyGraph + +::: lancedb.graph.AsyncMaterializedVirtualPropertyGraph + ## Indices (Asynchronous) Indices can be created on a table to speed up queries. This section diff --git a/python/python/lancedb/graph.py b/python/python/lancedb/graph.py index dc5ffc19b..73f95e89f 100644 --- a/python/python/lancedb/graph.py +++ b/python/python/lancedb/graph.py @@ -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 diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index d366bd821..5d39a2e3a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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", diff --git a/python/src/graph.rs b/python/src/graph.rs index 9e81889a1..7a4d862fd 100644 --- a/python/src/graph.rs +++ b/python/src/graph.rs @@ -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) -> PyResult { - 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())) } diff --git a/rust/lancedb/src/graph.rs b/rust/lancedb/src/graph.rs index b811ac404..59a8e0580 100644 --- a/rust/lancedb/src/graph.rs +++ b/rust/lancedb/src/graph.rs @@ -34,6 +34,34 @@ fn to_json(value: &impl Serialize) -> Result { }) } +/// `schema` in the JSON Arrow form a graph definition carries. +/// +/// The form spells no timestamp or duration unit, no time zone and no +/// dictionary, and cannot spell a time or an interval at all, so a schema that +/// would not read back as itself is refused here rather than reaching the +/// server as another schema. +pub fn schema_to_json(schema: &Schema) -> Result { + let json = lance_namespace::schema::arrow_schema_to_json(schema)?; + for (field, encoded) in schema.fields().iter().zip(&json.fields) { + let read_back = lance_namespace::schema::convert_json_arrow_field(encoded).ok(); + if read_back.as_ref() != Some(field.as_ref()) { + let as_read = match read_back { + Some(read_back) => format!("it would be read as {}", read_back.data_type()), + None => "it cannot be read back".to_string(), + }; + return Err(Error::InvalidInput { + message: format!( + "property '{}' of type {} cannot be carried by a property graph \ + definition: {as_read}", + field.name(), + field.data_type() + ), + }); + } + } + Ok(json) +} + /// A property graph's node and edge types. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct PropertyGraphDefinition { @@ -86,7 +114,7 @@ impl NodeType { Ok(Self { label: label.into(), key: key.into(), - schema: lance_namespace::schema::arrow_schema_to_json(schema)?, + schema: schema_to_json(schema)?, }) } @@ -125,9 +153,7 @@ impl EdgeType { label: label.into(), source, destination, - schema: schema - .map(lance_namespace::schema::arrow_schema_to_json) - .transpose()?, + schema: schema.map(schema_to_json).transpose()?, }) } @@ -297,7 +323,7 @@ impl MaterializedVirtualPropertyGraphDescription { #[cfg(test)] mod tests { - use arrow_schema::{DataType, Field}; + use arrow_schema::{DataType, Field, TimeUnit}; use super::*; @@ -417,6 +443,49 @@ mod tests { assert_eq!(decoded, description); } + /// A type the JSON form would read back as another type, or not at all, + /// is refused before any request carries it, in node and edge schemas. + #[test] + fn test_a_schema_the_definition_cannot_carry_is_refused() { + let endpoint = |column: &str| EndpointType { + label: "Person".to_string(), + column: column.to_string(), + }; + for data_type in [ + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + DataType::Timestamp(TimeUnit::Second, None), + DataType::Duration(TimeUnit::Nanosecond), + DataType::Time32(TimeUnit::Second), + DataType::Time64(TimeUnit::Microsecond), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::new_list(DataType::Timestamp(TimeUnit::Millisecond, None), true), + ] { + let schema = Schema::new(vec![ + Field::new("person_id", DataType::Int64, false), + Field::new("at", data_type.clone(), true), + ]); + let error = NodeType::new("Person", "person_id", &schema) + .unwrap_err() + .to_string(); + assert!(error.contains("property 'at'"), "{data_type}: {error}"); + let error = EdgeType::new("KNOWS", endpoint("src"), endpoint("dst"), Some(&schema)) + .unwrap_err() + .to_string(); + assert!(error.contains("property 'at'"), "{data_type}: {error}"); + } + + // What the form spells exactly still round-trips. + let schema = Schema::new(vec![ + Field::new("person_id", DataType::Int64, false), + Field::new("at", DataType::Timestamp(TimeUnit::Microsecond, None), true), + Field::new("took", DataType::Duration(TimeUnit::Microsecond), true), + Field::new("price", DataType::Decimal128(18, 4), true), + ]); + let person = NodeType::new("Person", "person_id", &schema).unwrap(); + assert_eq!(person.arrow_schema().unwrap(), schema); + } + #[tokio::test] async fn test_local_databases_refuse_property_graphs() { let dir = tempfile::tempdir().unwrap();