如何使用arrow-rs比较两个Apache Arrow RecordBatch？

核心思路 #

因为两个向量的批次划分可能不同，所以首先要把所有批次合并为完整的数据集，消除批次差异后再做Schema和内容校验。同时要实现类似assert_eq!的错误提示，明确指出失败原因。

实现步骤 #

1. 合并批次到单个RecordBatch #

先写一个辅助函数把分散的批次合并，这样批次划分的差异就不会影响后续比较：

use arrow::record_batch::RecordBatch;
use arrow::array::ArrayRef;
use arrow::error::ArrowError;

fn merge_batches(batches: &[RecordBatch]) -> Result<RecordBatch, ArrowError> {
    if batches.is_empty() {
        return Err(ArrowError::InvalidArgumentError("Empty batch list".to_string()));
    }

    let schema = batches[0].schema().clone();
    let num_columns = schema.fields().len();

    let mut merged_arrays = Vec::with_capacity(num_columns);
    for col_idx in 0..num_columns {
        let arrays: Vec<ArrayRef> = batches.iter().map(|batch| batch.column(col_idx).clone()).collect();
        let merged_array = arrow::compute::concat(&arrays)?;
        merged_arrays.push(merged_array);
    }

    RecordBatch::try_new(schema, merged_arrays)
}

2. 校验Schema并输出差异 #

对比合并后的Schema，不一致时明确列出两边的字段信息：

fn compare_schemas(a: &RecordBatch, b: &RecordBatch) -> Result<(), String> {
    if a.schema() != b.schema() {
        let a_fields = a.schema().fields().iter()
            .map(|f| format!("{}: {}", f.name(), f.data_type()))
            .collect::<Vec<_>>().join("\n");
        let b_fields = b.schema().fields().iter()
            .map(|f| format!("{}: {}", f.name(), f.data_type()))
            .collect::<Vec<_>>().join("\n");
        return Err(format!(
            "Schema mismatch:\nLeft schema:\n{}\nRight schema:\n{}",
            a_fields, b_fields
        ));
    }
    Ok(())
}

3. 逐列校验内容 #

对比数据内容，先检查行数是否一致，再逐列用Arrow内置的equals方法校验，出错时指明具体列：

fn compare_content(a: &RecordBatch, b: &RecordBatch) -> Result<(), String> {
    if a.num_rows() != b.num_rows() {
        return Err(format!(
            "Row count mismatch: left has {} rows, right has {} rows",
            a.num_rows(), b.num_rows()
        ));
    }

    for (col_idx, field) in a.schema().fields().iter().enumerate() {
        let a_col = a.column(col_idx);
        let b_col = b.column(col_idx);
        if !a_col.equals(b_col) {
            return Err(format!(
                "Column '{}' (index {}) content mismatch",
                field.name(), col_idx
            ));
        }
    }
    Ok(())
}

4. 封装成断言宏 #

把上述步骤打包成宏，用起来和assert_eq!一样方便：

#[macro_export]
macro_rules! assert_record_batches_eq {
    ($a:expr, $b:expr) => {
        match (merge_batches($a), merge_batches($b)) {
            (Ok(merged_a), Ok(merged_b)) => {
                if let Err(e) = compare_schemas(&merged_a, &merged_b) {
                    panic!("{}", e);
                }
                if let Err(e) = compare_content(&merged_a, &merged_b) {
                    panic!("{}", e);
                }
            }
            (Err(e), _) => panic!("Failed to merge left batches: {}", e),
            (_, Err(e)) => panic!("Failed to merge right batches: {}", e),
        }
    };
}

使用示例 #

// 假设有两个Vec<RecordBatch> batch1和batch2
assert_record_batches_eq!(&batch1, &batch2);

一旦比较失败，会直接panic并输出具体错误，比如Schema字段不匹配、行数差异或者某列内容不一致。

注意事项 #

• 超大批次合并会占用较多内存，若需优化可改为逐列分批比较，不过实现会更复杂。
• Arrow的equals方法会正确处理nullable字段的NULL值比较。

内容的提问来源于stack exchange，提问作者aljazerzen