diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 5d6ee3d..f53246e 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -22,7 +22,7 @@ use arrow::ffi_stream::ArrowArrayStreamReader; use arrow_array::{RecordBatch, RecordBatchReader}; use arrow_schema::Schema; use datafusion::datasource::{DefaultTableSource, MemTable, TableProvider}; -use datafusion::execution::context::SessionContext; +use datafusion::execution::context::{SessionConfig, SessionContext}; use datafusion::execution::session_state::SessionStateBuilder; use futures::StreamExt; use lance::datafusion::LanceTableProvider; @@ -40,9 +40,9 @@ use lance_graph::{ SqlDialect as RustSqlDialect, SqlQuery as RustSqlQuery, VectorSearch as RustVectorSearch, }; use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::{PyDict, PyInt, PyList}, + types::{PyBool, PyDict, PyInt, PyList}, IntoPyObject, }; use serde_json::Value as JsonValue; @@ -1537,12 +1537,18 @@ fn build_engine_state( config: &RustGraphConfig, sources: HashMap, use_csr: bool, + target_partitions: Option, previous: Option<&EngineState>, ) -> PyResult { + let mut session_config = SessionConfig::new(); + if let Some(partitions) = target_partitions { + session_config = session_config.with_target_partitions(partitions); + } // The CSR planner only changes plans that contain CSR nodes, so install it always. let ctx = SessionContext::new_with_state( SessionStateBuilder::new() .with_default_features() + .with_config(session_config) .with_query_planner(Arc::new(CsrQueryPlanner)) // Move target-only filters into CSR nodes first, so pruning can then drop // the filter columns from the take. @@ -2109,6 +2115,7 @@ pub struct CypherEngine { config: RustGraphConfig, state: RwLock>, use_csr: bool, + target_partitions: Option, } impl CypherEngine { @@ -2213,23 +2220,48 @@ impl CypherEngine { /// datasets : dict /// Dictionary mapping table names to Lance datasets or PyArrow tables. /// Lance datasets are not read here; PyArrow inputs are held in memory. + /// target_partitions : int or None, optional + /// Positive target for DataFusion query partitioning, preserved by refresh(). + /// None keeps DataFusion's default. This is not a thread or memory limit. /// /// Raises /// ------ /// ValueError /// If no datasets are given, a configured id/endpoint column is missing, or a /// Lance dataset cannot be reopened by URI (for example, namespace-vended - /// credentials) + /// credentials), or target_partitions is zero, negative or out of range + /// TypeError + /// If target_partitions is a boolean or a non-integer value other than None /// RuntimeError /// If a dataset cannot be opened or registered #[new] - #[pyo3(signature = (config, datasets, use_csr = false))] + #[pyo3(signature = (config, datasets, use_csr = false, *, target_partitions = None))] fn new( py: Python, config: &GraphConfig, datasets: &Bound<'_, PyDict>, use_csr: bool, + target_partitions: Option<&Bound<'_, PyAny>>, ) -> PyResult { + let target_partitions = target_partitions + .map(|value| { + if value.is_instance_of::() || !value.is_instance_of::() { + return Err(PyTypeError::new_err( + "target_partitions must be an integer or None, not a boolean", + )); + } + value + .extract::() + .ok() + .filter(|value| *value > 0) + .ok_or_else(|| { + PyValueError::new_err(format!( + "target_partitions must be between 1 and {}", + usize::MAX + )) + }) + }) + .transpose()?; if datasets.is_empty() { return Err(PyValueError::new_err("No input datasets provided")); } @@ -2239,11 +2271,13 @@ impl CypherEngine { let source = python_value_to_source(py, &name, &value)?; sources.insert(name, source); } - let state = build_engine_state(py, &config.inner, sources, use_csr, None)?; + let state = + build_engine_state(py, &config.inner, sources, use_csr, target_partitions, None)?; Ok(Self { config: config.inner.clone(), state: RwLock::new(Arc::new(state)), use_csr, + target_partitions, }) } @@ -2372,8 +2406,14 @@ impl CypherEngine { } } } - let state = - build_engine_state(py, &self.config, sources, self.use_csr, Some(¤t))?; + let state = build_engine_state( + py, + &self.config, + sources, + self.use_csr, + self.target_partitions, + Some(¤t), + )?; if self.try_publish_refresh(¤t, state) { return Ok(()); } @@ -2751,6 +2791,48 @@ mod tests { } } + #[test] + fn target_partitions_survive_refresh_without_changing_other_defaults() { + Python::initialize(); + Python::attach(|py| { + let defaults = SessionConfig::new(); + let config = RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for target_partitions in [None, Some(1), Some(2)] { + let state = build_engine_state( + py, + &config, + test_state(1).sources, + false, + target_partitions, + None, + ) + .unwrap(); + let engine = CypherEngine { + config: config.clone(), + state: RwLock::new(Arc::new(state)), + use_csr: false, + target_partitions, + }; + let before = engine.snapshot(); + engine.refresh(py).unwrap(); + for snapshot in [before, engine.snapshot()] { + let state = snapshot.context.state(); + assert_eq!( + state.config().target_partitions(), + target_partitions.unwrap_or(defaults.target_partitions()) + ); + assert_eq!( + state.config().round_robin_repartition(), + defaults.round_robin_repartition() + ); + } + } + }); + } + #[test] fn test_plan_cache_reuses_bounds_and_refreshes() { let engine = CypherEngine { @@ -2760,6 +2842,7 @@ mod tests { .unwrap(), state: RwLock::new(Arc::new(test_state(1))), use_csr: false, + target_partitions: None, }; let snapshot = engine.snapshot(); let text = "MATCH (p:Person) WHERE p.id = $value RETURN p.id"; @@ -2835,6 +2918,7 @@ mod tests { .unwrap(), state: RwLock::new(Arc::new(test_state(1))), use_csr: false, + target_partitions: None, }; let reader_snapshot = engine.snapshot(); let (started_tx, started_rx) = mpsc::channel(); diff --git a/docs/api-reference.md b/docs/api-reference.md index dd2e790..968ffb1 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -123,7 +123,7 @@ See [Graph schema](graph-schema.md). ```python -CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) +CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False, *, target_partitions: int | None = None) ``` Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index @@ -138,10 +138,12 @@ See [Repeated queries](lance.md#repeated-queries) for the scope and limits. | `config` | The graph mapping. | | `datasets` | See [Datasets](#datasets). Lance datasets are pinned to the version passed in and not read here. | | `use_csr` | Build the experimental [CSR index](csr.md) for relationship types with integer endpoints, and node indexes for labels with unique integer keys. Default `False`. | +| `target_partitions` | Optional positive integer target for DataFusion query partitioning. `None` keeps DataFusion's default. Keyword-only; preserved by `refresh()`. It is not a thread or memory limit. See [Query parallelism](lance.md#query-parallelism). | Raises `ValueError` if `datasets` is empty, a configured key or endpoint column is missing, or a Lance dataset was opened through a namespace client and can't be reopened by URI. Raises `RuntimeError` if a dataset can't -be opened. +be opened. For `target_partitions`, booleans and non-integers other than `None` raise `TypeError`; +zero, negative values and integers outside the platform's unsigned integer range raise `ValueError`. `#!python execute(query: str, params: dict | None = None, *, timeout: float | None = None) -> pyarrow.Table` : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, @@ -160,7 +162,8 @@ be opened. `#!python refresh() -> None` : Moves every Lance dataset to its latest version and updates the CSR, node indexes and statistics. All-or-nothing; running queries keep their snapshot. If another refresh publishes first, retries from - that snapshot so an older candidate cannot replace it. See [Lance datasets](lance.md#refresh). + that snapshot so an older candidate cannot replace it. Preserves `target_partitions`. + See [Lance datasets](lance.md#refresh). `#!python dataset_versions() -> dict[str, int]` : The pinned version of each Lance dataset. PyArrow inputs are not listed. diff --git a/docs/cypher.md b/docs/cypher.md index df005b4..210f248 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -40,12 +40,13 @@ engine = CypherEngine(config, datasets, use_csr=True) ## `CypherEngine` ```python -CypherEngine(config, datasets, use_csr=False) +CypherEngine(config, datasets, use_csr=False, target_partitions=None) ``` The constructor validates the datasets against the config, registers them with a DataFusion session, and, with `use_csr=True`, builds the [CSR index](csr.md). After that, every query reuses the same catalog and -session. +session. The optional keyword-only `target_partitions` sets a positive target for query partitioning; +`None` keeps the current default. See [Query parallelism](lance.md#query-parallelism) before tuning it. !!! tip "Share one engine" diff --git a/docs/lance.md b/docs/lance.md index 4e09671..5a6a83b 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -95,6 +95,28 @@ change the plan. Queries whose operators cannot produce a reusable template fall new snapshot, including a refresh whose datasets did not change; active readers retain their existing plans. These rules apply with and without CSR. They do not promise a speedup for a particular workload. +## Query parallelism + +`CypherEngine(..., target_partitions=None)` keeps DataFusion's default query partitioning. If you have +measured your workload, you can choose a positive target explicitly: + +```python +serial_engine = CypherEngine( + config, + {"Person": lance.dataset("person.lance"), "KNOWS": lance.dataset("knows.lance")}, + use_csr=True, + target_partitions=1, +) +print(serial_engine.execute("MATCH (p:Person) RETURN count(*) AS n").to_pylist()) # [{'n': 4}] +serial_engine.refresh() # retains target_partitions=1 +``` + +The target is a planning preference, not a thread or memory limit; operators may use a different number +of partitions. A lower target can help short queries while making larger joins and scans slower. Compare +your full workload in the storage mode you use before choosing a value. The default remains unchanged, +and `refresh()` preserves your choice. Booleans and non-integers other than `None` raise `TypeError`; +zero, negative and out-of-range integers raise `ValueError`. + ## Streaming large results `execute()` returns one table, held in memory. For big outputs, stream: diff --git a/examples/lance_snapshot_refresh.py b/examples/lance_snapshot_refresh.py index 1588a0c..efae887 100644 --- a/examples/lance_snapshot_refresh.py +++ b/examples/lance_snapshot_refresh.py @@ -46,6 +46,8 @@ def main() -> None: "Person": lance.dataset(person_uri), "KNOWS": lance.dataset(root / "knows.lance"), }, + # Optional: choose this target after measuring your own queries; refresh keeps it. + target_partitions=2, ) count = "MATCH (p:Person) RETURN count(*) AS n" diff --git a/python/DEVELOPMENT.md b/python/DEVELOPMENT.md index 7f57095..d5e73a7 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -58,14 +58,16 @@ test executable needs an explicit link to the active Python library. From `pytho ```shell lance_python_library=$(uv run --no-sync python -c 'import pathlib, sysconfig; print(pathlib.Path(sysconfig.get_config_var("LIBDIR")) / sysconfig.get_config_var("LDLIBRARY"))') +lance_python_home=$(uv run --no-sync python -c 'import sys; print(sys.base_prefix)') lance_test_artifacts=$(mktemp) cargo rustc --manifest-path ../crates/lance-graph-python/Cargo.toml --tests --message-format=json -- -C "link-arg=$lance_python_library" > "$lance_test_artifacts" lance_test_binary=$(uv run --no-sync python -c 'import json, sys; artifacts = [json.loads(line) for line in open(sys.argv[1])]; print(next(a["executable"] for a in artifacts if a.get("reason") == "compiler-artifact" and a.get("profile", {}).get("test") and a.get("executable")))' "$lance_test_artifacts") -"$lance_test_binary" +PYTHONHOME="$lance_python_home" "$lance_test_binary" rm "$lance_test_artifacts" ``` -This uses Cargo's emitted executable path rather than selecting an older binary from `target/`. It does not +`PYTHONHOME` lets tests that initialize the interpreter find its standard library, including with a +relocatable Python installation. This uses Cargo's emitted executable path rather than selecting an older binary from `target/`. It does not change extension packaging or replace the Python tests. Use a separate `CARGO_TARGET_DIR` for each worktree when validating multiple revisions, so each test binary is built from that worktree's sources. diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi index e6544f7..0c921a6 100644 --- a/python/python/lance_csr/_internal/graph.pyi +++ b/python/python/lance_csr/_internal/graph.pyi @@ -115,7 +115,12 @@ class CypherQuery: @final class CypherEngine: def __init__( - self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False + self, + config: GraphConfig, + datasets: dict[str, _DatasetT], + use_csr: bool = False, + *, + target_partitions: int | None = None, ) -> None: ... def execute( self, diff --git a/python/python/tests/test_target_partitions.py b/python/python/tests/test_target_partitions.py new file mode 100644 index 0000000..bd608af --- /dev/null +++ b/python/python/tests/test_target_partitions.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""The optional partition target preserves query results and snapshot behavior.""" + +import inspect + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture +def graph(): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return config, { + "Person": pa.table({"id": [1, 2, 3, 4]}), + "KNOWS": pa.table({"src": [1, 1, 2], "dst": [2, 3, 4]}), + } + + +@pytest.mark.parametrize("partitions", [None, 1, 2]) +@pytest.mark.parametrize("use_csr", [False, True]) +@pytest.mark.parametrize("storage", ["arrow", "lance"]) +def test_partition_targets_preserve_queries_and_refresh( + graph, partitions, use_csr, storage, tmp_path +): + config, data = graph + if storage == "lance": + data = { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in data.items() + } + engine = CypherEngine(config, data, use_csr, target_partitions=partitions) + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) WHERE a.id = $start " + "RETURN b.id AS id ORDER BY id" + ) + for start, ids in [(1, [2, 3]), (2, [4]), (1, [2, 3])]: + expected = [{"id": value} for value in ids] + assert engine.execute(query, {"start": start}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"start": start}).read_all().to_pylist() + == expected + ) + assert engine.execute( + "MATCH (p:Person) RETURN count(*) AS n, sum(p.id) AS total" + ).to_pylist() == [{"n": 4, "total": 10}] + scan = "MATCH (p:Person) RETURN p.id AS id ORDER BY id" + reader = engine.execute_batches(scan) + if storage == "lance": + lance.write_dataset( + pa.table({"id": [5]}), tmp_path / "Person.lance", mode="append" + ) + engine.refresh() + assert reader.read_all().to_pylist() == [{"id": value} for value in range(1, 5)] + count = 5 if storage == "lance" else 4 + assert engine.execute(scan).to_pylist() == [ + {"id": value} for value in range(1, count + 1) + ] + + +@pytest.mark.parametrize("value", [True, False, 1.0, "1", [], {}]) +def test_partition_target_rejects_non_integers(graph, value): + config, data = graph + with pytest.raises(TypeError, match="target_partitions"): + CypherEngine(config, data, target_partitions=value) + + +@pytest.mark.parametrize("value", [0, -1, -(1 << 128), 1 << 128]) +def test_partition_target_rejects_nonpositive_and_overflow(graph, value): + config, data = graph + with pytest.raises(ValueError, match="target_partitions"): + CypherEngine(config, data, target_partitions=value) + + +def test_partition_target_is_keyword_only_and_default_is_none(graph): + config, data = graph + parameter = inspect.signature(CypherEngine).parameters["target_partitions"] + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + assert parameter.default is None + with pytest.raises(TypeError): + CypherEngine(config, data, False, 1) + assert CypherEngine(config, data, False).execute("RETURN 1 AS n").to_pylist() == [ + {"n": 1} + ] diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py index c4e3704..27f2965 100644 --- a/python/python/tests/typing/public_api.py +++ b/python/python/tests/typing/public_api.py @@ -66,6 +66,8 @@ def public_api( CypherEngine(config, {"Person": batch}) CypherEngine(config, {"Person": reader}) CypherEngine(config, {"Person": dataset}) + CypherEngine(config, datasets, target_partitions=1) + CypherEngine(config, datasets, target_partitions=None) one = CypherQuery(query).with_config(config).with_parameter("id", 1) assert_type(one, CypherQuery) @@ -107,6 +109,8 @@ def public_api( assert_type(SqlEngine(datasets).execute(sql.sql()), Table) # Rejected calls are checked too: removing a type restriction makes these unused. + CypherEngine(config, datasets, target_partitions="1") # pyright: ignore[reportArgumentType] + CypherEngine(config, datasets, False, 1) # pyright: ignore[reportCallIssue] engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] vector.top_k("5") # pyright: ignore[reportArgumentType] one.with_parameter("ids", {1, 2}) # pyright: ignore[reportArgumentType]