diff --git a/crates/lance-graph-benches/README.md b/crates/lance-graph-benches/README.md index ce63a57..47c729a 100644 --- a/crates/lance-graph-benches/README.md +++ b/crates/lance-graph-benches/README.md @@ -26,3 +26,16 @@ cargo bench -p lance-graph-benches --bench graph_execution This crate is not published to crates.io and is excluded from releases. It exists solely for performance testing during development. + +## Comprehension counts + +Compare `size([x IN list WHERE predicate])` with the equivalent identity-mapped comprehension, which keeps the +ordinary list-building path. Both queries run against 1,000 input rows with list lengths 16 and 256; results, +row count and summed counts must match before timing. This measures planning plus execution after parsing. + +```bash +cargo bench -p lance-graph-benches --bench graph_execution -- comprehension_count --warm-up-time 1 +``` + +The count path removes ordered list aggregation. It still numbers and unnests rows, evaluates predicates and +groups by input row, so these measurements do not establish a general bound on planning or memory use. diff --git a/crates/lance-graph-benches/benches/graph_execution.rs b/crates/lance-graph-benches/benches/graph_execution.rs index 2b6831b..253fac3 100644 --- a/crates/lance-graph-benches/benches/graph_execution.rs +++ b/crates/lance-graph-benches/benches/graph_execution.rs @@ -336,5 +336,44 @@ fn bench_cypher_execution(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_cypher_execution); +fn bench_comprehension_counts(c: &mut Criterion) { + use arrow_array::Int64Array; + let rt = tokio::runtime::Runtime::new().unwrap(); + let config = GraphConfig::builder() + .node_table("Person", "person_id") + .build() + .unwrap(); + let data = HashMap::from([("Person".to_string(), make_people_batch(1000))]); + let mut group = c.benchmark_group("comprehension_count"); + for length in [16, 256] { + let query = |map| { + CypherQuery::new(&format!( + "MATCH (p:Person) RETURN p.person_id AS id, size([x IN range(0, {}) WHERE x % 2 = 0 {map}]) AS n ORDER BY id", length - 1 + )).unwrap().with_config(config.clone()) + }; + let count = query(""); + let reference = query("| x"); + let actual = execute_cypher_query(&rt, &count, data.clone()); + let expected = execute_cypher_query(&rt, &reference, data.clone()); + assert_eq!(actual.columns(), expected.columns()); + assert_eq!(actual.num_rows(), 1000); + let counts = actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(counts.values().iter().sum::(), 1000 * length / 2); + for (name, query) in [ + ("filtered_count", &count), + ("materialized_reference", &reference), + ] { + group.bench_with_input(BenchmarkId::new(name, length), &length, |b, _| { + b.iter(|| black_box(execute_cypher_query(&rt, query, data.clone()))); + }); + } + } + group.finish(); +} + +criterion_group!(benches, bench_cypher_execution, bench_comprehension_counts); criterion_main!(benches); diff --git a/crates/lance-graph/README.md b/crates/lance-graph/README.md index e5172fb..f7fdad1 100644 --- a/crates/lance-graph/README.md +++ b/crates/lance-graph/README.md @@ -19,6 +19,10 @@ Replace a directly constructed `Some(n)` with `Some(ast::RowCount::Literal(n))`; use `ast::RowCount::Parameter(name)`. Resolve parameter values before building a logical plan, using the normal `CypherQuery` parameter-binding methods. Existing Cypher text with literal pagination is unchanged. +`LogicalOperator::ListComprehension` also has a `count_only` field. Direct constructors that produce +list values must set it to `false`; the planner sets it to `true` only for an unmapped comprehension +whose size is requested. Its count output is Int64 and non-nullable, including for empty or null lists. + ## Quick Start ```rust diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs index 43b77d2..93ae693 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -9,6 +9,7 @@ use arrow_schema::DataType; use datafusion::common::Column; use datafusion::functions::core::expr_fn::coalesce; use datafusion::functions_aggregate::array_agg::array_agg; +use datafusion::functions_aggregate::count::count; use datafusion::functions_nested::expr_fn::{array_element, cardinality, gen_series, make_array}; use datafusion::functions_window::expr_fn::row_number; use datafusion::logical_expr::{ @@ -42,6 +43,7 @@ impl DataFusionPlanner { keep: Option<&BooleanExpression>, map: &ValueExpression, output: &str, + count_only: bool, ) -> Result { let error = |what: &str| { let what = what.to_string(); @@ -132,12 +134,18 @@ impl DataFusionPlanner { if let Some(keep) = keep { kept = kept.and(to_df_boolean_expr(&bind_boolean(keep, elements.schema()))); } - let value = to_df_value_expr(&bind_value(map, elements.schema())); - let collected = array_agg(value) - .filter(kept) - .order_by(vec![col(&pos).sort(true, false)]) - .build() - .map_err(error("Failed to collect comprehension values"))? + let aggregate = if count_only { + // Count rows, not the mapped value: null elements still occupy a slot. + count(lit(1_i64)).filter(kept).build() + } else { + let value = to_df_value_expr(&bind_value(map, elements.schema())); + array_agg(value) + .filter(kept) + .order_by(vec![col(&pos).sort(true, false)]) + .build() + }; + let collected = aggregate + .map_err(error("Failed to aggregate comprehension values"))? .alias(output); // One group per input row: its number plus every input column, all equal within it. let names: HashSet = numbered @@ -160,7 +168,12 @@ impl DataFusionPlanner { .filter(|f| *f.name() != row) .map(|f| Expr::Column(Column::new_unqualified(f.name()))) .collect(); - out.push(coalesce(vec![col(output), make_array(vec![])]).alias(output)); + let value = if count_only { + col(output) + } else { + coalesce(vec![col(output), make_array(vec![])]) + }; + out.push(value.alias(output)); LogicalPlanBuilder::from(elements) .aggregate(group, vec![collected]) .and_then(|b| b.project(out)) @@ -168,3 +181,89 @@ impl DataFusionPlanner { .map_err(error("Failed to build list comprehension")) } } + +#[cfg(test)] +mod tests { + use arrow_array::{types::Int64Type, ArrayRef, Int64Array, ListArray, RecordBatch}; + use std::collections::HashMap; + use std::sync::Arc; + + fn data() -> HashMap { + HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([ + ( + "id", + Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, + ), + ( + "values", + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), None, Some(1), Some(2)]), + Some(vec![]), + None, + Some(vec![None]), + ])) as ArrayRef, + ), + ]) + .unwrap(), + )]) + } + + #[tokio::test] + async fn comprehension_counts_keep_nulls_duplicates_and_empty_rows() { + let config = crate::GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for (predicate, expected) in [ + ("", vec![4, 0, 0, 1]), + ("WHERE x > 1", vec![1, 0, 0, 0]), + ("WHERE x IS NULL", vec![1, 0, 0, 1]), + ("WHERE x IS NULL OR x > 0", vec![4, 0, 0, 1]), + ] { + // Duplicated input rows must remain distinct; an identity map is the + // unoptimized reference and also retains null elements in ARRAY_AGG. + let make_query = |map| { + crate::CypherQuery::new(&format!( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate RETURN p.id AS id, \ + size([x IN p.values {predicate} {map}]) AS n ORDER BY id" + )) + .unwrap() + .with_config(config.clone()) + }; + let query = make_query(""); + let actual = query.execute(data(), None).await.unwrap(); + let reference = make_query("| x").execute(data(), None).await.unwrap(); + assert_eq!(actual.columns(), reference.columns(), "{predicate}"); + assert_eq!( + actual.schema().field(1).data_type(), + reference.schema().field(1).data_type() + ); + assert_eq!(actual.schema().field(0).name(), "id"); + assert_eq!(actual.schema().field(1).name(), "n"); + assert!(!actual.schema().field(1).is_nullable()); + let counts = actual + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + counts.values().to_vec(), + expected + .into_iter() + .flat_map(|n| [n, n]) + .collect::>() + ); + let plan = query.explain(data()).await.unwrap().to_lowercase(); + assert!(plan.contains("count("), "{plan}"); + assert!(!plan.contains("array_agg("), "{plan}"); + let mapped = make_query("| x") + .explain(data()) + .await + .unwrap() + .to_lowercase(); + assert!(mapped.contains("array_agg("), "{mapped}"); + } + } +} diff --git a/crates/lance-graph/src/datafusion_planner/builder/mod.rs b/crates/lance-graph/src/datafusion_planner/builder/mod.rs index c2ca6a4..5f07235 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/mod.rs @@ -117,6 +117,7 @@ impl DataFusionPlanner { keep, map, output, + count_only, } => self.build_list_comprehension( ctx, input, @@ -126,6 +127,7 @@ impl DataFusionPlanner { keep.as_ref(), map, output, + *count_only, ), LogicalOperator::ComprehensionElement => Ok(ctx .comprehension_elements diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index a85a642..e8ccc63 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -172,6 +172,9 @@ pub enum LogicalOperator { keep: Option, map: ValueExpression, output: String, + /// Count retained elements without building an ordered list. + #[serde(default)] + count_only: bool, }, /// The element rows of the enclosing [`LogicalOperator::ListComprehension`]. @@ -754,6 +757,24 @@ impl<'a> LogicalPlanner<'a> { None => Ok(input), } } + // The unmapped comprehension retains every passing element, including + // nulls and duplicates. Its size needs a filtered count, not ARRAY_AGG. + // Keep mapped comprehensions unchanged so their evaluation/errors remain. + ValueExpression::ScalarFunction { name, args } + if name.eq_ignore_ascii_case("size") + && matches!( + args.as_slice(), + [ValueExpression::ListComprehension { map: None, .. }] + ) => + { + let mut lowered = self.rewrite_exists_in_value(input, &mut args[0])?; + let LogicalOperator::ListComprehension { count_only, .. } = &mut lowered else { + unreachable!("a list comprehension lowers to its own operator") + }; + *count_only = true; + *expr = args.remove(0); + Ok(lowered) + } ValueExpression::ScalarFunction { args, .. } | ValueExpression::AggregateFunction { args, .. } => { for arg in args.iter_mut() { @@ -901,6 +922,7 @@ impl<'a> LogicalPlanner<'a> { keep, map: value, output: output.clone(), + count_only: false, }) })(); self.variables = before; diff --git a/docs/api-reference.md b/docs/api-reference.md index d28bb6b..bfb9a55 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -132,6 +132,8 @@ be opened. : Runs `query` and returns the whole result. `params` binds `$name` parameters: `int`, `float`, `str`, `bool`, `None`, lists, string-keyed nested maps, `datetime.datetime` and `datetime.date`. Integers must fit signed 64-bit values; nesting beyond 128 levels, cycles and unsupported values raise `ValueError`. + Direct counts from `size()` of an unmapped list comprehension have a non-nullable Arrow field; + see [Cypher functions](cypher.md#functions). `#!python execute_batches(query: str, params: dict | None = None) -> pyarrow.RecordBatchReader` : Streams the result. Planning errors raise immediately; execution errors raise while reading. The reader diff --git a/docs/cypher.md b/docs/cypher.md index f59fa82..007f724 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -269,6 +269,13 @@ Any other function raises `NotImplementedError: Cypher function 'name' is not im missing: math functions (`abs`, `round`, `sqrt`, ...), most string and list functions, `reduce`, the list predicates `any` / `all` / `none` / `single`, and `relationships(p)`. +When `size()` consumes an unmapped list comprehension, such as `size([x IN values WHERE x > 0])`, the planner +counts retained elements directly. Each duplicate and each retained null element still contributes one; +an empty or null input list produces zero, as for the existing comprehension. A mapped comprehension +(`| expression`) or a pattern comprehension keeps the ordinary list-building path. The count path still +unnests and groups rows; `explain()` shows a filtered count in place of ordered list aggregation. A direct +count output is marked non-nullable in its Arrow schema; the previous list-size expression was marked nullable. + ### Example: a little of everything ```python diff --git a/python/python/tests/test_comprehension_count.py b/python/python/tests/test_comprehension_count.py new file mode 100644 index 0000000..ad04322 --- /dev/null +++ b/python/python/tests/test_comprehension_count.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Count-only comprehensions preserve multiplicity and predicate semantics.""" + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table( + { + "id": [1, 2, 3, 4], + "values": pa.array( + [[1, None, 1, 2], [], None, [None]], type=pa.list_(pa.int64()) + ), + } + ), + # A duplicate edge affects the list's multiplicity; repeated matches + # inside EXISTS must still contribute a single truth value per element. + "KNOWS": pa.table({"src": [1, 1, 1, 2, 2], "dst": [2, 2, 3, 3, 3]}), + }, + use_csr=request.param, + ) + + +@pytest.mark.parametrize( + ("predicate", "expected"), + [ + ("", [4, 0, 0, 1]), + ("WHERE x > 1", [1, 0, 0, 0]), + ("WHERE x IS NULL", [1, 0, 0, 1]), + ("WHERE x IS NULL OR x > 0", [4, 0, 0, 1]), + ], +) +def test_count_matches_list_reference(engine, predicate, expected): + query = ( + "MATCH (p:Person) UNWIND [1, 1] AS duplicate RETURN p.id AS id, " + f"size([x IN p.values {predicate}]) AS n ORDER BY id" + ) + reference = query.replace(f"{predicate}])", f"{predicate} | x])") + rows = [ + {"id": index, "n": n} for index, n in enumerate(expected, 1) for _ in range(2) + ] + assert engine.execute(query).to_pylist() == rows + assert engine.execute_batches(query).read_all().to_pylist() == rows + assert engine.execute(reference).to_pylist() == rows + assert "array_agg(" not in engine.explain(query).lower() + assert "array_agg(" in engine.explain(reference).lower() + + +def test_count_with_node_lookup_and_exists_keeps_duplicate_elements(engine): + query = ( + "MATCH (p:Person) OPTIONAL MATCH (p)-[:KNOWS]->(f:Person) " + "WITH p, collect(f) AS friends " + "RETURN p.id AS id, size([f IN friends WHERE (f)-[:KNOWS]->(:Person)]) AS n " + "ORDER BY id" + ) + expected = [ + {"id": 1, "n": 2}, + {"id": 2, "n": 0}, + {"id": 3, "n": 0}, + {"id": 4, "n": 0}, + ] + assert engine.execute(query).to_pylist() == expected + assert ( + engine.execute(query.replace("(:Person)])", "(:Person) | f])")).to_pylist() + == expected + ) + + +def test_mapped_and_pattern_comprehensions_keep_normal_list_semantics(engine): + query = ( + "MATCH (p:Person) RETURN p.id AS id, " + "size([x IN p.values | null]) AS n ORDER BY id" + ) + assert engine.execute(query).to_pylist() == [ + {"id": 1, "n": 4}, + {"id": 2, "n": 0}, + {"id": 3, "n": 0}, + {"id": 4, "n": 1}, + ] + assert "array_agg(" in engine.explain(query).lower() + pattern = ( + "MATCH (p:Person) RETURN p.id AS id, " + "size([(p)-[:KNOWS]->(f) | f.id]) AS n ORDER BY id" + ) + assert engine.execute(pattern).to_pylist() == [ + {"id": 1, "n": 3}, + {"id": 2, "n": 2}, + {"id": 3, "n": 0}, + {"id": 4, "n": 0}, + ] + assert "array_agg(" in engine.explain(pattern).lower()