Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions python/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions python/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
2 changes: 2 additions & 0 deletions python/python/lance_graph/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
23 changes: 22 additions & 1 deletion python/python/tests/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"}
52 changes: 48 additions & 4 deletions python/src/graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -34,6 +34,7 @@ use pyo3::{
};
use serde_json::Value as JsonValue;

use crate::namespace::PyDirNamespace;
use crate::RT;

/// Execution strategy for Cypher queries
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(())
Expand Down
1 change: 1 addition & 0 deletions python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
24 changes: 24 additions & 0 deletions python/src/namespace.rs
Original file line number Diff line number Diff line change
@@ -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()
}
}
62 changes: 32 additions & 30 deletions rust/lance-graph/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading