feat(python): allow an explicit query partition target

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<String, EngineSource>,
     use_csr: bool,
+    target_partitions: Option<usize>,
     previous: Option<&EngineState>,
 ) -> PyResult<EngineState> {
+    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<Arc<EngineState>>,
     use_csr: bool,
+    target_partitions: Option<usize>,
 }
 
 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<Self> {
+        let target_partitions = target_partitions
+            .map(|value| {
+                if value.is_instance_of::<PyBool>() || !value.is_instance_of::<PyInt>() {
+                    return Err(PyTypeError::new_err(
+                        "target_partitions must be an integer or None, not a boolean",
+                    ));
+                }
+                value
+                    .extract::<usize>()
+                    .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(&current))?;
+            let state = build_engine_state(
+                py,
+                &self.config,
+                sources,
+                self.use_csr,
+                self.target_partitions,
+                Some(&current),
+            )?;
             if self.try_publish_refresh(&current, 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();

docs/api-reference.md

@@ -123,7 +123,7 @@ See [Graph schema](graph-schema.md).
 
 <!-- skip-check -->
 ```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.

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"
 

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:

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"
 

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.
 

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,