如何在Rust中按指定列对arrow2 Chunk进行排序？

要实现按第一列对arrow2的Chunk排序，推荐使用compute::sort模块下的sort_by_column函数，它可以直接指定列索引完成排序，适合大型Chunk场景。以下是修正后的完整代码：

fn main() {
    use arrow2::{
        array::{Int64Array, Int64Vec},
        chunk::Chunk,
        compute::sort::{sort_by_column, SortOptions},
    };

    // 构建测试列数据
    let mut col1: Int64Vec = Int64Vec::new();
    col1.push(Some(0));
    col1.push(Some(5));
    col1.push(Some(3));
    col1.push(Some(2));

    let mut col2: Int64Vec = Int64Vec::new();
    col2.push(Some(1));
    col2.push(Some(2));
    col2.push(Some(3));
    col2.push(Some(4));

    let chunk = Chunk::new(vec![col1.into_arc(), col2.into_arc()]);

    // 按第一列（索引0）升序排序，SortOptions可配置排序规则
    let sorted_chunk = sort_by_column(
        &chunk,
        &0,
        &SortOptions::default(),
    ).unwrap();

    // 验证排序结果
    let sorted_col1 = sorted_chunk.arrays()[0].as_any().downcast_ref::<Int64Array>().unwrap();
    let sorted_col2 = sorted_chunk.arrays()[1].as_any().downcast_ref::<Int64Array>().unwrap();
    
    println!("Sorted col1: {:?}", sorted_col1.values());
    println!("Sorted col2: {:?}", sorted_col2.values());
}

关键说明 #

• sort_by_column是arrow2专为多列Chunk设计的排序API，参数包括：
  • 待排序的Chunk引用
  • 目标列的索引（第一列对应索引0）
  • SortOptions：可通过SortOptions { descending: true, nulls_first: false }配置降序、空值排序规则
• 若需手动控制排序索引，可使用indices_sorted_by_key生成排序索引后，通过chunk.take(&indices)重排Chunk，示例：

use arrow2::compute::sort::indices_sorted_by_key;

// 获取第一列作为排序键
let sort_key = chunk.arrays()[0].as_any().downcast_ref::<Int64Array>().unwrap();
// 生成排序索引
let indices = indices_sorted_by_key(sort_key, &SortOptions::default()).unwrap();
// 重排Chunk
let sorted_chunk = chunk.take(&indices).unwrap();

内容的提问来源于stack exchange，提问作者André Fontenele