diff --git a/docs/src/python/python.md b/docs/src/python/python.md index ccc44a55e..6efb278fb 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -202,6 +202,22 @@ listing a storage directory. ::: lancedb.view.ViewDescription +## Property Graphs + +A property graph reads tables as the nodes and edges of a graph. Manage one +with [create_property_graph][lancedb.DBConnection.create_property_graph], +[describe_property_graph][lancedb.DBConnection.describe_property_graph], +[list_property_graphs][lancedb.DBConnection.list_property_graphs] and +[drop_property_graph][lancedb.DBConnection.drop_property_graph]. + +::: lancedb.graph.NodeTable + +::: lancedb.graph.EdgeTable + +::: lancedb.graph.Endpoint + +::: lancedb.graph.PropertyGraphDescription + ## Expressions Type-safe expression builder for filters and projections. Use these instead diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index d34b541a8..ec93cbd5d 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -48,6 +48,7 @@ from .materialized_view import ( MaterializedViewDefinition, ) from .view import ViewDescription as ViewDescription +from .graph import PropertyGraphDescription as PropertyGraphDescription from .table import AsyncTable, Table from .types import BaseTokenizerType from ._lancedb import Session @@ -584,6 +585,7 @@ __all__ = [ "MaterializedView", "MaterializedViewDefinition", "ViewDescription", + "PropertyGraphDescription", "connect", "connect_async", "tokenize", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 1f26ad6f4..1df17d503 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -186,6 +186,21 @@ class Connection(object): async def list_views( self, namespace_path: Optional[List[str]] = None ) -> List[str]: ... + async def create_property_graph( + self, + name: str, + definition_json: str, + namespace_path: Optional[List[str]] = None, + ) -> str: ... + async def describe_property_graph( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> str: ... + async def drop_property_graph( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> None: ... + async def list_property_graphs( + self, namespace_path: Optional[List[str]] = None + ) -> List[str]: ... async def list_jobs(self) -> List[JobInfo]: ... async def cancel_job(self, job_id: str) -> bool: ... async def pause_job(self, job_id: str) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 02de75460..045a67e72 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -49,6 +49,13 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from . import __version__ from ._lancedb import connect as lancedb_connect # type: ignore from .functions import FunctionVersion, UdfDefinition +from .graph import ( + EdgeTable, + NodeTable, + PropertyGraphDescription, + _definition_json, + _description_from_json, +) from .job import AsyncJob, Job, _typed_job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery @@ -1018,6 +1025,79 @@ class DBConnection(EnforceOverrides): "View operations are not supported for this connection type" ) + def create_property_graph( + self, + name: str, + nodes: List[NodeTable], + edges: List[EdgeTable], + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraphDescription: + """Create a property graph: tables in its namespace read as the nodes + and edges of a graph. + + Each node table names its key column and label. Each edge table names + its label and the node keys its source and destination columns + reference. A name already taken is an error. Local connections raise + ``NotImplementedError``. + + >>> import lancedb + >>> from lancedb.graph import EdgeTable, Endpoint, NodeTable + >>> db = lancedb.connect("db://my_database") # doctest: +SKIP + >>> graph = db.create_property_graph( + ... "social", + ... nodes=[NodeTable("person", key="person_id", label="Person")], + ... edges=[ + ... EdgeTable( + ... "knows", + ... label="KNOWS", + ... source=Endpoint("src_id", references=("person", "person_id")), + ... destination=Endpoint( + ... "dst_id", references=("person", "person_id") + ... ), + ... ) + ... ], + ... ) # doctest: +SKIP + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def describe_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraphDescription: + """What this database records about a property graph: its node and + edge tables, its size, and the table versions it was built from. + + Local connections raise ``NotImplementedError``. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + """Drop a property graph and wait for its data to be deleted. + + The tables it reads are untouched. Local connections raise + ``NotImplementedError``. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + + def list_property_graphs( + self, *, namespace_path: Optional[List[str]] = None + ) -> List[str]: + """The names of the property graphs in one namespace. + + Local connections raise ``NotImplementedError``. + """ + raise NotImplementedError( + "Property graph operations are not supported for this connection type" + ) + def open_job(self, job_id: str) -> Job: """Open a server-side job by id, returning a handle with its record already populated. @@ -1883,6 +1963,41 @@ class LanceDBConnection(DBConnection): def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: return LOOP.run(self._conn.list_views(namespace_path=namespace_path)) + @override + def create_property_graph( + self, + name: str, + nodes: List[NodeTable], + edges: List[EdgeTable], + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraphDescription: + return LOOP.run( + self._conn.create_property_graph( + name, nodes, edges, namespace_path=namespace_path + ) + ) + + @override + def describe_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraphDescription: + return LOOP.run( + self._conn.describe_property_graph(name, namespace_path=namespace_path) + ) + + @override + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run(self._conn.drop_property_graph(name, namespace_path=namespace_path)) + + @override + def list_property_graphs( + self, *, namespace_path: Optional[List[str]] = None + ) -> List[str]: + return LOOP.run(self._conn.list_property_graphs(namespace_path=namespace_path)) + @override def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" @@ -2910,6 +3025,45 @@ class AsyncConnection(object): """The names of the views in one namespace.""" return await self._inner.list_views(list(namespace_path or [])) + async def create_property_graph( + self, + name: str, + nodes: List[NodeTable], + edges: List[EdgeTable], + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraphDescription: + """Create a property graph over tables in its namespace. + + See + [DBConnection.create_property_graph][lancedb.DBConnection.create_property_graph]. + """ + return _description_from_json( + await self._inner.create_property_graph( + name, _definition_json(nodes, edges), list(namespace_path or []) + ) + ) + + async def describe_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraphDescription: + """What this database records about a property graph.""" + return _description_from_json( + await self._inner.describe_property_graph(name, list(namespace_path or [])) + ) + + async def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + """Drop a property graph and wait for its data to be deleted.""" + await self._inner.drop_property_graph(name, list(namespace_path or [])) + + async def list_property_graphs( + self, *, namespace_path: Optional[List[str]] = None + ) -> List[str]: + """The names of the property graphs in one namespace.""" + return await self._inner.list_property_graphs(list(namespace_path or [])) + async def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" return await self._inner.list_jobs() diff --git a/python/python/lancedb/graph.py b/python/python/lancedb/graph.py new file mode 100644 index 000000000..0d6c8bd8c --- /dev/null +++ b/python/python/lancedb/graph.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Property graphs: tables read as the nodes and edges of a graph. + +A property graph is a catalog object in a namespace, next to the tables it +reads. Its definition says which tables hold nodes and which hold edges: each +node table's key and label, each edge table's label, and how its source and +destination columns reference node keys -- the shape of SQL/PGQ's +``CREATE PROPERTY GRAPH``. The tables stay ordinary tables. See +``DBConnection.create_property_graph``. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple + + +@dataclass(frozen=True) +class NodeTable: + """A table whose rows are the nodes of one label.""" + + table: str + """The table, in the graph's namespace.""" + key: str + """The column that identifies a node.""" + label: str + """The label its nodes carry.""" + properties: Optional[List[str]] = None + """The columns exposed as node properties; every column when ``None``.""" + + +@dataclass(frozen=True) +class Endpoint: + """One end of an edge: a column of the edge table and the node key it holds.""" + + column: str + """The edge table's column.""" + references: Tuple[str, str] + """The node table and its key column, as ``(table, column)``.""" + + +@dataclass(frozen=True) +class EdgeTable: + """A table whose rows are the edges of one label.""" + + table: str + """The table, in the graph's namespace.""" + label: str + """The label its edges carry.""" + source: Endpoint + """The column naming each edge's source node.""" + destination: Endpoint + """The column naming each edge's destination node.""" + properties: Optional[List[str]] = None + """The columns exposed as edge properties; every column when ``None``.""" + + +@dataclass(frozen=True) +class PropertyGraphDescription: + """What a database records about one property graph.""" + + name: str + """The graph's name within its namespace.""" + namespace_path: List[str] + """The namespace holding the graph; empty is the root namespace.""" + nodes: List[NodeTable] + """The node tables.""" + edges: List[EdgeTable] + """The edge tables.""" + vertex_count: int + """How many nodes the graph held when it was built.""" + edge_count: int + """How many edges the graph held when it was built.""" + sources: Dict[str, int] + """The version of each table the graph was built from, by table name.""" + + +def _definition_json(nodes: Sequence[NodeTable], edges: Sequence[EdgeTable]) -> str: + return json.dumps( + { + "nodes": [_node_to_json(node) for node in nodes], + "edges": [_edge_to_json(edge) for edge in edges], + } + ) + + +def _description_from_json(text: str) -> PropertyGraphDescription: + described = json.loads(text) + return PropertyGraphDescription( + name=described["name"], + namespace_path=list(described.get("namespace", [])), + nodes=[_node_from_json(node) for node in described["nodes"]], + edges=[_edge_from_json(edge) for edge in described.get("edges", [])], + vertex_count=described["vertex_count"], + edge_count=described["edge_count"], + sources={ + source["table"]: source["version"] + for source in described.get("sources", []) + }, + ) + + +def _node_to_json(node: NodeTable) -> Dict[str, Any]: + encoded: Dict[str, Any] = { + "table": node.table, + "key": node.key, + "label": node.label, + } + if node.properties is not None: + encoded["properties"] = list(node.properties) + return encoded + + +def _edge_to_json(edge: EdgeTable) -> Dict[str, Any]: + encoded: Dict[str, Any] = { + "table": edge.table, + "label": edge.label, + "source": _endpoint_to_json(edge.source), + "destination": _endpoint_to_json(edge.destination), + } + if edge.properties is not None: + encoded["properties"] = list(edge.properties) + return encoded + + +def _endpoint_to_json(endpoint: Endpoint) -> Dict[str, Any]: + table, column = endpoint.references + return { + "column": endpoint.column, + "references": {"table": table, "column": column}, + } + + +def _node_from_json(node: Dict[str, Any]) -> NodeTable: + return NodeTable( + table=node["table"], + key=node["key"], + label=node["label"], + properties=node.get("properties"), + ) + + +def _edge_from_json(edge: Dict[str, Any]) -> EdgeTable: + return EdgeTable( + table=edge["table"], + label=edge["label"], + source=_endpoint_from_json(edge["source"]), + destination=_endpoint_from_json(edge["destination"]), + properties=edge.get("properties"), + ) + + +def _endpoint_from_json(endpoint: Dict[str, Any]) -> Endpoint: + references = endpoint["references"] + return Endpoint( + column=endpoint["column"], + references=(references["table"], references["column"]), + ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 69c5a62b0..a657e77bb 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -36,6 +36,7 @@ import pyarrow as pa from ..common import DATA from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition +from ..graph import EdgeTable, NodeTable, PropertyGraphDescription from ..job import AsyncJob, Job from ..sql import Query as SqlQuery from ..sql import QueryDescription @@ -942,6 +943,41 @@ class RemoteDBConnection(DBConnection): def list_views(self, *, namespace_path: Optional[List[str]] = None) -> List[str]: return LOOP.run(self._conn.list_views(namespace_path=namespace_path)) + @override + def create_property_graph( + self, + name: str, + nodes: List[NodeTable], + edges: List[EdgeTable], + *, + namespace_path: Optional[List[str]] = None, + ) -> PropertyGraphDescription: + return LOOP.run( + self._conn.create_property_graph( + name, nodes, edges, namespace_path=namespace_path + ) + ) + + @override + def describe_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> PropertyGraphDescription: + return LOOP.run( + self._conn.describe_property_graph(name, namespace_path=namespace_path) + ) + + @override + def drop_property_graph( + self, name: str, *, namespace_path: Optional[List[str]] = None + ) -> None: + LOOP.run(self._conn.drop_property_graph(name, namespace_path=namespace_path)) + + @override + def list_property_graphs( + self, *, namespace_path: Optional[List[str]] = None + ) -> List[str]: + return LOOP.run(self._conn.list_property_graphs(namespace_path=namespace_path)) + @override def list_jobs(self) -> List["JobInfo"]: """List server-side jobs across the database's tables.""" diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index e32dbe8ca..ed4b3b60c 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -17,6 +17,7 @@ from packaging.version import Version import lancedb from lancedb.conftest import MockTextEmbeddingFunction +from lancedb.graph import EdgeTable, Endpoint, NodeTable, PropertyGraphDescription from lancedb.query import ColumnOrdering from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError @@ -2828,3 +2829,98 @@ def test_view_crud_addresses_its_own_routes(): ("GET", "/v1/namespace/analytics/view/list"), ("POST", "/v1/view/analytics$adults/drop"), ] + + +def test_property_graph_crud_addresses_its_own_routes(): + # The graph verbs are their own routes: the definition is the create body, + # and a description adds the graph's size and the table versions it was + # built from. + calls = [] + person = {"table": "person", "column": "person_id"} + definition = { + "nodes": [ + { + "table": "person", + "key": "person_id", + "label": "Person", + "properties": ["name"], + } + ], + "edges": [ + { + "table": "knows", + "label": "KNOWS", + "source": {"column": "src_id", "references": person}, + "destination": {"column": "dst_id", "references": person}, + } + ], + } + described = { + "name": "social", + "namespace": ["analytics"], + **definition, + "vertex_count": 4, + "edge_count": 5, + "sources": [ + {"table": "person", "version": 3}, + {"table": "knows", "version": 2}, + ], + } + + def handler(request): + length = int(request.headers.get("Content-Length") or 0) + body = json.loads(request.rfile.read(length)) if length else None + calls.append((request.command, request.path, body)) + if request.path.endswith("/property_graph/list"): + response = {"property_graphs": ["social"], "page_token": None} + elif request.path.endswith("/drop"): + response = {} + else: + response = described + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(response).encode()) + + nodes = [NodeTable("person", key="person_id", label="Person", properties=["name"])] + edges = [ + EdgeTable( + "knows", + label="KNOWS", + source=Endpoint("src_id", references=("person", "person_id")), + destination=Endpoint("dst_id", references=("person", "person_id")), + ) + ] + with mock_lancedb_connection(handler) as db: + graph = db.create_property_graph( + "social", nodes, edges, namespace_path=["analytics"] + ) + assert graph == PropertyGraphDescription( + name="social", + namespace_path=["analytics"], + nodes=nodes, + edges=edges, + vertex_count=4, + edge_count=5, + sources={"person": 3, "knows": 2}, + ) + assert db.describe_property_graph("social", namespace_path=["analytics"]) == ( + graph + ) + assert db.list_property_graphs(namespace_path=["analytics"]) == ["social"] + db.drop_property_graph("social", namespace_path=["analytics"]) + + assert calls == [ + ("POST", "/v1/property_graph/analytics$social/create", definition), + ("POST", "/v1/property_graph/analytics$social/describe", None), + ("GET", "/v1/namespace/analytics/property_graph/list", None), + ("POST", "/v1/property_graph/analytics$social/drop", None), + ] + + +def test_local_connections_refuse_property_graphs(tmp_path): + db = lancedb.connect(tmp_path) + with pytest.raises(NotImplementedError, match="Property graph operations"): + db.list_property_graphs() + with pytest.raises(NotImplementedError, match="Property graph operations"): + db.describe_property_graph("social") diff --git a/python/src/connection.rs b/python/src/connection.rs index 2959fe037..707bb1e16 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -965,6 +965,76 @@ impl Connection { }) } + #[pyo3(signature = (name, definition_json, namespace_path=None))] + pub fn create_property_graph( + self_: PyRef<'_, Self>, + name: String, + definition_json: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + let definition = + lancedb::graph::PropertyGraphDefinition::from_json(&definition_json).infer_error()?; + future_into_py(self_.py(), async move { + inner + .create_property_graph(name, &definition, &namespace_path) + .await + .infer_error()? + .to_json() + .infer_error() + }) + } + + #[pyo3(signature = (name, namespace_path=None))] + pub fn describe_property_graph( + self_: PyRef<'_, Self>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .describe_property_graph(name, &namespace_path) + .await + .infer_error()? + .to_json() + .infer_error() + }) + } + + #[pyo3(signature = (name, namespace_path=None))] + pub fn drop_property_graph( + self_: PyRef<'_, Self>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .drop_property_graph(name, &namespace_path) + .await + .infer_error() + }) + } + + #[pyo3(signature = (namespace_path=None))] + pub fn list_property_graphs( + self_: PyRef<'_, Self>, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let namespace_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .list_property_graphs(&namespace_path) + .await + .infer_error() + }) + } + pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 0e5bbf5da..5134f1bc8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -28,6 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -38,8 +39,8 @@ use crate::remote::{ }; use crate::secrets::SecretInfo; use crate::utils::{ - validate_namespace, validate_secret_component, validate_secret_reference, - validate_view_reference, + validate_namespace, validate_property_graph_reference, validate_secret_component, + validate_secret_reference, validate_view_reference, }; use crate::view::ViewDescription; use lance::io::ObjectStoreParams; @@ -866,6 +867,92 @@ impl Connection { self.internal.list_views(namespace_path).await } + /// Create a property graph: tables in its namespace read as the nodes and + /// edges of a graph. + /// + /// The definition names the node tables with their keys and labels, and + /// the edge tables with the node keys their source and destination columns + /// reference. A name already taken is an error. Local databases return + /// [`Error::NotSupported`]. + /// + /// ```no_run + /// # use lancedb::graph::{EdgeTable, Endpoint, EndpointReference, NodeTable, + /// # PropertyGraphDefinition}; + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// let person = EndpointReference { + /// table: "person".to_string(), + /// column: "person_id".to_string(), + /// }; + /// let definition = PropertyGraphDefinition { + /// nodes: vec![NodeTable { + /// table: "person".to_string(), + /// key: "person_id".to_string(), + /// label: "Person".to_string(), + /// properties: None, + /// }], + /// edges: vec![EdgeTable { + /// table: "knows".to_string(), + /// label: "KNOWS".to_string(), + /// source: Endpoint { column: "src_id".to_string(), references: person.clone() }, + /// destination: Endpoint { column: "dst_id".to_string(), references: person }, + /// properties: None, + /// }], + /// }; + /// let graph = connection.create_property_graph("social", &definition, &[]).await?; + /// println!("{} nodes, {} edges", graph.vertex_count, graph.edge_count); + /// # Ok(()) + /// # } + /// ``` + pub async fn create_property_graph( + &self, + name: impl AsRef, + definition: &PropertyGraphDefinition, + namespace_path: &[String], + ) -> Result { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .create_property_graph(name.as_ref(), definition, namespace_path) + .await + } + + /// What this database records about one property graph: its definition, + /// its size, and the table versions it was built from. Local databases + /// return [`Error::NotSupported`]. + pub async fn describe_property_graph( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .describe_property_graph(name.as_ref(), namespace_path) + .await + } + + /// Drop a property graph and wait for its data to be deleted. + /// + /// The tables it reads are untouched. Local databases return + /// [`Error::NotSupported`]. + pub async fn drop_property_graph( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result<()> { + validate_property_graph_reference(name.as_ref(), namespace_path)?; + self.internal + .drop_property_graph(name.as_ref(), namespace_path) + .await + } + + /// The names of the property graphs in one namespace. + /// + /// The client walks all server pages before returning. Local databases + /// return [`Error::NotSupported`]. + pub async fn list_property_graphs(&self, namespace_path: &[String]) -> Result> { + validate_namespace(namespace_path)?; + self.internal.list_property_graphs(namespace_path).await + } + /// Rename a table in the database. /// /// This is only supported in LanceDB Cloud. diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 224fd8770..c1de55756 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -28,6 +28,7 @@ use lance_namespace::models::{ use crate::data::scannable::Scannable; use crate::error::Result; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; use crate::job::Job; use crate::materialized_view::CreateMaterializedViewRequest; use crate::secrets::SecretInfo; @@ -288,6 +289,12 @@ fn view_ops_not_supported() -> Result { }) } +fn property_graph_ops_not_supported() -> Result { + Err(crate::error::Error::NotSupported { + message: "Property graph operations are not supported by this database".to_string(), + }) +} + /// The `Database` trait defines the interface for database implementations. /// /// A database is responsible for managing tables and their metadata. @@ -507,6 +514,42 @@ pub trait Database: async fn list_views(&self, _namespace_path: &[String]) -> Result> { view_ops_not_supported() } + /// Create a property graph over tables in its namespace. + async fn create_property_graph( + &self, + _name: &str, + _definition: &PropertyGraphDefinition, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// What the database records about one property graph. + async fn describe_property_graph( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// Drop a property graph and wait for its data to be deleted. The tables it + /// reads are untouched. + async fn drop_property_graph(&self, _name: &str, _namespace_path: &[String]) -> Result<()> { + property_graph_ops_not_supported() + } + /// Drop a property graph and return the job deleting its data, without + /// waiting. + #[doc(hidden)] + async fn drop_property_graph_async( + &self, + _name: &str, + _namespace_path: &[String], + ) -> Result { + property_graph_ops_not_supported() + } + /// The names of the property graphs in one namespace. + async fn list_property_graphs(&self, _namespace_path: &[String]) -> Result> { + property_graph_ops_not_supported() + } /// Open a job by id, returning a handle with its record already /// populated. Fails with [`crate::Error::JobNotFound`] when the server has /// no such job. diff --git a/rust/lancedb/src/graph.rs b/rust/lancedb/src/graph.rs new file mode 100644 index 000000000..7b82b8a0b --- /dev/null +++ b/rust/lancedb/src/graph.rs @@ -0,0 +1,210 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Property graphs: tables read as the nodes and edges of a graph. +//! +//! A property graph is a catalog object in a namespace, next to the tables it +//! reads. Its definition says which tables hold nodes and which hold edges: +//! each node table's key and label, each edge table's label, and how its +//! source and destination columns reference node keys -- the shape of SQL/PGQ's +//! `CREATE PROPERTY GRAPH`. The tables themselves stay ordinary tables. +//! +//! The verbs live on [`crate::connection::Connection`]. + +use serde::{Deserialize, Serialize}; + +use crate::error::{Error, Result}; + +/// A table whose rows are the nodes of one label. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct NodeTable { + /// The table, in the graph's namespace. + pub table: String, + /// The column that identifies a node. + pub key: String, + /// The label its nodes carry. + pub label: String, + /// The columns exposed as node properties; every column when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub properties: Option>, +} + +/// A table whose rows are the edges of one label. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct EdgeTable { + /// The table, in the graph's namespace. + pub table: String, + /// The label its edges carry. + pub label: String, + /// The column naming each edge's source node. + pub source: Endpoint, + /// The column naming each edge's destination node. + pub destination: Endpoint, + /// The columns exposed as edge properties; every column when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub properties: Option>, +} + +/// One end of an edge: a column of the edge table and the node key it holds. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct Endpoint { + /// The edge table's column. + pub column: String, + /// The node table and key column its values match. + pub references: EndpointReference, +} + +/// The node table and key column an endpoint's values match. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct EndpointReference { + /// A node table of the same graph. + pub table: String, + /// That node table's key column. + pub column: String, +} + +/// The node and edge tables that make up a property graph. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropertyGraphDefinition { + /// The node tables, one per label. + pub nodes: Vec, + /// The edge tables, one per label. + #[serde(default)] + pub edges: Vec, +} + +impl PropertyGraphDefinition { + /// Parse a definition from its JSON form. + pub fn from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(|source| Error::InvalidInput { + message: format!("invalid property graph definition: {source}"), + }) + } +} + +/// The version of one table a graph was built from. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct GraphSourceVersion { + /// The table. + pub table: String, + /// Its version when the graph was built. + pub version: u64, +} + +/// What a database records about one property graph. +/// +/// Returned by [`crate::connection::Connection::describe_property_graph`], and +/// by `create_property_graph`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct PropertyGraphDescription { + /// The graph's name within its namespace. + pub name: String, + /// The namespace holding the graph; empty is the root namespace. + #[serde(rename = "namespace", default)] + pub namespace_path: Vec, + /// The node tables. + pub nodes: Vec, + /// The edge tables. + #[serde(default)] + pub edges: Vec, + /// How many nodes the graph held when it was built. + pub vertex_count: u64, + /// How many edges the graph held when it was built. + pub edge_count: u64, + /// The table versions the graph was built from. + #[serde(default)] + pub sources: Vec, +} + +impl PropertyGraphDescription { + /// The description in its JSON form. + pub fn to_json(&self) -> Result { + serde_json::to_string(self).map_err(|source| Error::Runtime { + message: format!("could not encode a property graph description: {source}"), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn social() -> PropertyGraphDefinition { + let person = EndpointReference { + table: "person".to_string(), + column: "person_id".to_string(), + }; + PropertyGraphDefinition { + nodes: vec![NodeTable { + table: "person".to_string(), + key: "person_id".to_string(), + label: "Person".to_string(), + properties: Some(vec!["name".to_string()]), + }], + edges: vec![EdgeTable { + table: "knows".to_string(), + label: "KNOWS".to_string(), + source: Endpoint { + column: "src_id".to_string(), + references: person.clone(), + }, + destination: Endpoint { + column: "dst_id".to_string(), + references: person, + }, + properties: None, + }], + } + } + + #[test] + fn test_definition_json_is_the_sql_pgq_shape() { + let json = serde_json::to_value(social()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "nodes": [{"table": "person", "key": "person_id", "label": "Person", + "properties": ["name"]}], + "edges": [{"table": "knows", "label": "KNOWS", + "source": {"column": "src_id", + "references": {"table": "person", "column": "person_id"}}, + "destination": {"column": "dst_id", + "references": {"table": "person", + "column": "person_id"}}}] + }) + ); + assert_eq!( + PropertyGraphDefinition::from_json(&json.to_string()).unwrap(), + social() + ); + assert!(PropertyGraphDefinition::from_json(r#"{"edges": []}"#).is_err()); + } + + #[tokio::test] + async fn test_local_databases_refuse_property_graphs() { + let dir = tempfile::tempdir().unwrap(); + let conn = crate::connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let refused = |error: Error| { + assert!( + matches!(error, Error::NotSupported { ref message } + if message.contains("Property graph operations")), + "{error}" + ); + }; + refused( + conn.create_property_graph("social", &social(), &[]) + .await + .unwrap_err(), + ); + refused( + conn.describe_property_graph("social", &[]) + .await + .unwrap_err(), + ); + refused(conn.drop_property_graph("social", &[]).await.unwrap_err()); + refused(conn.list_property_graphs(&[]).await.unwrap_err()); + } +} diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 530328cbb..8f9bce116 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -183,6 +183,7 @@ pub mod embeddings; pub mod error; pub mod expr; pub mod function; +pub mod graph; pub mod index; pub mod io; pub mod ipc; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 39bb5a528..ecbc410d8 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -28,6 +28,7 @@ use crate::function::{ FunctionArtifactRequest, FunctionRegistrationRequest, FunctionSignature, FunctionVersion, PythonRuntimeSpec, }; +use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription}; use crate::job::Job; use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::job::{PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client}; @@ -839,6 +840,14 @@ struct RemoteListViewsResponse { page_token: Option, } +#[derive(serde::Deserialize)] +struct RemoteListPropertyGraphsResponse { + #[serde(default)] + property_graphs: Vec, + #[serde(default)] + page_token: Option, +} + /// Bound on `list_jobs` page walking; a warning is logged when the listing /// is truncated at this many pages. const MAX_LIST_JOBS_PAGES: usize = 100; @@ -1278,6 +1287,111 @@ impl Database for RemoteDatabase { Ok(views) } + async fn create_property_graph( + &self, + name: &str, + definition: &PropertyGraphDefinition, + namespace_path: &[String], + ) -> Result { + let graph_id = build_object_identifier("Property graph name", name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/create")) + .json(definition); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + response.json().await.err_to_http(request_id) + } + + async fn describe_property_graph( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + let graph_id = build_object_identifier("Property graph name", name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/describe")); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + response.json().await.err_to_http(request_id) + } + + async fn drop_property_graph(&self, name: &str, namespace_path: &[String]) -> Result<()> { + self.drop_property_graph_async(name, namespace_path) + .await? + .wait() + .await + } + + async fn drop_property_graph_async( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + let graph_id = build_object_identifier("Property graph name", name, namespace_path)?; + let req = self + .client + .post(&format!("/v1/property_graph/{graph_id}/drop")); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + match status { + // Nothing was bound to the name, so nothing is being deleted. + StatusCode::OK => Ok(Job::new_done()), + StatusCode::ACCEPTED => { + let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { + source: "property graph drop response did not contain a valid job_id".into(), + request_id, + status_code: Some(status), + })?; + Ok(Job::new(Box::new(RemoteJob::new( + self.client.clone(), + job_id, + )))) + } + _ => Err(Error::Http { + source: "property graph drop must return 200 OK or 202 Accepted".into(), + request_id, + status_code: Some(status), + }), + } + } + + async fn list_property_graphs(&self, namespace_path: &[String]) -> Result> { + let namespace_id = build_namespace_identifier(namespace_path)?; + let path = format!("/v1/namespace/{namespace_id}/property_graph/list"); + let mut graphs = Vec::new(); + let mut page_token: Option = None; + let mut seen_page_tokens = HashSet::new(); + loop { + let mut req = self.client.get(&path); + if let Some(token) = &page_token { + req = req.query(&[("page_token", token)]); + } + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let response: RemoteListPropertyGraphsResponse = + response.json().await.err_to_http(request_id.clone())?; + graphs.extend(response.property_graphs); + let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) + else { + break; + }; + if !seen_page_tokens.insert(next_page_token.clone()) { + return Err(Error::Http { + source: "Property graph listing response repeated a page_token".into(), + request_id, + status_code: Some(status), + }); + } + page_token = Some(next_page_token); + } + Ok(graphs) + } + async fn open_job(&self, job_id: &str) -> Result { let handle = super::job::RemoteJob::new(self.client.clone(), job_id.to_string()); match crate::job::JobHandle::describe(&handle).await { @@ -4280,6 +4394,192 @@ mod tests { } } + fn social_graph() -> crate::graph::PropertyGraphDefinition { + crate::graph::PropertyGraphDefinition::from_json( + r#"{"nodes": [{"table": "person", "key": "person_id", "label": "Person"}], + "edges": [{"table": "knows", "label": "KNOWS", + "source": {"column": "src_id", + "references": {"table": "person", "column": "person_id"}}, + "destination": {"column": "dst_id", + "references": {"table": "person", + "column": "person_id"}}, + "properties": ["since"]}]}"#, + ) + .unwrap() + } + + fn property_graph_description_body(name: &str, namespace: &[&str]) -> String { + let mut body = serde_json::to_value(social_graph()).unwrap(); + let fields = body.as_object_mut().unwrap(); + fields.insert("name".into(), name.into()); + fields.insert("namespace".into(), namespace.into()); + fields.insert("vertex_count".into(), 4.into()); + fields.insert("edge_count".into(), 5.into()); + fields.insert( + "sources".into(), + serde_json::json!([{"table": "person", "version": 3}, {"table": "knows", "version": 2}]), + ); + body.to_string() + } + + #[tokio::test] + async fn test_create_property_graph_posts_the_definition() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!( + request.url().path(), + "/v1/property_graph/analytics$social/create" + ); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body, serde_json::to_value(social_graph()).unwrap()); + http::Response::builder() + .status(200) + .body(property_graph_description_body("social", &["analytics"])) + .unwrap() + }); + let graph = conn + .create_property_graph("social", &social_graph(), &["analytics".into()]) + .await + .unwrap(); + assert_eq!(graph.name, "social"); + assert_eq!(graph.namespace_path, vec!["analytics".to_string()]); + assert_eq!(graph.nodes, social_graph().nodes); + assert_eq!(graph.edges, social_graph().edges); + assert_eq!((graph.vertex_count, graph.edge_count), (4, 5)); + assert_eq!( + graph.sources[1], + crate::graph::GraphSourceVersion { + table: "knows".to_string(), + version: 2 + } + ); + } + + #[tokio::test] + async fn test_describe_and_drop_address_the_property_graph_in_the_path() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/property_graph/social/describe"); + assert!(request.body().is_none(), "{:?}", request.body()); + http::Response::builder() + .status(200) + .body(property_graph_description_body("social", &[])) + .unwrap() + }); + let graph = conn.describe_property_graph("social", &[]).await.unwrap(); + assert!(graph.namespace_path.is_empty()); + assert_eq!(graph.edges[0].properties, Some(vec!["since".to_string()])); + + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!( + request.url().path(), + "/v1/property_graph/analytics$social/drop" + ); + assert!(request.body().is_none(), "{:?}", request.body()); + http::Response::builder().status(200).body("{}").unwrap() + }); + conn.drop_property_graph("social", &["analytics".into()]) + .await + .unwrap(); + } + + #[tokio::test] + async fn test_drop_property_graph_async_reports_the_cleanup_job() { + let db = super::RemoteDatabase::new_mock(|_| { + http::Response::builder() + .status(202) + .body(r#"{"job_id":"j1-do-graph"}"#) + .unwrap() + }); + let job = db.drop_property_graph_async("social", &[]).await.unwrap(); + assert_eq!(job.id(), Some("j1-do-graph")); + + let db = super::RemoteDatabase::new_mock(|_| { + http::Response::builder().status(200).body("{}").unwrap() + }); + let job = db.drop_property_graph_async("social", &[]).await.unwrap(); + assert_eq!(job.id(), None); + job.wait().await.unwrap(); + + for (status, body, expected) in [ + (202, "{}", "valid job_id"), + (204, "", "200 OK or 202 Accepted"), + ] { + let db = super::RemoteDatabase::new_mock(move |_| { + http::Response::builder().status(status).body(body).unwrap() + }); + let error = db + .drop_property_graph_async("social", &[]) + .await + .err() + .unwrap(); + assert!(error.to_string().contains(expected), "{error}"); + } + } + + #[tokio::test] + async fn test_list_property_graphs_walks_pages() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!( + request.url().path(), + "/v1/namespace/analytics/property_graph/list" + ); + let page = request + .url() + .query_pairs() + .find(|(key, _)| key == "page_token") + .map(|(_, value)| value.into_owned()); + let body = match page.as_deref() { + None => r#"{"property_graphs":["social"],"page_token":"p2"}"#, + Some("p2") => r#"{"property_graphs":["payments"],"page_token":null}"#, + Some(other) => panic!("unexpected page token: {other}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + assert_eq!( + conn.list_property_graphs(&["analytics".into()]) + .await + .unwrap(), + vec!["social".to_string(), "payments".to_string()] + ); + + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(200) + .body(r#"{"property_graphs":["social"],"page_token":"same"}"#) + .unwrap() + }); + let error = conn.list_property_graphs(&[]).await.unwrap_err(); + assert!( + error.to_string().contains("repeated a page_token"), + "{error}" + ); + } + + /// A name carrying the delimiter would split back apart as a different + /// graph, so it is refused before it reaches a route. + #[tokio::test] + async fn test_property_graph_names_that_would_resplit_are_refused() { + let conn = Connection::new_with_handler(|_| -> http::Response { + panic!("an invalid identifier must not reach the service") + }); + for (name, namespace) in [ + ("analytics$social", vec![]), + ("social", vec!["ana$lytics".to_string()]), + ("", vec![]), + ("..", vec![]), + ] { + let error = match conn.describe_property_graph(name, &namespace).await { + Ok(_) => panic!("accepted {name:?} in {namespace:?}"), + Err(error) => error.to_string(), + }; + assert!(error.contains("property graph"), "{error}"); + } + } + #[tokio::test] async fn test_create_function_async_sends_canonical_request_and_decodes_typed_job() { const REQUEST: &str = include_str!( diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 52b7bd5e8..368edfb3e 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -195,6 +195,15 @@ pub fn validate_view_component(what: &str, value: &str) -> Result<()> { }) } +/// Validate a property graph's name and the namespace path holding it, under +/// the object-name rules views follow. +pub fn validate_property_graph_reference(name: &str, namespace_path: &[String]) -> Result<()> { + for segment in namespace_path { + validate_view_component("property graph namespace path segment", segment)?; + } + validate_view_component("property graph name", name) +} + /// Validate all components of a namespace /// /// Iterates through all namespace components and validates each one.