diff --git a/python/Cargo.lock b/python/Cargo.lock index f44716cfe..ccaa2badf 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -3435,6 +3435,7 @@ dependencies = [ "arrow", "arrow-array", "arrow-schema", + "async-trait", "datafusion", "datafusion-common", "datafusion-expr", @@ -3443,6 +3444,7 @@ dependencies = [ "futures", "lance", "lance-linalg", + "lance-namespace", "nom 7.1.3", "serde", "serde_json", diff --git a/python/pyproject.toml b/python/pyproject.toml index b7e5c03be..a38cb9265 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -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', diff --git a/python/python/lance_graph/__init__.py b/python/python/lance_graph/__init__.py index ecd166fe1..3e2d0b654 100644 --- a/python/python/lance_graph/__init__.py +++ b/python/python/lance_graph/__init__.py @@ -73,6 +73,7 @@ def _load_dev_build() -> ModuleType: CypherQuery = _bindings.graph.CypherQuery VectorSearch = _bindings.graph.VectorSearch DistanceMetric = _bindings.graph.DistanceMetric +DirNamespace = _bindings.graph.DirNamespace __all__ = [ "GraphConfig", @@ -80,6 +81,7 @@ def _load_dev_build() -> ModuleType: "CypherQuery", "VectorSearch", "DistanceMetric", + "DirNamespace", ] __version__ = _bindings.__version__ diff --git a/python/python/tests/test_graph.py b/python/python/tests/test_graph.py index c699d3384..4fca84fab 100644 --- a/python/python/tests/test_graph.py +++ b/python/python/tests/test_graph.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"} diff --git a/python/src/graph.rs b/python/src/graph.rs index 5fb76cb44..c1051d132 100644 --- a/python/src/graph.rs +++ b/python/src/graph.rs @@ -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, + ) -> PyResult { + 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::()?; graph_module.add_class::()?; graph_module.add_class::()?; + graph_module.add_class::()?; parent_module.add_submodule(&graph_module)?; Ok(()) diff --git a/python/src/lib.rs b/python/src/lib.rs index 32973a540..d9f440fbe 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -4,6 +4,7 @@ use pyo3::prelude::*; mod executor; mod graph; +mod namespace; pub(crate) static RT: LazyLock = LazyLock::new(executor::BackgroundExecutor::new); diff --git a/python/src/namespace.rs b/python/src/namespace.rs new file mode 100644 index 000000000..491a6f071 --- /dev/null +++ b/python/src/namespace.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, +} + +#[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() + } +} diff --git a/rust/lance-graph/Cargo.lock b/rust/lance-graph/Cargo.lock index dc0cff0d1..3b879717a 100644 --- a/rust/lance-graph/Cargo.lock +++ b/rust/lance-graph/Cargo.lock @@ -2360,9 +2360,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33c3b76530d0755759d6537bdb19b849b2b108b9013a379be1be6036d2cd210c" +checksum = "5ffdff7a2d68d22afc0657eddde3e946371ce7cfe730a3f78a5ed44ea5b1cb2e" dependencies = [ "arrow-array", "rand 0.9.2", @@ -3295,9 +3295,9 @@ dependencies = [ [[package]] name = "lance" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "993c0199f4c8291dec6e545a523d2ac97875f455bf6f27b8e60504a68c7654d3" +checksum = "e8c439decbc304e180748e34bb6d3df729069a222e83e74e2185c38f107136e9" dependencies = [ "arrow", "arrow-arith", @@ -3361,9 +3361,9 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "590272235088a62bf7fb7d3b13820af76c0d9092a6cffab8989e93c628fc8451" +checksum = "f4ee5508b225456d3d56998eaeef0d8fbce5ea93856df47b12a94d2e74153210" dependencies = [ "arrow-array", "arrow-buffer", @@ -3381,9 +3381,9 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75a1e0c990610227ec0d048f4c2cda80d7474de5c63498694a90408d62dcac32" +checksum = "d1c065fb3bd4a8cc4f78428443e990d4921aa08f707b676753db740e0b402a21" dependencies = [ "arrayref", "paste", @@ -3392,9 +3392,9 @@ dependencies = [ [[package]] name = "lance-core" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea4a8f6f2c25f6a74f460a58caecce62a288abaad9172c199e1ec3c004e5d74e" +checksum = "e8856abad92e624b75cd57a04703f6441948a239463bdf973f2ac1924b0bcdbe" dependencies = [ "arrow-array", "arrow-buffer", @@ -3430,9 +3430,9 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6cc44f88d68610cc7df1308cd852ba1ae5c341ac6132a9fce1873843957c2985" +checksum = "4c8835308044cef5467d7751be87fcbefc2db01c22370726a8704bd62991693f" dependencies = [ "arrow", "arrow-array", @@ -3462,9 +3462,9 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fef9a3d6436f234e98a5e8cbb292fdce0d78febea152d738117d3d70f86cbb58" +checksum = "612de1e888bb36f6bf51196a6eb9574587fdf256b1759a4c50e643e00d5f96d0" dependencies = [ "arrow", "arrow-array", @@ -3481,9 +3481,9 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31f32004e88abb4819d6d6386cf26e3dfc2258d4e3d0b2748364fb279b6787a6" +checksum = "2b456b29b135d3c7192602e516ccade38b5483986e121895fa43cf1fdb38bf60" dependencies = [ "arrow-arith", "arrow-array", @@ -3520,9 +3520,9 @@ dependencies = [ [[package]] name = "lance-file" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d261cb07db8aac64b8926e61cef73caf6d0aa3bd49e9fc8401eff1768fb9aa0" +checksum = "ab1538d14d5bb3735b4222b3f5aff83cfa59cc6ef7cdd3dd9139e4c77193c80b" dependencies = [ "arrow-arith", "arrow-array", @@ -3554,9 +3554,9 @@ dependencies = [ [[package]] name = "lance-geo" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5144f9f1ce0faaac0c791df55d5c9a0945a2a48489f108e12915f8e0b9e7ce2f" +checksum = "a5a69a2f3b55703d9c240ad7c5ffa2c755db69e9cf8aa05efe274a212910472d" dependencies = [ "datafusion", "geo-types", @@ -3572,6 +3572,7 @@ dependencies = [ "arrow", "arrow-array", "arrow-schema", + "async-trait", "criterion", "datafusion", "datafusion-common", @@ -3583,6 +3584,7 @@ dependencies = [ "lance-arrow", "lance-index", "lance-linalg", + "lance-namespace", "nom 7.1.3", "serde", "serde_json", @@ -3593,9 +3595,9 @@ dependencies = [ [[package]] name = "lance-index" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30dfdff485125e321f411f17a12a7878ef51234554237a2f416a1f8988a0a403" +checksum = "0ea84613df6fa6b9168a1f056ba4f9cb73b90a1b452814c6fd4b3529bcdbfc78" dependencies = [ "arrow", "arrow-arith", @@ -3656,9 +3658,9 @@ dependencies = [ [[package]] name = "lance-io" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b44d19e7ee99c3cdb417785ded67bac1fd8a9c9cbd9d1751fba0348294e6f3f" +checksum = "6b3fc4c1d941fceef40a0edbd664dbef108acfc5d559bb9e7f588d0c733cbc35" dependencies = [ "arrow", "arrow-arith", @@ -3698,9 +3700,9 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28e275333cfcd97cabfe7bbd77bad4df48a324499ebdf6c05febeef050efc685" +checksum = "b62ffbc5ce367fbf700a69de3fe0612ee1a11191a64a632888610b6bacfa0f63" dependencies = [ "arrow-array", "arrow-buffer", @@ -3716,9 +3718,9 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb58614e43573a083ca6425a3f22e20c67ce69d3259428b9cda25bc8b68ccc69" +checksum = "791bbcd868ee758123a34e07d320a1fb99379432b5ecc0e78d6b4686e999b629" dependencies = [ "arrow", "async-trait", @@ -3743,9 +3745,9 @@ dependencies = [ [[package]] name = "lance-table" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ef592598d62f12b11117a937859d180e4e46df85643988300e6a7bca44a714f" +checksum = "6fdb2d56bfa4d1511c765fa0cc00fdaa37e5d2d1cd2f57b3c6355d9072177052" dependencies = [ "arrow", "arrow-array", diff --git a/rust/lance-graph/Cargo.toml b/rust/lance-graph/Cargo.toml index 8ce0a6d3c..7e272ffdb 100644 --- a/rust/lance-graph/Cargo.toml +++ b/rust/lance-graph/Cargo.toml @@ -28,8 +28,10 @@ datafusion-expr = "50.3" datafusion-sql = "50.3" datafusion-functions-aggregate = "50.3" futures = "0.3" +async-trait = "0.1" lance = "1.0.0" lance-linalg = "1.0.0" +lance-namespace = "1.0.1" nom = "7.1" serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/rust/lance-graph/src/lib.rs b/rust/lance-graph/src/lib.rs index 864423100..1ff47ddfa 100644 --- a/rust/lance-graph/src/lib.rs +++ b/rust/lance-graph/src/lib.rs @@ -42,6 +42,7 @@ pub mod error; pub mod lance_native_planner; pub mod lance_vector_search; pub mod logical_plan; +pub mod namespace; pub mod parser; pub mod query; pub mod query_processor; @@ -55,4 +56,5 @@ pub const MAX_VARIABLE_LENGTH_HOPS: u32 = 20; pub use config::{GraphConfig, NodeMapping, RelationshipMapping}; pub use error::{GraphError, Result}; pub use lance_vector_search::VectorSearch; +pub use namespace::DirNamespace; pub use query::{CypherQuery, ExecutionStrategy}; diff --git a/rust/lance-graph/src/namespace/directory.rs b/rust/lance-graph/src/namespace/directory.rs new file mode 100644 index 000000000..541a9b2aa --- /dev/null +++ b/rust/lance-graph/src/namespace/directory.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 + +use async_trait::async_trait; +use lance_namespace::models::{DescribeTableRequest, DescribeTableResponse}; +use lance_namespace::{Error as NamespaceError, LanceNamespace, Result}; +use snafu::location; + +/// A namespace that resolves table names relative to a base directory or URI. +#[derive(Debug, Clone)] +pub struct DirNamespace { + base_uri: String, +} + +impl DirNamespace { + /// Create a new directory-backed namespace rooted at `base_uri`. + /// + /// The URI is normalized so that it does not end with a trailing slash. + pub fn new(base_uri: impl Into) -> Self { + let uri = base_uri.into(); + let clean_uri = uri.trim_end_matches('/').to_string(); + Self { + base_uri: clean_uri, + } + } + + /// Return the normalized base URI. + pub fn base_uri(&self) -> &str { + &self.base_uri + } +} + +#[async_trait] +impl LanceNamespace for DirNamespace { + fn namespace_id(&self) -> String { + format!("DirNamespace {{ base_uri: '{}' }}", self.base_uri) + } + + async fn describe_table(&self, request: DescribeTableRequest) -> Result { + let id = request.id.ok_or_else(|| { + NamespaceError::invalid_input( + "DirNamespace requires the table identifier to be provided", + location!(), + ) + })?; + + if id.len() != 1 { + return Err(NamespaceError::invalid_input( + format!( + "DirNamespace expects identifiers with a single component, got {:?}", + id + ), + location!(), + )); + } + + let table_name = &id[0]; + let location = format!("{}/{}.lance", self.base_uri, table_name); + + let mut response = DescribeTableResponse::new(); + response.location = Some(location); + response.storage_options = None; + Ok(response) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn describe_table_returns_clean_location() { + let namespace = DirNamespace::new("s3://bucket/path/"); + let mut request = DescribeTableRequest::new(); + request.id = Some(vec!["users".to_string()]); + + let response = namespace.describe_table(request).await.unwrap(); + assert_eq!( + response.location.as_deref(), + Some("s3://bucket/path/users.lance") + ); + } + + #[tokio::test] + async fn describe_table_rejects_missing_identifier() { + let namespace = DirNamespace::new("file:///tmp"); + let request = DescribeTableRequest::new(); + + let err = namespace.describe_table(request).await.unwrap_err(); + assert!( + err.to_string() + .contains("DirNamespace requires the table identifier"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn describe_table_rejects_multi_component_identifier() { + let namespace = DirNamespace::new("memory://namespace"); + let mut request = DescribeTableRequest::new(); + request.id = Some(vec!["foo".into(), "bar".into()]); + + let err = namespace.describe_table(request).await.unwrap_err(); + assert!( + err.to_string() + .contains("expects identifiers with a single component"), + "unexpected error: {err}" + ); + } +} diff --git a/rust/lance-graph/src/namespace/mod.rs b/rust/lance-graph/src/namespace/mod.rs new file mode 100644 index 000000000..fe6c89cd3 --- /dev/null +++ b/rust/lance-graph/src/namespace/mod.rs @@ -0,0 +1,3 @@ +pub mod directory; + +pub use directory::DirNamespace; diff --git a/rust/lance-graph/src/query.rs b/rust/lance-graph/src/query.rs index 8028eb0fd..52d5840ef 100644 --- a/rust/lance-graph/src/query.rs +++ b/rust/lance-graph/src/query.rs @@ -7,11 +7,13 @@ use crate::ast::CypherQuery as CypherAST; use crate::config::GraphConfig; use crate::error::{GraphError, Result}; use crate::logical_plan::LogicalPlanner; +use crate::namespace::DirNamespace; use crate::parser::parse_cypher_query; use crate::simple_executor::{ to_df_boolean_expr_simple, to_df_order_by_expr_simple, to_df_value_expr_simple, PathExecutor, }; -use std::collections::HashMap; +use lance_namespace::models::DescribeTableRequest; +use std::collections::{HashMap, HashSet}; /// Execution strategy for Cypher queries #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] @@ -151,6 +153,57 @@ impl CypherQuery { } } + /// Execute the query using a namespace-backed table resolver. + /// + /// The namespace is provided by value and will be shared internally as needed. + pub async fn execute_with_namespace( + &self, + namespace: DirNamespace, + strategy: Option, + ) -> Result { + self.execute_with_namespace_arc(std::sync::Arc::new(namespace), strategy) + .await + } + + /// Execute the query using a shared namespace instance. + pub async fn execute_with_namespace_arc( + &self, + namespace: std::sync::Arc, + strategy: Option, + ) -> Result { + let namespace_trait: std::sync::Arc = + namespace; + self.execute_with_namespace_internal(namespace_trait, strategy) + .await + } + + async fn execute_with_namespace_internal( + &self, + namespace: std::sync::Arc, + strategy: Option, + ) -> Result { + let strategy = strategy.unwrap_or_default(); + match strategy { + ExecutionStrategy::DataFusion => { + let (catalog, ctx) = self + .build_catalog_and_context_from_namespace(namespace) + .await?; + self.execute_with_catalog_and_context(std::sync::Arc::new(catalog), ctx) + .await + } + ExecutionStrategy::Simple => Err(GraphError::UnsupportedFeature { + feature: + "Simple execution strategy is not supported for namespace-backed execution" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + ExecutionStrategy::LanceNative => Err(GraphError::UnsupportedFeature { + feature: "Lance native execution strategy is not yet implemented".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }), + } + } + /// Explain the query execution plan using in-memory datasets /// /// Returns a formatted string showing the query execution plan at different stages: @@ -522,6 +575,119 @@ impl CypherQuery { Ok((catalog, ctx)) } + /// Helper to build catalog and context using a namespace resolver + async fn build_catalog_and_context_from_namespace( + &self, + namespace: std::sync::Arc, + ) -> Result<( + crate::source_catalog::InMemoryCatalog, + datafusion::execution::context::SessionContext, + )> { + use crate::source_catalog::InMemoryCatalog; + use datafusion::datasource::{DefaultTableSource, TableProvider}; + use datafusion::execution::context::SessionContext; + use lance::datafusion::LanceTableProvider; + use std::sync::Arc; + + let config = self.require_config()?; + + let mut required_tables: HashSet = HashSet::new(); + required_tables.extend(config.node_mappings.keys().cloned()); + required_tables.extend(config.relationship_mappings.keys().cloned()); + + if required_tables.is_empty() { + return Err(GraphError::ConfigError { + message: + "Graph configuration does not reference any node labels or relationship types" + .to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + + let ctx = SessionContext::new(); + let mut catalog = InMemoryCatalog::new(); + let mut providers: HashMap> = HashMap::new(); + + for table_name in required_tables { + let mut request = DescribeTableRequest::new(); + request.id = Some(vec![table_name.clone()]); + + let response = + namespace + .describe_table(request) + .await + .map_err(|e| GraphError::ConfigError { + message: format!( + "Namespace failed to resolve table '{}': {}", + table_name, e + ), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let location = response.location.ok_or_else(|| GraphError::ConfigError { + message: format!( + "Namespace did not provide a location for table '{}'", + table_name + ), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let dataset = lance::dataset::Dataset::open(&location) + .await + .map_err(|e| GraphError::ConfigError { + message: format!("Failed to open dataset for table '{}': {}", table_name, e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let dataset = Arc::new(dataset); + let provider: Arc = + Arc::new(LanceTableProvider::new(dataset.clone(), true, true)); + + ctx.register_table(&table_name, provider.clone()) + .map_err(|e| GraphError::PlanError { + message: format!( + "Failed to register table '{}' in SessionContext: {}", + table_name, e + ), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + providers.insert(table_name, provider); + } + + for label in config.node_mappings.keys() { + let provider = providers + .get(label) + .ok_or_else(|| GraphError::ConfigError { + message: format!( + "Namespace did not resolve dataset for node label '{}'", + label + ), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let table_source = Arc::new(DefaultTableSource::new(provider.clone())); + catalog = catalog.with_node_source(label, table_source); + } + + for rel_type in config.relationship_mappings.keys() { + let provider = providers + .get(rel_type) + .ok_or_else(|| GraphError::ConfigError { + message: format!( + "Namespace did not resolve dataset for relationship type '{}'", + rel_type + ), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let table_source = Arc::new(DefaultTableSource::new(provider.clone())); + catalog = catalog.with_relationship_source(rel_type, table_source); + } + + Ok((catalog, ctx)) + } + /// Internal helper to explain the query execution plan with explicit catalog and session context async fn explain_internal( &self, @@ -1960,4 +2126,109 @@ mod tests { // We check for "p.name" which is the expected output alias assert!(sql.contains("p.name")); } + + async fn write_lance_dataset(path: &std::path::Path, batch: arrow_array::RecordBatch) { + use arrow_array::{RecordBatch, RecordBatchIterator}; + use lance::dataset::{Dataset, WriteParams}; + + let schema = batch.schema(); + let batches: Vec> = + vec![std::result::Result::Ok(batch)]; + let reader = RecordBatchIterator::new(batches.into_iter(), schema); + + Dataset::write(reader, path.to_str().unwrap(), None::) + .await + .expect("write lance dataset"); + } + + fn build_people_batch() -> arrow_array::RecordBatch { + use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + let schema = Arc::new(Schema::new(vec![ + Field::new("person_id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("age", DataType::Int32, false), + ])); + + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, + Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "David"])) as ArrayRef, + Arc::new(Int32Array::from(vec![28, 34, 29, 42])) as ArrayRef, + ]; + + RecordBatch::try_new(schema, columns).expect("valid person batch") + } + + fn build_friendship_batch() -> arrow_array::RecordBatch { + use arrow_array::{ArrayRef, Int64Array, RecordBatch}; + use arrow_schema::{DataType, Field, Schema}; + use std::sync::Arc; + + let schema = Arc::new(Schema::new(vec![ + Field::new("person1_id", DataType::Int64, false), + Field::new("person2_id", DataType::Int64, false), + ])); + + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![1, 1, 2, 3])) as ArrayRef, + Arc::new(Int64Array::from(vec![2, 3, 4, 4])) as ArrayRef, + ]; + + RecordBatch::try_new(schema, columns).expect("valid friendship batch") + } + + #[tokio::test] + async fn executes_against_directory_namespace() { + use arrow_array::StringArray; + use tempfile::tempdir; + + let tmp_dir = tempdir().unwrap(); + write_lance_dataset(&tmp_dir.path().join("Person.lance"), build_people_batch()).await; + write_lance_dataset( + &tmp_dir.path().join("FRIEND_OF.lance"), + build_friendship_batch(), + ) + .await; + + let config = GraphConfig::builder() + .with_node_label("Person", "person_id") + .with_relationship("FRIEND_OF", "person1_id", "person2_id") + .build() + .expect("valid graph config"); + + let query = CypherQuery::new("MATCH (p:Person) WHERE p.age > 30 RETURN p.name") + .expect("query parses") + .with_config(config); + + let namespace = DirNamespace::new(tmp_dir.path().to_string_lossy().into_owned()); + + let result = query + .execute_with_namespace(namespace.clone(), None) + .await + .expect("namespace execution succeeds"); + + use arrow_array::Array; + let names = result + .column(0) + .as_any() + .downcast_ref::() + .expect("string column"); + + let mut values: Vec = (0..names.len()) + .map(|i| names.value(i).to_string()) + .collect(); + values.sort(); + assert_eq!(values, vec!["Bob".to_string(), "David".to_string()]); + + let err = query + .execute_with_namespace(namespace, Some(ExecutionStrategy::Simple)) + .await + .expect_err("simple strategy not supported"); + assert!( + matches!(err, GraphError::UnsupportedFeature { .. }), + "expected unsupported feature error, got {err:?}" + ); + } }