freeeve / rustychickpeas / 27909064898

//! GraphSnapshot Python wrapper

use crate::direction::Direction;

use crate::node::Node;

use crate::relationship::Relationship;

use crate::utils::{py_to_property_value, value_id_to_pyobject};

use pyo3::ffi;

use pyo3::prelude::*;

use pyo3::IntoPyObjectExt;

use roaring::RoaringBitmap;

use rustychickpeas_core::bitmap::NodeSet;

use rustychickpeas_core::types::PropertyKey;

use rustychickpeas_core::{

    AggOp, ColumnDtype, GraphSnapshot as CoreGraphSnapshot, Label, RelationshipRef,

    RelationshipType, ValueId,

};

use std::os::raw::{c_char, c_int, c_void};

use std::sync::{Arc, Mutex, PoisonError};

/// Python wrapper for GraphSnapshot

#[pyclass(name = "GraphSnapshot")]

pub struct GraphSnapshot {

    pub(crate) snapshot: std::sync::Arc<CoreGraphSnapshot>,

}

/// Iterator over the node IDs of a GraphSnapshot

///

/// Node IDs in a finalized snapshot are dense in `0..n_nodes` (the same range

/// accepted by `GraphSnapshot.node()`), so iteration only needs the bounds.

#[pyclass(name = "NodeIdIter")]

pub struct NodeIdIter {

    current: u32,

    end: u32,

}

#[pymethods]

impl NodeIdIter {

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {

        slf

    }

    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<u32> {

        if slf.current < slf.end {

            let id = slf.current;

            slf.current += 1;

            Some(id)

        } else {

            None

        }

    }

}

impl GraphSnapshot {

    /// Get string ID from string

    /// Uses the reverse index in Atoms for O(1) lookup

    fn get_string_id(&self, s: &str) -> Option<u32> {

        self.snapshot.atoms.get_id(s)

    }

    /// Get label from string

    fn label_from_str(&self, s: &str) -> Option<Label> {

        self.get_string_id(s).map(Label::new)

    }

    /// Get relationship type from string

    fn rel_type_from_str(&self, s: &str) -> Option<RelationshipType> {

        self.get_string_id(s).map(RelationshipType::new)

    }

    /// Get property key from string

    fn property_key_from_str(&self, s: &str) -> Option<PropertyKey> {

        self.get_string_id(s)

    }

    /// Convert a Python value to a core [`ValueId`], resolving a string to its

    /// interned atom id. Returns `Ok(None)` when a string value is not interned

    /// in this snapshot — no node can carry it, so predicate/lookup callers

    /// short-circuit to "no match" rather than erroring.

    fn py_value_to_id_opt(&self, value: &Bound<'_, PyAny>) -> PyResult<Option<ValueId>> {

        use rustychickpeas_core::PropertyValue;

        Ok(match py_to_property_value(value)? {

            PropertyValue::String(s) => self.get_string_id(&s).map(ValueId::Str),

            PropertyValue::Integer(i) => Some(ValueId::I64(i)),

            PropertyValue::Float(f) => Some(ValueId::from_f64(f)),

            PropertyValue::Boolean(b) => Some(ValueId::Bool(b)),

            PropertyValue::InternedString(_) => {

                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(

                    "InternedString not supported here",

                ));

            }

        })

    }

}

impl GraphSnapshot {

    /// Internal constructor (not exposed to Python)

    /// Takes ownership of a GraphSnapshot

    pub(crate) fn new(snapshot: CoreGraphSnapshot) -> Self {

        Self {

            snapshot: std::sync::Arc::new(snapshot),

        }

    }

    /// Internal constructor from Arc (for manager.get_graph_snapshot)

    pub(crate) fn from_arc(snapshot: std::sync::Arc<CoreGraphSnapshot>) -> Self {

        Self { snapshot }

    }

}

#[pymethods]

impl GraphSnapshot {

    fn __repr__(&self) -> String {

        let version = self

            .snapshot

            .version()

            .map(|v| v.to_string())

            .unwrap_or_else(|| "None".to_string());

        format!(

            "GraphSnapshot(nodes={}, rels={}, version={})",

            self.snapshot.n_nodes, self.snapshot.n_rels, version

        )

    }

    fn __len__(&self) -> usize {

        self.snapshot.n_nodes as usize

    }

    /// Iterate over all node IDs in the snapshot (0..n_nodes)

    ///

    /// Yields exactly the node IDs accepted by `node()`.

    fn __iter__(&self) -> NodeIdIter {

        NodeIdIter {

            current: 0,

            end: self.snapshot.n_nodes,

        }

    }

    /// Get number of nodes

    fn node_count(&self) -> u32 {

        self.snapshot.n_nodes

    }

    /// Get number of relationships

    fn relationship_count(&self) -> u64 {

        self.snapshot.n_rels

    }

    /// Get node labels

    fn node_labels(&self, node_id: u32) -> PyResult<Vec<String>> {

        // GraphSnapshot doesn't store labels per node directly

        // We need to iterate through label_index to find which labels contain this node

        let mut labels = Vec::new();

        for (label, node_set) in &self.snapshot.label_index {

            if node_set.contains(node_id) {

                if let Some(label_str) = self.snapshot.resolve_string(label.id()) {

                    labels.push(label_str.to_string());

                }

            }

        }

        Ok(labels)

    }

