lance_arrow/
deepcopy.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_array::{make_array, Array, RecordBatch};
7use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
8use arrow_data::{ArrayData, ArrayDataBuilder};
9
10pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer {
11    Buffer::from(buffer.as_slice())
12}
13
14fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
15    let nulls = nulls?;
16    let bit_buffer = deep_copy_buffer(nulls.inner().inner());
17    Some(unsafe {
18        NullBuffer::new_unchecked(
19            BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()),
20            nulls.null_count(),
21        )
22    })
23}
24
25pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData {
26    let data_type = data.data_type().clone();
27    let len = data.len();
28    let nulls = deep_copy_nulls(data.nulls());
29    let offset = data.offset();
30    let buffers = data
31        .buffers()
32        .iter()
33        .map(deep_copy_buffer)
34        .collect::<Vec<_>>();
35    let child_data = data
36        .child_data()
37        .iter()
38        .map(deep_copy_array_data)
39        .collect::<Vec<_>>();
40    unsafe {
41        ArrayDataBuilder::new(data_type)
42            .len(len)
43            .nulls(nulls)
44            .offset(offset)
45            .buffers(buffers)
46            .child_data(child_data)
47            .build_unchecked()
48    }
49}
50
51pub fn deep_copy_array(array: &dyn Array) -> Arc<dyn Array> {
52    let data = array.to_data();
53    let data = deep_copy_array_data(&data);
54    make_array(data)
55}
56
57pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result<RecordBatch> {
58    let arrays = batch
59        .columns()
60        .iter()
61        .map(|array| deep_copy_array(array))
62        .collect::<Vec<_>>();
63    RecordBatch::try_new(batch.schema(), arrays)
64}
65
66#[cfg(test)]
67pub mod tests {
68    use std::sync::Arc;
69
70    use arrow_array::{Array, Int32Array};
71
72    #[test]
73    fn test_deep_copy_sliced_array_with_nulls() {
74        let array = Arc::new(Int32Array::from(vec![
75            Some(1),
76            None,
77            Some(3),
78            None,
79            Some(5),
80        ]));
81        let sliced_array = array.slice(1, 3);
82        let copied_array = super::deep_copy_array(&sliced_array);
83        assert_eq!(sliced_array.len(), copied_array.len());
84        assert_eq!(sliced_array.nulls(), copied_array.nulls());
85    }
86}
lance_arrow/deepcopy.rs

lance_arrow/
deepcopy.rs
