feat(namespace): add directory-backed execution (lance-format#96)

Author: beinan

python/Cargo.lock

@@ -77,6 +77,7 @@ markers = [
     "slow",
     "torch: tests which rely on pytorch being installed",
     "recurring: marks tests as recurring tests",
+    "requires_lance: tests that need the Lance Python bindings available",
 ]
 filterwarnings = [
     'error::FutureWarning',

python/pyproject.toml

@@ -73,13 +73,15 @@ def _load_dev_build() -> ModuleType:
 CypherQuery = _bindings.graph.CypherQuery
 VectorSearch = _bindings.graph.VectorSearch
 DistanceMetric = _bindings.graph.DistanceMetric
+DirNamespace = _bindings.graph.DirNamespace
 
 __all__ = [
     "GraphConfig",
     "GraphConfigBuilder",
     "CypherQuery",
     "VectorSearch",
     "DistanceMetric",
+    "DirNamespace",
 ]
 
 __version__ = _bindings.__version__

python/python/lance_graph/__init__.py

@@ -3,7 +3,7 @@
 
 import pyarrow as pa
 import pytest
-from lance_graph import CypherQuery, GraphConfig
+from lance_graph import CypherQuery, DirNamespace, GraphConfig
 
 
 @pytest.fixture
@@ -194,3 +194,24 @@ def test_distinct_clause(graph_env):
 
     assert len(data["c.company_name"]) == 3
     assert set(data["c.company_name"]) == {"TechCorp", "DataInc", "CloudSoft"}
+
+
+@pytest.mark.requires_lance
+def test_execute_with_directory_namespace(graph_env, tmp_path):
+    config, datasets, _ = graph_env
+
+    from lance import write_dataset
+
+    for name, table in datasets.items():
+        write_dataset(table, tmp_path / f"{name}.lance")
+
+    namespace = DirNamespace(str(tmp_path))
+
+    query = CypherQuery("MATCH (p:Person) WHERE p.age > 30 RETURN p.name").with_config(
+        config
+    )
+
+    result = query.execute_with_namespace(namespace)
+    data = result.to_pydict()
+
+    assert set(data["p.name"]) == {"Bob", "David"}

python/python/tests/test_graph.py

@@ -22,9 +22,9 @@ use arrow::ffi_stream::ArrowArrayStreamReader;
 use arrow_array::{RecordBatch, RecordBatchReader};
 use arrow_schema::Schema;
 use lance_graph::{
-    ExecutionStrategy as RustExecutionStrategy, CypherQuery as RustCypherQuery,
-    GraphConfig as RustGraphConfig, GraphError as RustGraphError,
-    VectorSearch as RustVectorSearch, ast::DistanceMetric as RustDistanceMetric,
+    ast::DistanceMetric as RustDistanceMetric, CypherQuery as RustCypherQuery,
+    ExecutionStrategy as RustExecutionStrategy, GraphConfig as RustGraphConfig,
+    GraphError as RustGraphError, VectorSearch as RustVectorSearch,
 };
 use pyo3::{
     exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
@@ -34,6 +34,7 @@ use pyo3::{
 };
 use serde_json::Value as JsonValue;
 
+use crate::namespace::PyDirNamespace;
 use crate::RT;
 
 /// Execution strategy for Cypher queries
@@ -537,6 +538,45 @@ impl CypherQuery {
         record_batch_to_python_table(py, &result_batch)
     }
 
+    /// Execute query using a namespace resolver.
+    ///
+    /// Parameters
+    /// ----------
+    /// namespace : DirNamespace
+    ///     Directory-backed namespace that resolves table names to Lance datasets.
+    /// strategy : ExecutionStrategy, optional
+    ///     Execution strategy to use (defaults to DataFusion)
+    ///
+    /// Returns
+    /// -------
+    /// pyarrow.Table
+    ///     Query results as Arrow table
+    ///
+    /// Raises
+    /// ------
+    /// RuntimeError
+    ///     If query execution fails
+    #[pyo3(signature = (namespace, strategy=None))]
+    fn execute_with_namespace(
+        &self,
+        py: Python,
+        namespace: &Bound<'_, PyDirNamespace>,
+        strategy: Option<ExecutionStrategy>,
+    ) -> PyResult<PyObject> {
+        let rust_strategy = strategy.map(|s| s.into());
+        let inner_query = self.inner.clone();
+        let namespace_arc = namespace.borrow().inner.clone();
+
+        let result_batch = RT
+            .block_on(
+                Some(py),
+                inner_query.execute_with_namespace_arc(namespace_arc, rust_strategy),
+            )?
+            .map_err(graph_error_to_pyerr)?;
+
+        record_batch_to_python_table(py, &result_batch)
+    }
+
     /// Explain query using the DataFusion planner with in-memory datasets
     ///
     /// Parameters
@@ -641,7 +681,10 @@ impl CypherQuery {
 
         // Execute via runtime
         let result = RT
-            .block_on(Some(py), inner_query.execute_with_vector_rerank(arrow_datasets, vs))?
+            .block_on(
+                Some(py),
+                inner_query.execute_with_vector_rerank(arrow_datasets, vs),
+            )?
             .map_err(graph_error_to_pyerr)?;
 
         record_batch_to_python_table(py, &result)
@@ -840,6 +883,7 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) ->
     graph_module.add_class::<GraphConfigBuilder>()?;
     graph_module.add_class::<CypherQuery>()?;
     graph_module.add_class::<VectorSearch>()?;
+    graph_module.add_class::<PyDirNamespace>()?;
 
     parent_module.add_submodule(&graph_module)?;
     Ok(())

python/src/graph.rs

@@ -4,6 +4,7 @@ use pyo3::prelude::*;
 
 mod executor;
 mod graph;
+mod namespace;
 
 pub(crate) static RT: LazyLock<executor::BackgroundExecutor> =
     LazyLock::new(executor::BackgroundExecutor::new);

python/src/lib.rs

@@ -0,0 +1,24 @@
+use std::sync::Arc;
+
+use lance_graph::namespace::DirNamespace;
+use pyo3::prelude::*;
+
+#[pyclass(name = "DirNamespace", module = "lance.graph")]
+pub struct PyDirNamespace {
+    pub(crate) inner: Arc<DirNamespace>,
+}
+
+#[pymethods]
+impl PyDirNamespace {
+    #[new]
+    fn new(base_uri: String) -> Self {
+        Self {
+            inner: Arc::new(DirNamespace::new(base_uri)),
+        }
+    }
+
+    #[getter]
+    fn base_uri(&self) -> String {
+        self.inner.base_uri().to_string()
+    }
+}