nordisk/nornir

crates/mimir-graph/src/arrow.rs

//! CSR straight from Arrow edge batches (plan §11: sorted scan of
//! `code_edges` + cumsum over `from` → `GrB_Matrix`).
//!
//! Input: one or more `RecordBatch`es with two non-null `UInt64` columns
//! holding node ordinals `0..n_nodes`. Three paths, picked by one scan:
//!
//! - **sorted by (from, to), no duplicates** (what a sorted `code_edges` scan
//!   yields): offsets by counting + cumsum written into a malloc'd array, the
//!   `to` buffer copied once, `GxB_Matrix_pack_CSR` (no further copy).
//! - **sorted by from only**: same, rows marked jumbled; a duplicate check
//!   runs and, if any `(from, to)` repeats, falls through to the build path.
//! - **unsorted**: `GxB_Matrix_build_Scalar` (GraphBLAS' parallel sort,
//!   duplicates collapse). With one batch it reads the Arrow buffers in place.

use arrow_array::{Array, RecordBatch, UInt64Array};
use arrow_schema::DataType;

use crate::csr::{MallocBuf, build_iso, pack, validate};
use crate::error::{Error, Result};
use crate::matrix::Matrix;

/// Which import path [`from_arrow_edge_batches`] took (for benches and logs).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImportPath {
    /// Sorted by (from, to), packed without sorting.
    SortedPack,
    /// Sorted by from, rows jumbled, packed (GraphBLAS sorts rows lazily).
    FromSortedPack,
    /// Unsorted or duplicated: built by GraphBLAS.
    Build,
}

fn u64_column<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a [u64]> {
    let col = batch.column_by_name(name).ok_or_else(|| {
        Error::Arrow(format!(
            "no column `{name}` in batch {:?}",
            batch
                .schema()
                .fields()
                .iter()
                .map(|f| f.name())
                .collect::<Vec<_>>()
        ))
    })?;
    if col.data_type() != &DataType::UInt64 {
        return Err(Error::Arrow(format!(
            "column `{name}` is {}, want UInt64",
            col.data_type()
        )));
    }
    if col.null_count() > 0 {
        return Err(Error::Arrow(format!(
            "column `{name}` has {} nulls",
            col.null_count()
        )));
    }
    let a = col
        .as_any()
        .downcast_ref::<UInt64Array>()
        .ok_or_else(|| Error::Arrow(format!("column `{name}` is not a UInt64Array")))?;
    Ok(a.values().as_ref())
}

/// `n_nodes × n_nodes` boolean adjacency from one edge batch.
pub fn from_arrow_edges(
    batch: &RecordBatch,
    from_col: &str,
    to_col: &str,
    n_nodes: u64,
) -> Result<Matrix<bool>> {
    Ok(from_arrow_edge_batches(std::slice::from_ref(batch), from_col, to_col, n_nodes)?.0)
}

/// Same over a scan's batches (treated as one concatenated edge list); also
/// returns which path was taken.
pub fn from_arrow_edge_batches(
    batches: &[RecordBatch],
    from_col: &str,
    to_col: &str,
    n_nodes: u64,
) -> Result<(Matrix<bool>, ImportPath)> {
    let mut chunks = Vec::with_capacity(batches.len());
    for b in batches {
        chunks.push((u64_column(b, from_col)?, u64_column(b, to_col)?));
    }
    from_edge_chunks(&chunks, n_nodes)
}

/// The engine behind the Arrow entry points, over plain `(from, to)` slices.
pub fn from_edge_chunks(
    chunks: &[(&[u64], &[u64])],
    n_nodes: u64,
) -> Result<(Matrix<bool>, ImportPath)> {
    let mut m = 0usize;
    for (f, t) in chunks {
        if f.len() != t.len() {
            return Err(Error::Arrow(format!(
                "from has {} rows, to has {}",
                f.len(),
                t.len()
            )));
        }
        m += f.len();
    }

    // One scan: bounds, sorted by from, sorted by (from, to) strictly.
    let mut from_sorted = true;
    let mut pair_sorted = true;
    let mut prev: Option<(u64, u64)> = None;
    for (f, t) in chunks {
        for (&a, &b) in f.iter().zip(t.iter()) {
            if a >= n_nodes || b >= n_nodes {
                return Err(Error::invalid(format!(
                    "edge ({a}, {b}) out of range for {n_nodes} nodes"
                )));
            }
            if let Some((pa, pb)) = prev {
                if a < pa {
                    from_sorted = false;
                    pair_sorted = false;
                } else if a == pa && b <= pb {
                    pair_sorted = false;
                }
            }
            prev = Some((a, b));
        }
        if !from_sorted {
            break;
        }
    }

    if from_sorted {
        // Counting + cumsum straight into the array GraphBLAS will own.
        let n = n_nodes as usize;
        let mut offsets = MallocBuf::<u64>::zeroed(n + 1)?;
        {
            let o = offsets.as_mut_slice();
            for (f, _) in chunks {
                for &a in f.iter() {
                    o[a as usize + 1] += 1;
                }
            }
            for i in 0..n {
                o[i + 1] += o[i];
            }
        }
        let mut cols = MallocBuf::<u64>::new(m)?;
        {
            let c = cols.as_mut_slice();
            let mut at = 0;
            for (_, t) in chunks {
                c[at..at + t.len()].copy_from_slice(t);
                at += t.len();
            }
        }
        if pair_sorted {
            let mat = pack::<bool>(n_nodes, n_nodes, offsets, cols, None, true)?;
            // a pack pins CSR: let GraphBLAS go hypersparse when rows are mostly empty
            mat.conform()?;
            return Ok((mat, ImportPath::SortedPack));
        }
        // Rows may be jumbled; duplicates are not allowed in a pack.
        if validate(n_nodes, n_nodes, offsets.as_slice(), cols.as_slice()).is_ok() {
            let mat = pack::<bool>(n_nodes, n_nodes, offsets, cols, None, false)?;
            mat.conform()?;
            return Ok((mat, ImportPath::FromSortedPack));
        }
    }

    let mat = Matrix::<bool>::new(n_nodes, n_nodes)?;
    if let [(f, t)] = chunks {
        build_iso(&mat, f.as_ptr(), t.as_ptr(), f.len() as u64)?;
    } else {
        let f: Vec<u64> = chunks.iter().flat_map(|c| c.0.iter().copied()).collect();
        let t: Vec<u64> = chunks.iter().flat_map(|c| c.1.iter().copied()).collect();
        build_iso(&mat, f.as_ptr(), t.as_ptr(), f.len() as u64)?;
    }
    Ok((mat, ImportPath::Build))
}