arrow_select::take

Function take_record_batch

Source
pub fn take_record_batch(
    record_batch: &RecordBatch,
    indices: &dyn Array,
) -> Result<RecordBatch, ArrowError>
Expand description

Take rows by index from RecordBatch and returns a new RecordBatch from those indexes.

This function will call take on each array of the RecordBatch and assemble a new RecordBatch.

ยงExample


let schema = Arc::new(Schema::new(vec![
    Field::new("a", DataType::Int32, true),
    Field::new("b", DataType::Utf8, true),
]));
let batch = RecordBatch::try_new(
    schema.clone(),
    vec![
        Arc::new(Int32Array::from_iter_values(0..20)),
        Arc::new(StringArray::from_iter_values(
            (0..20).map(|i| format!("str-{}", i)),
        )),
    ],
)
.unwrap();

let indices = UInt32Array::from(vec![1, 5, 10]);
let taken = take_record_batch(&batch, &indices).unwrap();

let expected = RecordBatch::try_new(
    schema,
    vec![
        Arc::new(Int32Array::from(vec![1, 5, 10])),
        Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
    ],
)
.unwrap();
assert_eq!(taken, expected);