fix(query): preserve empty-input distinct counts

crates/lance-csr/src/query.rs

@@ -108,6 +108,43 @@ pub(crate) fn normalize_record_batch(batch: &RecordBatch) -> Result<RecordBatch>
     })
 }
 
+/// 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<datafusion::dataframe::DataFrame> {
         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::<Int64Array>()
+                    .unwrap()
+                    .value(0),
+                i64::from(expected)
+            );
+            assert_eq!(
+                result
+                    .column(1)
+                    .as_any()
+                    .downcast_ref::<BooleanArray>()
+                    .unwrap()
+                    .value(0),
+                expected
+            );
+            assert_eq!(
+                result
+                    .column(2)
+                    .as_any()
+                    .downcast_ref::<Int64Array>()
+                    .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};

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

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}]