Count qualifying comprehension elements without building a list

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.

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::<Int64Array>()
+            .unwrap();
+        assert_eq!(counts.values().iter().sum::<i64>(), 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);

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

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<LogicalPlan> {
         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<String> = numbered
@@ -160,11 +168,102 @@ 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))
             .and_then(|b| b.build())
             .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<String, RecordBatch> {
+        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::<Int64Type, _, _>(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::<Int64Array>()
+                .unwrap();
+            assert_eq!(
+                counts.values().to_vec(),
+                expected
+                    .into_iter()
+                    .flat_map(|n| [n, n])
+                    .collect::<Vec<_>>()
+            );
+            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}");
+        }
+    }
+}

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

crates/lance-graph/src/logical_plan.rs

@@ -172,6 +172,9 @@ pub enum LogicalOperator {
         keep: Option<BooleanExpression>,
         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;

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

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