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::{Array, RecordBatch, make_array};
7use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer};
8use arrow_data::{ArrayData, ArrayDataBuilder, transform::MutableArrayData};
9
10pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer {
11    Buffer::from(buffer.as_slice())
12}
13
14pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
15    let nulls = nulls?;
16    let bit_buffer = deep_copy_buffer(nulls.inner().inner());
17    // SAFETY: `null_count` is taken from the source `NullBuffer`, which already
18    // upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over
19    // the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice`
20    // adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances
21    // the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s
22    // `as_slice()` view from byte 0) reproduces the exact bit pattern at the same
23    // bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new`
24    // panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the
25    // copy has the same length, so that check still passes.
26    Some(unsafe {
27        NullBuffer::new_unchecked(
28            BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()),
29            nulls.null_count(),
30        )
31    })
32}
33
34pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData {
35    let data_type = data.data_type().clone();
36    let len = data.len();
37    let nulls = deep_copy_nulls(data.nulls());
38    let offset = data.offset();
39    let buffers = data
40        .buffers()
41        .iter()
42        .map(deep_copy_buffer)
43        .collect::<Vec<_>>();
44    let child_data = data
45        .child_data()
46        .iter()
47        .map(deep_copy_array_data)
48        .collect::<Vec<_>>();
49    // SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract —
50    // `(data_type, len, offset, nulls, buffers, child_data)` must form a valid
51    // Arrow array. This call reproduces `data` structurally: `data_type`, `len`,
52    // and `offset` are forwarded unchanged; each buffer is replaced by a byte-
53    // identical copy of its offset-applied `as_slice()` view (the output buffer
54    // is `MutableBuffer`-allocated, at least as aligned as the source); `nulls`
55    // is deep-copied with the same bit offset/length and unset-bit count (see
56    // `deep_copy_nulls`); `child_data` is recursively cloned with the same
57    // guarantee. Every value-level invariant the source upheld — UTF-8 validity,
58    // monotonic offsets, in-bounds dictionary indices, run-end monotonicity,
59    // struct child-length matching — therefore transfers to the copy. If the
60    // source `ArrayData` was itself constructed via `new_unchecked` with an
61    // invalid payload, this function faithfully reproduces that invalidity.
62    unsafe {
63        ArrayDataBuilder::new(data_type)
64            .len(len)
65            .nulls(nulls)
66            .offset(offset)
67            .buffers(buffers)
68            .child_data(child_data)
69            .build_unchecked()
70    }
71}
72
73pub fn deep_copy_array(array: &dyn Array) -> Arc<dyn Array> {
74    let data = array.to_data();
75    let data = deep_copy_array_data(&data);
76    make_array(data)
77}
78
79pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result<RecordBatch> {
80    let arrays = batch
81        .columns()
82        .iter()
83        .map(|array| deep_copy_array(array))
84        .collect::<Vec<_>>();
85    RecordBatch::try_new(batch.schema(), arrays)
86}
87
88/// Deep copy array data, extracting only the sliced portion using MutableArrayData
89/// This is the most efficient and correct way to copy just the sliced data
90pub fn deep_copy_array_data_sliced(data: &ArrayData) -> ArrayData {
91    // Use MutableArrayData to efficiently copy just the slice
92    let mut mutable = MutableArrayData::new(vec![data], false, data.len());
93
94    // Copy from offset to offset+len (the visible slice)
95    mutable.extend(0, data.offset(), data.offset() + data.len());
96
97    // Freeze into immutable ArrayData
98    mutable.freeze()
99}
100
101/// Deep copy an array, extracting only the sliced portion using MutableArrayData
102pub fn deep_copy_array_sliced(array: &dyn Array) -> Arc<dyn Array> {
103    let data = array.to_data();
104    let data = deep_copy_array_data_sliced(&data);
105    make_array(data)
106}
107
108/// Deep copy a RecordBatch, extracting only the sliced portion using MutableArrayData
109pub fn deep_copy_batch_sliced(batch: &RecordBatch) -> crate::Result<RecordBatch> {
110    let arrays = batch
111        .columns()
112        .iter()
113        .map(|array| deep_copy_array_sliced(array))
114        .collect::<Vec<_>>();
115    RecordBatch::try_new(batch.schema(), arrays)
116}
117
118#[cfg(test)]
119mod tests {
120    use std::sync::Arc;
121
122    use arrow_array::{Array, Int32Array, RecordBatch, StringArray};
123    use arrow_schema::{DataType, Field, Schema};
124
125    #[test]
126    fn test_deep_copy_sliced_array_with_nulls() {
127        let array = Arc::new(Int32Array::from(vec![
128            Some(1),
129            None,
130            Some(3),
131            None,
132            Some(5),
133        ]));
134        let sliced_array = array.slice(1, 3);
135        let copied_array = super::deep_copy_array(&sliced_array);
136        assert_eq!(sliced_array.len(), copied_array.len());
137        assert_eq!(sliced_array.nulls(), copied_array.nulls());
138    }
139
140    #[test]
141    fn test_deep_copy_array_data_sliced() {
142        let array = Int32Array::from((0..1000).collect::<Vec<i32>>());
143        let sliced = array.slice(100, 10);
144
145        let sliced_data = sliced.to_data();
146        let copied_data = super::deep_copy_array_data_sliced(&sliced_data);
147
148        assert_eq!(copied_data.len(), 10);
149        assert_eq!(copied_data.offset(), 0);
150
151        // Verify data correctness
152        let copied_array = Int32Array::from(copied_data);
153        for i in 0..10 {
154            assert_eq!(copied_array.value(i), 100 + i as i32);
155        }
156    }
157
158    #[test]
159    fn test_deep_copy_array_sliced() {
160        let array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
161        let sliced = array.slice(1, 3);
162
163        let copied = super::deep_copy_array_sliced(&sliced);
164
165        assert_eq!(copied.len(), 3);
166        let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
167        assert_eq!(copied_int.value(0), 2);
168        assert_eq!(copied_int.value(1), 3);
169        assert_eq!(copied_int.value(2), 4);
170    }
171
172    #[test]
173    fn test_deep_copy_batch_sliced() {
174        let schema = Arc::new(Schema::new(vec![
175            Field::new("id", DataType::Int32, false),
176            Field::new("name", DataType::Utf8, false),
177        ]));
178
179        let id_array = Arc::new(Int32Array::from((0..100).collect::<Vec<i32>>()));
180        let name_array = Arc::new(StringArray::from(
181            (0..100)
182                .map(|i| format!("name_{}", i))
183                .collect::<Vec<String>>(),
184        ));
185
186        let batch = RecordBatch::try_new(
187            schema,
188            vec![id_array as Arc<dyn Array>, name_array as Arc<dyn Array>],
189        )
190        .unwrap();
191
192        let sliced = batch.slice(10, 5);
193        let copied = super::deep_copy_batch_sliced(&sliced).unwrap();
194
195        assert_eq!(copied.num_rows(), 5);
196        assert_eq!(copied.num_columns(), 2);
197
198        // Verify data correctness
199        let id_col = copied
200            .column(0)
201            .as_any()
202            .downcast_ref::<Int32Array>()
203            .unwrap();
204        let name_col = copied
205            .column(1)
206            .as_any()
207            .downcast_ref::<StringArray>()
208            .unwrap();
209
210        for i in 0..5 {
211            assert_eq!(id_col.value(i), 10 + i as i32);
212            assert_eq!(name_col.value(i), format!("name_{}", 10 + i));
213        }
214    }
215
216    #[test]
217    fn test_deep_copy_array_sliced_with_nulls() {
218        let array = Arc::new(Int32Array::from(vec![
219            Some(1),
220            None,
221            Some(3),
222            None,
223            Some(5),
224        ]));
225        let sliced = array.slice(1, 3); // [None, Some(3), None]
226
227        let copied = super::deep_copy_array_sliced(&sliced);
228
229        assert_eq!(copied.len(), 3);
230        assert_eq!(copied.null_count(), 2); // Two nulls in the slice
231
232        let copied_int = copied.as_any().downcast_ref::<Int32Array>().unwrap();
233        assert!(!copied_int.is_valid(0)); // None
234        assert!(copied_int.is_valid(1)); // Some(3)
235        assert!(!copied_int.is_valid(2)); // None
236        assert_eq!(copied_int.value(1), 3);
237    }
238}
lance_arrow/deepcopy.rs

lance_arrow/
deepcopy.rs
