feat(csr)!: keep named relationships on the csr with edge rows

crates/lance-csr-catalog/README.md

@@ -1,3 +1,6 @@
 # Lance Graph Catalog
 
 Catalog and namespace utilities shared by the Lance graph query engine.
+
+> **Experimental.** The Rust API is not stable yet: breaking changes are possible over the next few
+> weeks while the CSR work lands. Pin an exact version (`lance-csr-catalog = "=x.y.z"`) if you depend on it.

crates/lance-csr-catalog/src/lib.rs

@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
 //! Catalog and namespace utilities for Lance Graph.
+//!
+//! **Experimental.** The Rust API is not stable yet: breaking changes are possible over the
+//! next few weeks while the CSR work lands. Pin an exact version if you depend on it.
 
 pub mod namespace;
 pub mod source_catalog;

crates/lance-csr/README.md

@@ -5,6 +5,9 @@ with a CSR (compressed sparse row) adjacency index for fast multi-hop traversal.
 Arrow tables as property graphs. The [lance-csr Python package](https://pypi.org/project/lance-csr/) wraps this
 crate.
 
+> **Experimental.** The Rust API is not stable yet: breaking changes are possible over the next few
+> weeks while the CSR work lands. Pin an exact version (`lance-csr = "=x.y.z"`) if you depend on it.
+
 ## Features
 
 - Cypher query parsing and AST construction

crates/lance-csr/src/csr_index.rs

@@ -12,34 +12,46 @@
 //! ```text
 //! Offset Array:  [0, 3, 5, 5, 9, ...]   (one entry per vertex + 1 sentinel)
 //! Neighbor Array: [2, 5, 7, 1, 4, 0, 3, 6, 8, ...]  (destination vertex IDs)
+//! Edge rows:      [...]                            (one source row per neighbour)
 //! ```
 //!
-//! For vertex `v`, its neighbors are `neighbors[offsets[v]..offsets[v+1]]`.
+//! For vertex `v`, its neighbors are `neighbors[offsets[v]..offsets[v+1]]`, and
+//! `edge_rows` at the same positions holds the row each edge came from.
 
 use std::collections::HashMap;
 
 /// In-memory CSR adjacency index for fast neighbor lookup.
 ///
-/// Stores graph topology in two arrays:
+/// Stores graph topology in three arrays:
 /// - `offsets[v]` = start position of vertex v's neighbors in the neighbor array
 /// - `neighbors[offsets[v]..offsets[v+1]]` = destination vertex IDs
+/// - `edge_rows`, aligned with `neighbors` = the row (a Lance row address, or a
+///   row index for in-memory input) each edge came from
 #[derive(Debug, Clone)]
 pub struct CsrIndex {
     offsets: Vec<u64>,
     neighbors: Vec<u64>,
+    edge_rows: Vec<u64>,
 }
 
 impl CsrIndex {
-    /// Look up all neighbors of a vertex. Returns an empty slice for vertices
-    /// with no outgoing edges or vertex IDs beyond the index range.
-    pub fn neighbors(&self, vertex_id: u64) -> &[u64] {
+    fn range(&self, vertex_id: u64) -> std::ops::Range<usize> {
         let v = vertex_id as usize;
         if v >= self.offsets.len() - 1 {
-            return &[];
+            return 0..0;
         }
-        let start = self.offsets[v] as usize;
-        let end = self.offsets[v + 1] as usize;
-        &self.neighbors[start..end]
+        self.offsets[v] as usize..self.offsets[v + 1] as usize
+    }
+
+    /// Look up all neighbors of a vertex. Returns an empty slice for vertices
+    /// with no outgoing edges or vertex IDs beyond the index range.
+    pub fn neighbors(&self, vertex_id: u64) -> &[u64] {
+        &self.neighbors[self.range(vertex_id)]
+    }
+
+    /// The edge rows of a vertex's neighbors, in the same order as [`Self::neighbors`].
+    pub fn edge_rows(&self, vertex_id: u64) -> &[u64] {
+        &self.edge_rows[self.range(vertex_id)]
     }
 
     /// Return the total number of edges in the index.
@@ -50,11 +62,11 @@ impl CsrIndex {
 
 /// Builder for constructing a CSR index from edge data.
 ///
-/// Accepts edges as (source, destination) pairs and builds the compressed
-/// sparse row representation.
+/// Accepts edges as (source, destination, edge row) triples and builds the
+/// compressed sparse row representation.
 #[derive(Debug)]
 pub struct CsrIndexBuilder {
-    edges: Vec<(u64, u64)>,
+    edges: Vec<(u64, u64, u64)>,
     num_vertices: Option<u64>,
 }
 
@@ -73,9 +85,9 @@ impl CsrIndexBuilder {
         self
     }
 
-    /// Add a single directed edge from `src` to `dst`.
-    pub fn add_edge(mut self, src: u64, dst: u64) -> Self {
-        self.edges.push((src, dst));
+    /// Add a single directed edge from `src` to `dst`, read from row `row`.
+    pub fn add_edge(mut self, src: u64, dst: u64, row: u64) -> Self {
+        self.edges.push((src, dst, row));
         self
     }
 
@@ -86,22 +98,22 @@ impl CsrIndexBuilder {
         let num_vertices = self.num_vertices.unwrap_or_else(|| {
             self.edges
                 .iter()
-                .flat_map(|&(s, d)| [s, d])
+                .flat_map(|&(s, d, _)| [s, d])
                 .max()
                 .map(|m| m + 1)
                 .unwrap_or(0)
         });
 
         // Sort by source vertex for CSR construction
-        self.edges.sort_unstable_by_key(|&(src, _)| src);
+        self.edges.sort_unstable_by_key(|&(src, _, _)| src);
 
         // Build offset and neighbor arrays
         let mut offsets = vec![0u64; num_vertices as usize + 1];
         let mut neighbors = Vec::with_capacity(self.edges.len());
 
         // Count degrees
         let mut degree_map: HashMap<u64, u64> = HashMap::new();
-        for &(src, _) in &self.edges {
+        for &(src, _, _) in &self.edges {
             *degree_map.entry(src).or_insert(0) += 1;
         }
 
@@ -113,12 +125,14 @@ impl CsrIndexBuilder {
         }
         offsets[num_vertices as usize] = running;
 
-        // Fill neighbor array
-        for &(_, dst) in &self.edges {
-            neighbors.push(dst);
-        }
+        let edge_rows = self.edges.iter().map(|&(_, _, row)| row).collect();
+        neighbors.extend(self.edges.iter().map(|&(_, dst, _)| dst));
 
-        CsrIndex { offsets, neighbors }
+        CsrIndex {
+            offsets,
+            neighbors,
+            edge_rows,
+        }
     }
 }
 
@@ -140,11 +154,11 @@ mod tests {
     fn sample_index() -> CsrIndex {
         CsrIndexBuilder::new()
             .with_num_vertices(4)
-            .add_edge(0, 1)
-            .add_edge(0, 2)
-            .add_edge(0, 3)
-            .add_edge(1, 2)
-            .add_edge(2, 3)
+            .add_edge(0, 1, 0)
+            .add_edge(0, 2, 1)
+            .add_edge(0, 3, 2)
+            .add_edge(1, 2, 3)
+            .add_edge(2, 3, 4)
             .build()
     }
 
@@ -156,6 +170,9 @@ mod tests {
         assert_eq!(idx.neighbors(1), &[2]);
         assert_eq!(idx.neighbors(2), &[3]);
         assert_eq!(idx.neighbors(3), &[] as &[u64]);
+        assert_eq!(idx.edge_rows(0), &[0, 1, 2]);
+        assert_eq!(idx.edge_rows(2), &[4]);
+        assert_eq!(idx.edge_rows(99), &[] as &[u64]);
     }
 
     #[test]
@@ -181,7 +198,7 @@ mod tests {
     fn test_isolated_vertices() {
         let idx = CsrIndexBuilder::new()
             .with_num_vertices(5)
-            .add_edge(1, 3)
+            .add_edge(1, 3, 5)
             .build();
 
         assert_eq!(idx.neighbors(0), &[] as &[u64]);
@@ -193,7 +210,10 @@ mod tests {
 
     #[test]
     fn test_auto_inferred_num_vertices() {
-        let idx = CsrIndexBuilder::new().add_edge(0, 5).add_edge(3, 7).build();
+        let idx = CsrIndexBuilder::new()
+            .add_edge(0, 5, 6)
+            .add_edge(3, 7, 7)
+            .build();
 
         // Inferred from the max ID (7): vertices up to 7 are addressable.
         assert_eq!(idx.neighbors(7), &[] as &[u64]);
@@ -205,9 +225,9 @@ mod tests {
     fn test_self_loops() {
         let idx = CsrIndexBuilder::new()
             .with_num_vertices(3)
-            .add_edge(0, 0)
-            .add_edge(1, 1)
-            .add_edge(0, 1)
+            .add_edge(0, 0, 8)
+            .add_edge(1, 1, 9)
+            .add_edge(0, 1, 10)
             .build();
 
         assert_eq!(idx.neighbors(0), &[0, 1]);
@@ -218,9 +238,9 @@ mod tests {
     fn test_parallel_edges() {
         let idx = CsrIndexBuilder::new()
             .with_num_vertices(2)
-            .add_edge(0, 1)
-            .add_edge(0, 1)
-            .add_edge(0, 1)
+            .add_edge(0, 1, 11)
+            .add_edge(0, 1, 12)
+            .add_edge(0, 1, 13)
             .build();
 
         // CSR preserves multi-edges

crates/lance-csr/src/csr_traversal.rs

@@ -161,12 +161,12 @@ impl TraversalCursor {
             let next = (step.from == frame.label)
                 .then(|| {
                     step.adjacency
-                        .candidates_from(frame.key, &mut frame.position)
+                        .candidates_from::<false>(frame.key, &mut frame.position)
                         .next()
                 })
                 .flatten();
             let key = match next {
-                Some(Some(key)) => key as i64,
+                Some(Some((key, _))) => key as i64,
                 Some(None) => continue,
                 None => {
                     frame.step += 1;
@@ -298,7 +298,11 @@ mod tests {
         catalog.insert(
             "hasType",
             true,
-            Arc::new(base.with_changes(&[], &vec![(1, 1); 4999])),
+            Arc::new(base.with_changes(&[], &{
+                let mut removed = crate::csr_expand::RemovedRows::default();
+                removed.rows(0, 0..4999);
+                removed
+            })),
         );
         let walk = Traversal::new(
             &config,
@@ -351,10 +355,10 @@ mod tests {
     #[tokio::test]
     async fn walk_yields_while_no_candidate_can_match_the_target_label() {
         let (config, mut catalog) = graph().await;
-        let adjacency = catalog
-            .get("isSubclass", true)
-            .unwrap()
-            .with_changes(&vec![(1, 1); 4096], &[]);
+        let adjacency = catalog.get("isSubclass", true).unwrap().with_changes(
+            &(0..4096).map(|row| (1, 1, row)).collect::<Vec<_>>(),
+            &Default::default(),
+        );
         catalog.insert("isSubclass", true, Arc::new(adjacency));
         let walk = Traversal::new(
             &config,

crates/lance-csr/src/datafusion_planner/builder/aggregate_ops.rs

@@ -36,24 +36,61 @@ fn split_aggregates(expr: &mut ValueExpression, aggregates: &mut Vec<(String, Va
     });
 }
 
-/// `count(v)` lowers to `count(v__id)`, which only nodes and relationships have. For a
+/// `count(v)` lowers to `count(v__id)`, which nodes resolve to their key earlier. For a
 /// value without an id column (a path, or an UNWIND or WITH value), count the value itself.
+/// For a relationship, count the rows where any of its `v__*` columns is set: a matched
+/// relationship has non-null endpoints, an unmatched OPTIONAL one has none. `DISTINCT`
+/// compares all its columns, so parallel edges with identical properties count once.
 fn count_values_without_ids(
     args: &mut [datafusion::logical_expr::Expr],
+    distinct: bool,
     schema: &datafusion::common::DFSchema,
 ) {
+    use datafusion::logical_expr::{lit, when, Expr};
     for arg in args {
-        let datafusion::logical_expr::Expr::Column(column) = arg else {
+        let Expr::Column(column) = arg else {
             continue;
         };
         let Some(variable) = column.name.strip_suffix("__id") else {
             continue;
         };
-        if !schema.has_column_with_unqualified_name(&column.name)
-            && schema.has_column_with_unqualified_name(variable)
-        {
+        if schema.has_column_with_unqualified_name(&column.name) {
+            continue;
+        }
+        if schema.has_column_with_unqualified_name(variable) {
             *arg = col(datafusion::common::Column::new_unqualified(variable));
+            continue;
         }
+        let prefix = format!("{variable}__");
+        let columns: Vec<_> = schema
+            .iter()
+            .filter(|(_, f)| f.name().starts_with(&prefix))
+            .map(|(q, f)| {
+                (
+                    f.name()[prefix.len()..].to_string(),
+                    Expr::Column((q, f).into()),
+                )
+            })
+            .collect();
+        let Some(bound) = columns
+            .iter()
+            .map(|(_, c)| c.clone().is_not_null())
+            .reduce(Expr::or)
+        else {
+            continue;
+        };
+        let value = if distinct {
+            datafusion::functions::core::expr_fn::named_struct(
+                columns
+                    .into_iter()
+                    .flat_map(|(name, c)| [lit(name), c])
+                    .collect(),
+            )
+        } else {
+            lit(1_i64)
+        };
+        // Building a CASE can only fail on an empty WHEN list, which `bound` rules out.
+        *arg = when(bound, value).end().expect("CASE with one WHEN");
     }
 }
 
@@ -108,7 +145,8 @@ impl DataFusionPlanner {
                             f.params.order_by = ordering.to_vec();
                         }
                         if f.func.name() == "count" {
-                            count_values_without_ids(&mut f.params.args, &schema);
+                            let distinct = f.params.distinct;
+                            count_values_without_ids(&mut f.params.args, distinct, &schema);
                         }
                     }
                     agg_exprs.push(aggregate.alias(column));
@@ -231,4 +269,27 @@ mod tests {
             s
         );
     }
+
+    #[test]
+    fn count_of_a_relationship_counts_rows_where_it_matched() {
+        use arrow_schema::{DataType, Field, Schema};
+        let schema = datafusion::common::DFSchema::try_from(Schema::new(vec![
+            Field::new("r__src", DataType::Int64, true),
+            Field::new("r__since", DataType::Int64, true),
+            Field::new("p__id", DataType::Int64, true),
+        ]))
+        .unwrap();
+        let count = |distinct| {
+            let mut args = vec![datafusion::logical_expr::col("r__id")];
+            super::count_values_without_ids(&mut args, distinct, &schema);
+            args[0].to_string()
+        };
+        let plain = count(false);
+        assert!(
+            plain.contains("r__src IS NOT NULL OR r__since IS NOT NULL"),
+            "{plain}"
+        );
+        assert!(!plain.contains("p__id"), "{plain}");
+        assert!(count(true).contains("named_struct"), "{}", count(true));
+    }
 }