feat: manage property graphs from the client

Adds create, describe, list and drop for property graphs -- catalog objects
that read tables as the nodes and edges of a graph -- to the Rust Connection,
the remote database, and the Python sync and async connections. Local
connections refuse them, as they do views.

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",

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: ...

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()

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"]),
+    )

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."""

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")

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<Vec<String>>,
+    ) -> PyResult<Bound<'_, PyAny>> {
+        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<Vec<String>>,
+    ) -> PyResult<Bound<'_, PyAny>> {
+        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<Vec<String>>,
+    ) -> PyResult<Bound<'_, PyAny>> {
+        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<Vec<String>>,
+    ) -> PyResult<Bound<'_, PyAny>> {
+        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<Bound<'_, PyAny>> {
         let inner = self_.get_inner()?.clone();
         future_into_py(self_.py(), async move {