    /// Get nodes with a specific label

    fn nodes_with_label(&self, label: String) -> PyResult<Vec<u32>> {

        // Call Rust function - it returns None if label doesn't exist

        // We need to distinguish between "label doesn't exist" vs "label exists but no nodes"

        // Since nodes_with_label_id returns Option<&NodeSet>, None means label not in index

        let label_id = self.label_from_str(&label);

        if label_id.is_none() {

            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(

                "Label '{}' not found",

                label

            )));

        }

        // Label exists, now get nodes (will return Some even if empty)

        if let Some(node_set) = self.snapshot.nodes_with_label(&label) {

            Ok(node_set.iter().collect())

        } else {

            // This shouldn't happen if label exists, but handle it

            Ok(Vec::new())

        }

    }

    /// Get a Node object for the given node ID

    fn node(&self, node_id: u32) -> PyResult<Node> {

        if node_id >= self.snapshot.n_nodes {

            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(

                "Node ID {} out of range (max: {})",

                node_id,

                self.snapshot.n_nodes.saturating_sub(1)

            )));

        }

        Ok(Node {

            snapshot: self.snapshot.clone(),

            node_id,

        })

    }

    /// Get relationships (neighbors) of a node with optional type filtering

    ///

    /// # Arguments

    /// * `node_id` - The node ID

    /// * `direction` - Direction of relationships (Outgoing, Incoming, Both)

    /// * `rel_types` - Optional list of relationship types to filter by

    #[pyo3(signature = (node_id, direction, rel_types=None))]

    fn relationships(

        &self,

        node_id: u32,

        direction: Direction,

        rel_types: Option<Vec<String>>,

    ) -> PyResult<Vec<Relationship>> {

        use rustychickpeas_core::types::RelationshipType;

        if node_id >= self.snapshot.n_nodes {

            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(

                "Node ID {} out of range (max: {})",

                node_id,

                self.snapshot.n_nodes.saturating_sub(1)

            )));

        }

        // Convert string types to RelationshipType IDs via the public resolver.

        let rel_type_ids: Option<Vec<RelationshipType>> = rel_types.as_ref().and_then(|types| {

            let ids: Vec<RelationshipType> = types

                .iter()

                .filter_map(|s| self.snapshot.rel_type(s))

                .collect();

            if ids.is_empty() && !types.is_empty() {

                return None;

            }

            Some(ids)

        });

        let mut relationships = Vec::new();

        // Handle outgoing relationships

        if matches!(direction, Direction::Outgoing | Direction::Both) {

            let start = self.snapshot.out_offsets[node_id as usize] as usize;

            let end = self.snapshot.out_offsets[node_id as usize + 1] as usize;

            for (idx, (&_neighbor, &rel_type)) in self.snapshot.out_nbrs[start..end]

                .iter()

                .zip(self.snapshot.out_types[start..end].iter())

                .enumerate()

            {

                let rel_csr_index = start + idx;

                // Apply type filter if provided

                if let Some(ref type_ids) = rel_type_ids {

                    if !type_ids.contains(&rel_type) {

                        continue;

                    }

                } else if rel_types.as_ref().map(|t| !t.is_empty()).unwrap_or(false) {

                    // Filter was provided but no types found - skip

                    continue;

                }

                relationships.push(Relationship {

                    snapshot: self.snapshot.clone(),

                    rel_index: rel_csr_index as u32,

                    is_outgoing: true,

                });

            }

        }

        // Handle incoming relationships

        if matches!(direction, Direction::Incoming | Direction::Both) {

            let start = self.snapshot.in_offsets[node_id as usize] as usize;

            let end = self.snapshot.in_offsets[node_id as usize + 1] as usize;

            for (idx, (&_neighbor, &rel_type)) in self.snapshot.in_nbrs[start..end]

                .iter()

                .zip(self.snapshot.in_types[start..end].iter())

                .enumerate()

            {

                let rel_csr_index = start + idx;

                // Apply type filter if provided

                if let Some(ref type_ids) = rel_type_ids {

                    if !type_ids.contains(&rel_type) {

                        continue;

                    }

                } else if rel_types.as_ref().map(|t| !t.is_empty()).unwrap_or(false) {

                    // Filter was provided but no types found - skip

                    continue;

                }

                relationships.push(Relationship {

                    snapshot: self.snapshot.clone(),

                    rel_index: rel_csr_index as u32,

                    is_outgoing: false,

                });

            }

        }

        Ok(relationships)

    }

    /// Neighbor node IDs in `direction`, optionally restricted to `rel_types`

    /// (deduplicated, ascending, when types are given).

    #[pyo3(signature = (node_id, direction, rel_types=None))]

    fn neighbor_ids(

        &self,

        node_id: u32,

        direction: Direction,

        rel_types: Option<Vec<String>>,

    ) -> Vec<u32> {

        match rel_types {

            None => self.snapshot.neighbors(node_id, direction.into()).collect(),

            Some(types) => {

                let mut set = RoaringBitmap::new();

                for t in &types {

                    for n in self

                        .snapshot

                        .neighbors_by_type(node_id, direction.into(), t.as_str())

                    {

                        set.insert(n);

                    }

                }

                set.iter().collect()

            }

        }

    }

    /// Histogram of the neighbours reached from `sources` via `rel_type` rels in

    /// `direction`: for each source node, count how many of its `rel_type`

    /// neighbours land on each target. Returns a dict mapping target node id to

    /// count. The whole aggregation runs in Rust on a single call, so it is far

    /// faster than counting neighbours in a Python loop.

    fn neighbor_counts(

        &self,

        sources: Vec<u32>,

        direction: Direction,

        rel_type: &str,

    ) -> std::collections::HashMap<u32, usize> {

        // Core returns a hashbrown map; collect into a std map so PyO3 hands Python a dict.

        self.snapshot

            .neighbor_counts(sources, direction.into(), rel_type)

            .into_iter()

            .collect()

    }

    /// Get neighbors of a node as Node objects

    /// Returns a list of Node objects for neighbors in the specified direction

    #[pyo3(signature = (node_id, direction, rel_types=None))]

    fn neighbors(

        &self,

        node_id: u32,

        direction: Direction,

        rel_types: Option<Vec<String>>,

    ) -> PyResult<Vec<Node>> {

        // Get relationships (optionally type-filtered) and extract neighbor IDs.

        let rels = self.relationships(node_id, direction, rel_types)?;

        let mut neighbor_ids = Vec::new();

        for rel in rels {

            let neighbor_id = if rel.is_outgoing {

                // For outgoing relationships, the end node is in out_nbrs

                let idx = rel.rel_index as usize;

                if idx < self.snapshot.out_nbrs.len() {

                    self.snapshot.out_nbrs[idx]

                } else {

                    continue; // Skip invalid relationship

                }

            } else {

                // For incoming relationships, the start node (neighbor) is in in_nbrs

                let idx = rel.rel_index as usize;

                if idx < self.snapshot.in_nbrs.len() {

                    self.snapshot.in_nbrs[idx]

                } else {

                    continue; // Skip invalid relationship

                }

            };

            neighbor_ids.push(neighbor_id);

        }

        Ok(neighbor_ids

            .into_iter()

            .map(|id| Node {

                snapshot: self.snapshot.clone(),

                node_id: id,

            })

            .collect())

    }

    /// Degree of a node — O(1) from the CSR offsets when untyped; with

    /// `rel_type`, the count of neighbors reached via that type.

    #[pyo3(signature = (node_id, direction, rel_type=None))]

    fn degree(&self, node_id: u32, direction: Direction, rel_type: Option<&str>) -> usize {

        match rel_type {

            Some(rt) => self

                .snapshot

                .neighbors_by_type(node_id, direction.into(), rt)

                .count(),

            None => crate::utils::csr_degree(&self.snapshot, node_id, direction.into()),

        }

    }

    /// Whether `node_id` carries `label` — an O(1) label-index check, vs the

    /// `"X" in node_labels(n)` scan it replaces.

    fn has_label(&self, node_id: u32, label: &str) -> bool {

        self.snapshot.has_label(node_id, label)

    }

    /// Whether `node_id` has any neighbor via `rel_type` in `direction`

    /// (existence check, short-circuits on the first match).

    fn has_rel(&self, node_id: u32, direction: Direction, rel_type: &str) -> bool {

        self.snapshot.has_rel(node_id, direction.into(), rel_type)

    }

    /// Whether `node_id` has a neighbor (via `rel_type`, `direction`) whose node

    /// property `key` equals `value` — resolves the value once, then compares ids

    /// per neighbor (vs a per-neighbor Python property read).

    fn has_neighbor_with_property(

        &self,

        node_id: u32,

        direction: Direction,

        rel_type: &str,

        key: &str,

        value: &Bound<'_, PyAny>,

    ) -> PyResult<bool> {

        let Some(value_id) = self.py_value_to_id_opt(value)? else {

            return Ok(false);

        };

        Ok(self.snapshot.has_neighbor_with_property(

            node_id,

            direction.into(),

            rel_type,

            key,

            value_id,

        ))

    }

    /// The smallest node carrying `label` with property `key` == `value`, or

    /// `None` — collapses `nodes_with_property(..)[0]`, label-scoped so a `name`

    /// shared across labels stays unambiguous.

    fn node_with_label_property(

        &self,

        label: &str,

        key: &str,

        value: &Bound<'_, PyAny>,

    ) -> PyResult<Option<u32>> {

        let Some(value_id) = self.py_value_to_id_opt(value)? else {

            return Ok(None);

        };

        Ok(self.snapshot.node_with_label_property(label, key, value_id))

    }

    /// First neighbor of `node_id` via `rel_type` in `direction`, or `None` —

    /// returns the id (not a `Node`), short-circuiting the scan.

    fn first_neighbor(&self, node_id: u32, direction: Direction, rel_type: &str) -> Option<u32> {

        self.snapshot

            .first_neighbor(node_id, direction.into(), rel_type)

    }

    /// Follow a fixed chain of `(direction, rel_type)` steps from `start`, taking

    /// the first neighbor at each; `None` if a step has no neighbor. Returns the

    /// final node id (not a list of `Node`s).

    fn follow(&self, start: u32, steps: Vec<(Direction, String)>) -> Option<u32> {

        let steps: Vec<(rustychickpeas_core::Direction, &str)> = steps

            .iter()

            .map(|(d, r)| ((*d).into(), r.as_str()))

            .collect();

        self.snapshot.follow(start, &steps)

    }

    /// The root each node reaches by following the *functional* `rel` in `direction`

    /// (each node has one successor — e.g. a message's `replyOf` thread root); a node

    /// already terminal maps to itself. Returns a `NodeArray` array indexed by node id —

    /// `roots[node]` or `memoryview(roots)` in a hot loop. The forest-root array is

    /// built once and cached on the snapshot, so this is the bulk form to reach for

    /// over a per-node `root_via`. `None` if `rel` is unknown.

    fn roots_via(&self, rel: &str, direction: Direction) -> Option<NodeArray> {

        let rt = self.snapshot.relationship_type_from_str(rel)?;

        let inner = self.snapshot.chain_roots(direction.into(), rt);

        let len = inner.len() as ffi::Py_ssize_t;

        Some(NodeArray {

            inner,

            shape: [len],

            strides: [4],

        })

    }

    /// The root of a single `node` via the functional `rel` in `direction` (see

    /// `roots_via`). Convenience for a one-off lookup; in a per-node loop prefer

    /// `roots_via` and index it. `None` if `rel` is unknown.

    fn root_via(&self, node: u32, rel: &str, direction: Direction) -> Option<u32> {

        let rt = self.snapshot.relationship_type_from_str(rel)?;

        Some(self.snapshot.chain_root(node, direction.into(), rt))

    }

    /// The single neighbor each node reaches via the *functional* `rel` in `direction`

    /// (one hop — e.g. a message's `hasCreator` -> its creator). The depth-1 sibling of

    /// `roots_via`: where that follows the chain to its terminal, this takes one step.

    /// Returns a `NodeArray` indexed by node id; a node with no such neighbor maps to

    /// `u32::MAX` (4294967295). Built fresh each call (one `first_neighbor` per node,

    /// GIL released), so hold the result for a hot loop. `None` if `rel` is unknown.

    fn neighbor_via(&self, py: Python<'_>, rel: &str, direction: Direction) -> Option<NodeArray> {

        let rt = self.snapshot.relationship_type_from_str(rel)?;

        let dir: rustychickpeas_core::Direction = direction.into();

        let snapshot = self.snapshot.clone();

        let inner: Arc<[u32]> = py.allow_threads(move || {

            let v: Vec<u32> = (0..snapshot.n_nodes)

                .map(|node| snapshot.first_neighbor(node, dir, rt).unwrap_or(u32::MAX))

                .collect();

            v.into()

        });

        let len = inner.len() as ffi::Py_ssize_t;

        Some(NodeArray {

            inner,

            shape: [len],

            strides: [4],

        })

    }

    /// Fold relationship `rel` (in `direction`) into a `PairWeights` map by projecting

    /// both endpoints of each rel through `projection` (a `NodeArray`, e.g. from

    /// `neighbor_via` or `roots_via`) — the one-mode / bipartite projection ("network

    /// folding") of a relation onto a derived node set. For each `rel` rel `a -> b`,

    /// the unordered pair `(min, max)` of `projection[a]` / `projection[b]` gets one

    /// count; self-pairs and endpoints mapping to the `u32::MAX` sentinel are skipped.

    /// Runs the parallel core kernel with the GIL released. The result stays resident

    /// (no per-pair Python object) so it can drive a native weighted `dijkstra` without

    /// a per-rel callback; `to_dict()` materializes it. E.g. BI Q19's person

    /// interaction graph: ``g.fold_via("replyOf", Direction.Outgoing,

    /// g.neighbor_via("hasCreator", Direction.Incoming))``.

    fn fold_via(

        &self,

        py: Python<'_>,

        rel: &str,

        direction: Direction,

        projection: &NodeArray,

    ) -> PairWeights {

        let snapshot = self.snapshot.clone();

        let dir: rustychickpeas_core::Direction = direction.into();

        let rel = rel.to_owned();

        let proj = projection.inner.clone();

        let map: std::collections::HashMap<(u32, u32), u64> = py.allow_threads(move || {

            snapshot

                .fold_via(&rel, dir, proj.as_ref())

                .into_iter()

                .collect()

        });

        PairWeights {

            inner: Arc::new(map),

        }

    }

    /// Single-source weighted shortest paths (Dijkstra) from `source` along `rel` in

    /// `direction`, with rel costs derived from a resident `weights` map (`PairWeights`,

    /// e.g. from `fold_via`). The cost of rel `(u, v)` is `1.0 / (weights[(u, v)] + base)`;

    /// a pair absent from `weights` is untraversable when `prune_missing` (else costs

    /// `1.0 / base`). Returns `{node_id: cost}` for every node reached (the source maps to

    /// `0.0`); pass `target` to stop once it is settled. The weight lookup runs inside the

    /// native kernel with the GIL released — no per-rel Python callback. E.g. BI Q19's

    /// interaction path: ``g.dijkstra(p1, Direction.Outgoing, "knows", weights=interaction,

    /// base=0.0, prune_missing=True)``.

    #[allow(clippy::too_many_arguments)]

    #[pyo3(signature = (source, direction, rel, *, weights, base=0.0, prune_missing=false, target=None))]

    fn dijkstra(

        &self,

        py: Python<'_>,

        source: u32,

        direction: Direction,

        rel: &str,

        weights: &PairWeights,

        base: f64,

        prune_missing: bool,

        target: Option<u32>,

    ) -> std::collections::HashMap<u32, f64> {

        let snapshot = self.snapshot.clone();

        let dir: rustychickpeas_core::Direction = direction.into();

        let rel = rel.to_owned();

        let map = weights.inner.clone();

        py.allow_threads(move || {

            let paths = snapshot.dijkstra(source, dir, rel.as_str(), target, |from, r| {

                let key = if from < r.neighbor {

                    (from, r.neighbor)

                } else {

                    (r.neighbor, from)

                };

                match map.get(&key) {

                    Some(&w) => 1.0 / (w as f64 + base),

                    None if prune_missing => f64::INFINITY,

                    None => 1.0 / base,

                }

            });

            paths.into_distances().into_iter().collect()

        })

    }

    /// Bulk-read a node's `rel_type` rels in `direction` as aligned arrays: the

    /// neighbor ids plus one Python list per requested property key (read from the

    /// rel column by CSR position, each column resolved once). The property-bearing

    /// bulk sibling of `neighbor_ids` — avoids the per-rel `Relationship` object and

    /// per-key `get_property` of `relationships()`. Returns `(neighbors, [values,

    /// ...])` aligned by index; a property absent on a rel is `None`.

    fn rels_with_props(

        &self,

        py: Python<'_>,

        node_id: u32,

        direction: Direction,

        rel_type: &str,

        prop_keys: Vec<String>,

    ) -> PyResult<PyObject> {

        use pyo3::types::PyList;

        use rustychickpeas_core::{BoolCol, ColumnDtype, Direction as CoreDir, F64Col, I64Col};

        enum H<'a> {

            I64(I64Col<'a>),

            F64(F64Col<'a>),

            Bool(BoolCol<'a>),

            None,

        }

        let snap = &self.snapshot;

        let dir: CoreDir = direction.into();

        let neighbors = PyList::empty(py);

        let cols: Vec<Bound<PyList>> = (0..prop_keys.len()).map(|_| PyList::empty(py)).collect();

        let rt = match snap.rel_type(rel_type) {

            Some(rt) if (node_id as usize) + 1 < snap.out_offsets.len() => rt,

            _ => return (neighbors, cols).into_py_any(py),

        };

        // Resolve each rel column to a typed reader once (not per rel).

        let hoisted: Vec<H> = prop_keys

            .iter()

            .map(|k| match snap.rel_col(k).map(|c| c.dtype()) {

                Some(ColumnDtype::I64) => H::I64(snap.rel_col(k).unwrap().i64()),

                Some(ColumnDtype::F64) => H::F64(snap.rel_col(k).unwrap().f64()),

                Some(ColumnDtype::Bool) => H::Bool(snap.rel_col(k).unwrap().bool()),

                _ => H::None,

            })

            .collect();

        let n = node_id as usize;

        let mut ranges: Vec<(usize, usize, bool)> = Vec::new();

        if matches!(dir, CoreDir::Outgoing | CoreDir::Both) {

            ranges.push((

                snap.out_offsets[n] as usize,

                snap.out_offsets[n + 1] as usize,

                true,

            ));

        }

        if matches!(dir, CoreDir::Incoming | CoreDir::Both) {

            ranges.push((

                snap.in_offsets[n] as usize,

                snap.in_offsets[n + 1] as usize,

                false,

            ));

        }

        for (s, e, outgoing) in ranges {

            for i in s..e {

                let (rtype, nbr, pos) = if outgoing {

                    (snap.out_types[i], snap.out_nbrs[i], i as u32)

                } else {

                    (

                        snap.in_types[i],

                        snap.in_nbrs[i],

                        snap.in_to_out.get(i).copied().unwrap_or(i as u32),

                    )

                };

                if rtype != rt {

                    continue;

                }

                neighbors.append(nbr)?;

                for (j, h) in hoisted.iter().enumerate() {

                    let v: PyObject = match h {

                        H::I64(c) => c.get(pos).into_py_any(py)?,

                        H::F64(c) => c.get(pos).into_py_any(py)?,

                        H::Bool(c) => c.get(pos).into_py_any(py)?,

                        H::None => py.None(),

                    };

                    cols[j].append(v)?;

                }

            }

        }

        (neighbors, cols).into_py_any(py)

    }

    /// Like `rels_with_props` but returns a `RelView` of zero-copy buffer arrays

    /// (`memoryview`-able, no per-value Python boxing): `.neighbors` (u32) and

    /// `.col(key)` (i64/f64). The native gather runs with the GIL released; a

    /// reduction like `sum(memoryview(v.col("amt")))` then runs at C speed. Missing

    /// props default to 0 (a typed buffer can't hold `None`).

    fn rel_view(

        &self,

        py: Python<'_>,

        node_id: u32,

        direction: Direction,

        rel_type: String,

        prop_keys: Vec<String>,

    ) -> RelView {

        use rustychickpeas_core::{ColumnDtype, Direction as CoreDir, F64Col, I64Col};

        enum KeyAcc<'a> {

            I64(I64Col<'a>, Vec<i64>),

            F64(F64Col<'a>, Vec<f64>),

            None(Vec<i64>),

        }

        let snapshot = self.snapshot.clone();

        let dir: CoreDir = direction.into();

        let n = node_id as usize;

        let (neighbors, cols): (Vec<u32>, Vec<(String, RelArrayData)>) =

            py.allow_threads(move || {

                let snap = &snapshot;

                let mut nbrs: Vec<u32> = Vec::new();

                let mut accs: Vec<KeyAcc> = prop_keys

                    .iter()

                    .map(|k| match snap.rel_col(k).map(|c| c.dtype()) {

                        Some(ColumnDtype::I64) => {

                            KeyAcc::I64(snap.rel_col(k).unwrap().i64(), Vec::new())

                        }

                        Some(ColumnDtype::F64) => {

                            KeyAcc::F64(snap.rel_col(k).unwrap().f64(), Vec::new())

                        }

                        _ => KeyAcc::None(Vec::new()),

                    })

                    .collect();

                let rt = snap.rel_type(&rel_type);

                if let (Some(rt), true) = (rt, n + 1 < snap.out_offsets.len()) {

                    let mut ranges: Vec<(usize, usize, bool)> = Vec::new();

                    if matches!(dir, CoreDir::Outgoing | CoreDir::Both) {

                        ranges.push((

                            snap.out_offsets[n] as usize,

                            snap.out_offsets[n + 1] as usize,

                            true,

                        ));

                    }

                    if matches!(dir, CoreDir::Incoming | CoreDir::Both) {

                        ranges.push((

                            snap.in_offsets[n] as usize,

                            snap.in_offsets[n + 1] as usize,

                            false,

                        ));

                    }

                    for (s, e, outgoing) in ranges {

                        for i in s..e {

                            let (rtype, nbr, pos) = if outgoing {

                                (snap.out_types[i], snap.out_nbrs[i], i as u32)

                            } else {

                                (

                                    snap.in_types[i],

                                    snap.in_nbrs[i],

                                    snap.in_to_out.get(i).copied().unwrap_or(i as u32),

                                )

                            };

                            if rtype != rt {

                                continue;

                            }

                            nbrs.push(nbr);

                            for acc in accs.iter_mut() {

                                match acc {

                                    KeyAcc::I64(c, v) => v.push(c.get(pos).unwrap_or(0)),

                                    KeyAcc::F64(c, v) => v.push(c.get(pos).unwrap_or(0.0)),

                                    KeyAcc::None(v) => v.push(0),

                                }

                            }

                        }

                    }

                }

                let cols = prop_keys

                    .into_iter()

                    .zip(accs)

                    .map(|(k, acc)| match acc {

                        KeyAcc::I64(_, v) => (k, RelArrayData::I64(v.into())),

                        KeyAcc::F64(_, v) => (k, RelArrayData::F64(v.into())),

                        KeyAcc::None(v) => (k, RelArrayData::I64(v.into())),

                    })

                    .collect();

                (nbrs, cols)

            });

        RelView {

            neighbors: neighbors.into(),

            cols,

        }

    }

    /// Build a `NeighborGroups` query over each source node's `rel` neighbors (in

    /// `direction`): group each source's neighbors by a projected attribute and

    /// reduce per source. Nothing runs until a terminal (`.sizes()` /

    /// `.top_by_size(...)`). E.g. BI Q4's biggest single-country membership per

    /// forum: ``g.neighbor_groups(forums, "hasMember", Direction.Outgoing)

    /// .project([(Direction.Outgoing, "isLocatedIn"), (Direction.Outgoing, "isPartOf")])

    /// .top_by_size(100, tie="flid")``.

    fn neighbor_groups(

        &self,

        sources: Vec<u32>,

        rel: String,

        direction: Direction,

    ) -> NeighborGroups {

        NeighborGroups {

            snapshot: self.snapshot.clone(),

            sources,

            rel,

            direction: direction.into(),

            project: Vec::new(),

        }

    }

    /// The string property `key` of `node_id`, or `None` when absent **or empty**

    /// (a dense string column stores a missing value as `""`).

    fn prop_str(&self, node_id: u32, key: &str) -> Option<String> {

        self.snapshot.prop_str(node_id, key).map(str::to_string)

    }

    /// All nodes within `min_hops..=max_hops` of `seed`, expanding only along

    /// `rel_type` in `direction` — the typed k-hop neighborhood as a list of ids

    /// (excludes `seed`). `min_hops` defaults to 1.

    #[pyo3(signature = (seed, direction, rel_type, max_hops, min_hops=1))]

    fn neighborhood(

        &self,

        seed: u32,

        direction: Direction,

        rel_type: &str,

        max_hops: u32,

        min_hops: u32,

    ) -> Vec<u32> {

        self.snapshot

            .neighborhood(seed, direction.into(), rel_type, min_hops..=max_hops)

            .iter()

            .collect()

    }

    /// The dense `i64` column `key` as a list (one value per node id), or `None`

    /// when the column is absent or not a dense `i64` column. Built on the

    /// zero-copy slice reader; the slice is copied to cross PyO3.

    fn i64_column(&self, key: &str) -> Option<Vec<i64>> {

        self.snapshot

            .col(key)

            .map(|c| c.i64())

            .and_then(|c| c.as_slice().map(<[i64]>::to_vec))

    }

    /// The interned id (code) for `s` in this snapshot, or `None` if `s` was never

    /// interned (so no node can carry it). Resolve filter targets to codes once,

    /// then compare them against a string [`Column`]'s codes (dtype `'string'`)

    /// vectorized — e.g. `numpy.isin(numpy.asarray(g.column("lang")), [c1, c2])`.

    fn string_id(&self, s: &str) -> Option<u32> {

        self.get_string_id(s)

    }

    /// A dense property column as a self-describing [`Column`] (its dtype is

    /// intrinsic — no `.i64()` narrowing), or `None` when the key is absent or the

    /// column is not stored densely. The `Column` supports the buffer protocol, so

    /// `numpy.asarray(col)` / `pyarrow.py_buffer(col)` / `memoryview(col)` read it

    /// zero-copy, and `col.to_pylist()` gives a plain Python list.

    fn column(&self, key: &str) -> Option<Column> {

        Column::build(self.snapshot.clone(), key)

    }

    /// Low-level grouped reduction over dense `i64` node columns, run in Rust with

    /// the GIL released. Prefer the fluent [`GraphSnapshot::aggregate`] builder —

    /// this is the kernel it calls, kept public for direct use.

    ///

    /// Scans the nodes of each label in `labels`. A row counts toward `total` when

    /// it passes every `pre_filters` predicate `(column, op, value)` (op ∈

    /// `<,<=,>,>=,==,!=`); rows additionally passing every `group_filters` predicate

    /// are grouped. The group key is the label index (when `group_label`), then each

    /// `group_cols` value, then a bucket index per `group_bins` `(column, bounds)`

    /// (bucket = count of `bounds <= value`). Returns `(rows, total)` with

    /// `rows = [(key_tuple, count, sum), ...]`. All referenced columns must be dense

    /// `i64` columns; a missing label or non-dense/-i64 column raises `ValueError`.

    #[pyo3(signature = (labels, pre_filters=vec![], group_filters=vec![], group_label=false, group_cols=vec![], group_bins=vec![], sum_col=None))]

    #[allow(clippy::too_many_arguments, clippy::type_complexity)]

    fn group_reduce(

        &self,

        py: Python<'_>,

        labels: Vec<String>,

        pre_filters: Vec<(String, String, i64)>,

        group_filters: Vec<(String, String, i64)>,

        group_label: bool,

        group_cols: Vec<String>,

        group_bins: Vec<(String, Vec<i64>)>,

        sum_col: Option<String>,

    ) -> PyResult<(Vec<(Vec<i64>, u64, i64)>, u64)> {

        let mut group: Vec<GroupSpec> = group_cols.into_iter().map(GroupSpec::Col).collect();

        group.extend(group_bins.into_iter().map(|(c, e)| GroupSpec::Bin(c, e)));

        let agg = build_core_agg(

            &self.snapshot,

            &labels,

            &pre_filters,

            &group_filters,

            group_label,

            &group,

            sum_col.as_deref(),

            None,

            None,

            &[],

        )?;

        let res = py

            .allow_threads(|| agg.run())

            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;

        let rows = res

            .rows

            .into_iter()

            .map(|r| (r.key, r.count, r.sum))

            .collect();

        Ok((rows, res.total))

    }

    /// Start a fluent aggregation over the given node labels — the Pythonic front

    /// for [`group_reduce`](Self::group_reduce). Chain `.where(col, op, value)` /

    /// `.having(col, op, value)` / `.by(col)` / `.bin(col, bounds)` / `.by_label()` /

    /// `.sum(col)`, then `.run()` for a result with `.total` and self-describing dict

    /// `.rows` (the source label comes back as its name). The heavy scan runs in Rust

    /// with the GIL released — no numpy/pyarrow needed.

    #[pyo3(signature = (*labels))]

    fn aggregate(&self, labels: Vec<String>) -> Aggregation {

        Aggregation {

            snapshot: self.snapshot.clone(),

            labels,

            where_filters: Vec::new(),

            having_filters: Vec::new(),

            by_label: false,

            group: Vec::new(),

            sum_col: None,

            through: None,

            neighbor_filter: None,

            projected_filters: Vec::new(),

        }

    }

    /// Weighted (or unweighted) shortest-path cost from `source` to `target`, or

    /// `None` if `target` is unreachable.

    ///

    /// With `weight_property`, each followed rel costs that f64/i64 rel

    /// property (an rel that lacks it is skipped); without it, every rel costs

    /// 1.0 (a hop count). `rel_types` optionally restricts which relationship

    /// types are followed (all types when `None`). Weights must be non-negative;

    /// the search is a bidirectional Dijkstra.

    #[pyo3(signature = (source, target, direction=Direction::Both, rel_types=None, weight_property=None))]

    fn shortest_path(

        &self,

        py: Python<'_>,

        source: u32,

        target: u32,

        direction: Direction,

        rel_types: Option<Vec<String>>,

        weight_property: Option<String>,

    ) -> PyResult<Option<f64>> {

        if source >= self.snapshot.n_nodes || target >= self.snapshot.n_nodes {

            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(

                "Node ID out of range (max: {})",

                self.snapshot.n_nodes.saturating_sub(1)

            )));

        }

        let types: Vec<&str> = rel_types

            .as_ref()

            .map(|t| t.iter().map(|s| s.as_str()).collect())

            .unwrap_or_default();

        let snapshot = self.snapshot.clone();

        // No Python is called during the search, so release the GIL.

        let cost = py.allow_threads(move || {

            let weight = |_from: u32, rel: &RelationshipRef| -> f64 {

                match &weight_property {

                    Some(prop) => match snapshot.rel_prop(rel.pos, prop).map(|p| p.value()) {

                        Some(ValueId::F64(bits)) => f64::from_bits(bits),

                        Some(ValueId::I64(w)) => w as f64,

                        _ => f64::INFINITY,

                    },

                    None => 1.0,

                }

            };

            snapshot.weighted_shortest_path(source, target, direction.into(), &types[..], weight)

        });

        Ok(cost)

    }

    /// Get relationships by type using the type_index bitmap for O(1) lookup

    /// Returns Relationship objects for all relationships of the specified type

    fn relationships_with_type(&self, rel_type: String) -> PyResult<Vec<Relationship>> {

        let rel_type_id = self.rel_type_from_str(&rel_type).ok_or_else(|| {

            PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(

                "Relationship type '{}' not found",

                rel_type

            ))

        })?;

        if let Some(bitmap) = self.snapshot.type_index.get(&rel_type_id) {

            let relationships: Vec<Relationship> = bitmap

                .iter()

                .map(|idx| Relationship {

                    snapshot: self.snapshot.clone(),

                    rel_index: idx,

                    is_outgoing: true,

                })

                .collect();

            Ok(relationships)

        } else {

            Ok(Vec::new())

        }

    }

    /// Get all nodes (returns all node IDs that have data: labels, rels, or properties)

    fn all_nodes(&self) -> PyResult<Vec<u32>> {

        use std::collections::HashSet;

        let mut nodes = HashSet::new();

        // Add nodes with labels

        for (_label, node_set) in &self.snapshot.label_index {

            for node_id in node_set.iter() {

                nodes.insert(node_id);

            }

        }

        // Add nodes with rels (check CSR arrays)

        // Nodes with outgoing rels

        for node_id in 0..self.snapshot.out_offsets.len().saturating_sub(1) {

            let start = self.snapshot.out_offsets[node_id] as usize;

            let end = self.snapshot.out_offsets[node_id + 1] as usize;

            if start < end {

                nodes.insert(node_id as u32);

            }

        }

        // Nodes with incoming rels

        for node_id in 0..self.snapshot.in_offsets.len().saturating_sub(1) {

            let start = self.snapshot.in_offsets[node_id] as usize;

            let end = self.snapshot.in_offsets[node_id + 1] as usize;

            if start < end {

                nodes.insert(node_id as u32);

            }

        }

        // Add nodes with properties

        for column in self.snapshot.columns.values() {

            match column {

                rustychickpeas_core::Column::DenseI64(_)

                | rustychickpeas_core::Column::DenseF64(_)

                | rustychickpeas_core::Column::DenseBool(_)

                | rustychickpeas_core::Column::DenseStr(_) => {

                    // Dense columns: all nodes from 0 to n_nodes-1 have this property

                    // But we only want nodes that actually have data, so skip dense columns

                    // (they're dense because most nodes have the property, but we can't tell which ones)

                }

                rustychickpeas_core::Column::SparseI64(pairs) => {

                    for (node_id, _) in pairs {

                        nodes.insert(*node_id);

                    }

                }

                rustychickpeas_core::Column::SparseF64(pairs) => {

                    for (node_id, _) in pairs {

                        nodes.insert(*node_id);

                    }

                }

                rustychickpeas_core::Column::SparseBool(pairs) => {

                    for (node_id, _) in pairs {

                        nodes.insert(*node_id);

                    }

                }

                rustychickpeas_core::Column::SparseStr(pairs) => {

                    for (node_id, _) in pairs {

                        nodes.insert(*node_id);

                    }

                }

                rustychickpeas_core::Column::RankI64 { present, .. }

                | rustychickpeas_core::Column::RankF64 { present, .. }

                | rustychickpeas_core::Column::RankBool { present, .. }

                | rustychickpeas_core::Column::RankStr { present, .. } => {

                    for pos in present.iter_ones() {

                        nodes.insert(pos as u32);

                    }

                }

            }

        }

        let mut result: Vec<u32> = nodes.into_iter().collect();

        result.sort_unstable();

        Ok(result)

    }

    /// Get all relationships as Relationship objects

    /// Returns all relationships in the graph