From 80d4438884872ce9d18c528caa1519441a9460ce Mon Sep 17 00:00:00 2001 From: David Hughes Date: Mon, 26 Oct 2026 16:35:31 -0700 Subject: [PATCH] fix(query): preserve empty-input distinct counts Refs lance-format/lance-graph#123. Adapt the focused fix to the current 0.3.1 engine and crate layout. Submitted by David Hughes (@stugorf), with AI Collaboration. --- crates/lance-csr/src/query.rs | 127 +++++++++++++++-- docs/cypher.md | 12 ++ .../python/tests/test_aggregate_predicates.py | 131 ++++++++++++++++++ 3 files changed, 257 insertions(+), 13 deletions(-) create mode 100644 python/python/tests/test_aggregate_predicates.py diff --git a/crates/lance-csr/src/query.rs b/crates/lance-csr/src/query.rs index d62c68b..71702c3 100644 --- a/crates/lance-csr/src/query.rs +++ b/crates/lance-csr/src/query.rs @@ -108,6 +108,43 @@ pub(crate) fn normalize_record_batch(batch: &RecordBatch) -> Result }) } +/// Build a query-local frame without unsafe distinct-count statistics shortcuts. +fn query_dataframe( + ctx: &datafusion::execution::context::SessionContext, + plan: datafusion::logical_expr::LogicalPlan, +) -> datafusion::dataframe::DataFrame { + use datafusion::execution::session_state::SessionStateBuilder; + use datafusion::logical_expr::{Expr, LogicalPlan}; + + fn has_distinct_count(plan: &LogicalPlan) -> bool { + if let LogicalPlan::Aggregate(aggregate) = plan { + if aggregate.aggr_expr.iter().any(|expr| { + matches!(expr.clone().unalias(), Expr::AggregateFunction(f) + if f.func.name() == "count" && f.params.distinct) + }) { + return true; + } + } + plan.inputs().into_iter().any(has_distinct_count) + } + + let mut state = ctx.state(); + if has_distinct_count(&plan) { + // shortcut: DataFusion 54 treats filtered distinct estimates as exact; remove + // when upstream COUNT DISTINCT statistics preserve empty-input semantics. + let rules = state + .physical_optimizers() + .iter() + .filter(|rule| rule.name() != "aggregate_statistics") + .cloned() + .collect(); + state = SessionStateBuilder::from(state) + .with_physical_optimizer_rules(rules) + .build(); + } + datafusion::dataframe::DataFrame::new(state, plan) +} + /// Execution strategy for Cypher queries #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum ExecutionStrategy { @@ -796,12 +833,7 @@ impl CypherQuery { ) -> Result { let (_logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; let ctx = deadline::context(ctx); - ctx.execute_logical_plan(df_logical_plan) - .await - .map_err(|e| GraphError::ExecutionError { - message: format!("Failed to execute DataFusion plan: {}", e), - location: snafu::Location::new(file!(), line!(), column!()), - }) + Ok(query_dataframe(&ctx, df_logical_plan)) } /// Execute using the DataFusion planner with in-memory datasets @@ -1146,13 +1178,7 @@ impl CypherQuery { // Phase 4: DataFusion Physical Plan let ctx = deadline::context(ctx.clone()); - let df = ctx - .execute_logical_plan(df_logical_plan.clone()) - .await - .map_err(|e| GraphError::ExecutionError { - message: format!("Failed to execute DataFusion plan: {}", e), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + let df = query_dataframe(&ctx, df_logical_plan.clone()); let physical_plan = df.create_physical_plan() @@ -2348,6 +2374,81 @@ mod tests { assert_eq!(name_set, expected, "Should return Alice and Bob"); } + #[tokio::test] + async fn test_aggregate_predicates_on_empty_input() { + use arrow_array::{BooleanArray, Int64Array}; + use datafusion::execution::context::SessionContext; + + let ctx = SessionContext::new(); + ctx.register_batch( + "Person", + RecordBatch::try_from_iter(vec![("id", Arc::new(Int64Array::from(vec![1, 2])) as _)]) + .unwrap(), + ) + .unwrap(); + let rules_before: Vec<_> = ctx + .state() + .physical_optimizers() + .iter() + .map(|rule| rule.name().to_owned()) + .collect(); + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + for (predicate, expected) in [("p.id = 99", false), ("p.id = 1", true)] { + let query = CypherQuery::new(&format!( + "MATCH (p:Person) WHERE {predicate} RETURN count(DISTINCT p.id) AS n, \ + count(DISTINCT p.id) > 0 AS present, \ + CASE WHEN count(*) > 0 THEN 1 ELSE 0 END AS flag" + )) + .unwrap() + .with_config(config.clone()); + let result = query.execute_with_context(ctx.clone()).await.unwrap(); + assert_eq!(result.num_rows(), 1); + assert_eq!( + result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + i64::from(expected) + ); + assert_eq!( + result + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + expected + ); + assert_eq!( + result + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + i64::from(expected) + ); + } + let rules_after: Vec<_> = ctx + .state() + .physical_optimizers() + .iter() + .map(|rule| rule.name().to_owned()) + .collect(); + assert_eq!( + rules_before, rules_after, + "caller session must remain unchanged" + ); + assert!(rules_after + .iter() + .any(|name| name == "aggregate_statistics")); + } + #[tokio::test] async fn test_execute_with_context_simple_scan() { use arrow_array::{Int64Array, RecordBatch, StringArray}; diff --git a/docs/cypher.md b/docs/cypher.md index d0452fb..d10a70a 100644 --- a/docs/cypher.md +++ b/docs/cypher.md @@ -502,6 +502,18 @@ rejects graph patterns, comprehensions, and aggregation. Compute such a list bef These semantics follow the pinned [Cypher 4.4 list reference](https://neo4j.com/docs/cypher-manual/4.4/functions/list/#functions-reduce) and [predicate reference](https://neo4j.com/docs/cypher-manual/4.4/functions/predicate/). +Aggregate results can be compared in `RETURN` or `WITH`, combined with `AND`, `OR` and `NOT`, +or used in a `CASE` condition. For example: + +```cypher +MATCH (p:Person) WHERE p.id = 99 +RETURN count(DISTINCT p.id) > 0 AS present +``` + +If no rows match, this returns one row with `present = false`. A global `count` returns `0` +on empty input, and `count(DISTINCT value)` skips nulls and duplicate values. These rules +also apply when a filter constrains the counted property to a constant. + ### Example: a little of everything ```python diff --git a/python/python/tests/test_aggregate_predicates.py b/python/python/tests/test_aggregate_predicates.py new file mode 100644 index 0000000..f58841f --- /dev/null +++ b/python/python/tests/test_aggregate_predicates.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Aggregate predicates must preserve empty-input and null semantics.""" + +import lance +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, CypherQuery, GraphConfig + + +@pytest.fixture(params=["arrow", "lance"]) +def graph(request, tmp_path): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .node_table("Post", "id") + .rel_table("likePost", "Person", "Post", "src", "dst") + .build() + ) + tables = { + "Person": pa.table( + { + "id": [1, 2, 3], + "firstname": ["Bill", "Bill", "Other"], + "lastname": ["Moore", "Moore", "Person"], + "score": [None, 2, 2], + } + ), + "Post": pa.table({"id": [1]}), + "likePost": pa.table({"src": [1, 1, 2], "dst": [1, 1, 1]}), + } + sources = ( + tables + if request.param == "arrow" + else { + name: lance.write_dataset(table, tmp_path / f"{name}.lance") + for name, table in tables.items() + } + ) + return config, tables, sources + + +def execute(graph, mode, query): + config, tables, sources = graph + if mode == "legacy": + return CypherQuery(query).with_config(config).execute(tables) + engine = CypherEngine(config, sources, use_csr="csr" in mode) + if mode.startswith("stream"): + return engine.execute_batches(query).read_all() + return engine.execute(query) + + +@pytest.mark.parametrize("mode", ["legacy", "join", "csr", "stream-join", "stream-csr"]) +@pytest.mark.parametrize( + ("query", "expected"), + [ + ("MATCH (p:Person) RETURN count(DISTINCT p.id) > 0 AS x", [{"x": True}]), + ( + "MATCH (p:Person) WHERE p.id = 99 RETURN count(DISTINCT p.id) AS x", + [{"x": 0}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 RETURN count(DISTINCT p.id) > 0 AS x", + [{"x": False}], + ), + ( + "MATCH (p:Person) WHERE p.id = 1 RETURN count(DISTINCT p.id) > 0 AS x", + [{"x": True}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 " + "RETURN NOT (count(DISTINCT p.id) > 0) AS x", + [{"x": True}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 " + "RETURN CASE WHEN count(*) > 0 THEN 1 ELSE 0 END AS x", + [{"x": 0}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 " + "RETURN count(DISTINCT p.id) > 0 OR count(*) = 0 AS x", + [{"x": True}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 RETURN sum(p.score) IS NULL AS x", + [{"x": True}], + ), + ( + "MATCH (p:Person) WHERE p.id = 99 " + "WITH count(DISTINCT p.id) > 0 AS x RETURN x", + [{"x": False}], + ), + ( + "MATCH (p:Person) RETURN count(DISTINCT p.score) AS x, count(p.score) AS y", + [{"x": 1, "y": 2}], + ), + ( + "MATCH (p:Post)<-[:likePost]-(p2:Person) " + 'WHERE p2.firstname = "Bill" AND p2.lastname = "Moore" AND p.id = 2 ' + "RETURN COUNT(DISTINCT p.id) > 0 AS liked", + [{"liked": False}], + ), + ( + "MATCH (p:Post)<-[:likePost]-(p2:Person) " + 'WHERE p2.firstname = "Bill" AND p2.lastname = "Moore" AND p.id = 1 ' + "RETURN COUNT(DISTINCT p.id) AS liked", + [{"liked": 1}], + ), + ], +) +def test_aggregate_predicates(graph, mode, query, expected): + result = execute(graph, mode, query) + assert result.to_pylist() == expected + for field, value in zip(result.schema, expected[0].values()): + assert field.type == (pa.bool_() if isinstance(value, bool) else pa.int64()) + + +@pytest.mark.parametrize("use_csr", [False, True]) +def test_cached_aggregate_predicates(graph, use_csr): + config, _, sources = graph + engine = CypherEngine(config, sources, use_csr=use_csr) + query = "MATCH (p:Person) WHERE p.id = $id RETURN count(DISTINCT p.id) > 0 AS x" + for value, expected in [(1, True), (99, False), (2, True), (99, False)]: + result = ( + engine.execute_batches(query, {"id": value}).read_all() + if value == 99 + else engine.execute(query, {"id": value}) + ) + assert result.to_pylist() == [{"x": expected}]