nordisk/knut

crates/graphar/src/merge.rs

//! Merging row-aligned property-group batches into a single `RecordBatch`.
//!
//! GraphAr stores each property group in its own chunk files, but all groups
//! of a vertex (or edge) type are row-aligned: row `i` of every group refers
//! to the same vertex/edge. Merging is therefore a horizontal column zip.

use std::sync::Arc;

use arrow::compute::concat_batches;
use arrow_array::RecordBatch;
use arrow_schema::{Field, Schema};

use crate::error::{GraphArError, Result};

/// Merge per-group batch streams into a single batch containing the columns
/// of every group, in the order the groups were given.
///
/// Each element of `groups` is the full batch list for one property group.
/// All groups must cover the same number of rows. Column names must be
/// unique across groups (GraphAr forbids duplicate property names).
pub fn merge_property_groups(groups: &[Vec<RecordBatch>]) -> Result<RecordBatch> {
    let mut flattened = Vec::new();
    for batches in groups {
        if batches.is_empty() {
            continue;
        }
        let schema = batches[0].schema();
        flattened.push(concat_batches(&schema, batches)?);
    }
    if flattened.is_empty() {
        return Err(GraphArError::Other(
            "merge_property_groups: no batches to merge".into(),
        ));
    }

    let row_count = flattened[0].num_rows();
    let mut fields: Vec<Field> = Vec::new();
    let mut columns = Vec::new();
    for batch in &flattened {
        if batch.num_rows() != row_count {
            return Err(GraphArError::Other(format!(
                "merge_property_groups: row count mismatch ({} vs {})",
                batch.num_rows(),
                row_count
            )));
        }
        for (field, column) in batch.schema().fields().iter().zip(batch.columns()) {
            if fields.iter().any(|f: &Field| f.name() == field.name()) {
                return Err(GraphArError::Other(format!(
                    "merge_property_groups: duplicate column '{}'",
                    field.name()
                )));
            }
            fields.push(field.as_ref().clone());
            columns.push(column.clone());
        }
    }

    Ok(RecordBatch::try_new(
        Arc::new(Schema::new(fields)),
        columns,
    )?)
}