lance_encoding/previous/encodings/physical/
fixed_size_binary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_buffer::ScalarBuffer;
7use arrow_schema::DataType;
8use futures::{future::BoxFuture, FutureExt};
9use lance_core::Result;
10
11use crate::{
12    buffer::LanceBuffer,
13    data::{BlockInfo, DataBlock, FixedWidthDataBlock, VariableWidthBlock},
14    decoder::{PageScheduler, PrimitivePageDecoder},
15    format::ProtobufUtils,
16    previous::encoder::{ArrayEncoder, EncodedArray},
17    EncodingsIo,
18};
19
20/// A scheduler for fixed size binary data
21#[derive(Debug)]
22pub struct FixedSizeBinaryPageScheduler {
23    bytes_scheduler: Box<dyn PageScheduler>,
24    byte_width: u32,
25    bytes_per_offset: u32,
26}
27
28impl FixedSizeBinaryPageScheduler {
29    pub fn new(
30        bytes_scheduler: Box<dyn PageScheduler>,
31        byte_width: u32,
32        bytes_per_offset: u32,
33    ) -> Self {
34        Self {
35            bytes_scheduler,
36            byte_width,
37            bytes_per_offset,
38        }
39    }
40}
41
42impl PageScheduler for FixedSizeBinaryPageScheduler {
43    fn schedule_ranges(
44        &self,
45        ranges: &[std::ops::Range<u64>],
46        scheduler: &Arc<dyn EncodingsIo>,
47        top_level_row: u64,
48    ) -> BoxFuture<'static, Result<Box<dyn PrimitivePageDecoder>>> {
49        let expanded_ranges = ranges
50            .iter()
51            .map(|range| {
52                (range.start * self.byte_width as u64)..(range.end * self.byte_width as u64)
53            })
54            .collect::<Vec<_>>();
55
56        let bytes_page_decoder =
57            self.bytes_scheduler
58                .schedule_ranges(&expanded_ranges, scheduler, top_level_row);
59
60        let byte_width = self.byte_width as u64;
61        let bytes_per_offset = self.bytes_per_offset;
62
63        async move {
64            let bytes_decoder = bytes_page_decoder.await?;
65            Ok(Box::new(FixedSizeBinaryDecoder {
66                bytes_decoder,
67                byte_width,
68                bytes_per_offset,
69            }) as Box<dyn PrimitivePageDecoder>)
70        }
71        .boxed()
72    }
73}
74
75pub struct FixedSizeBinaryDecoder {
76    bytes_decoder: Box<dyn PrimitivePageDecoder>,
77    byte_width: u64,
78    bytes_per_offset: u32,
79}
80
81impl PrimitivePageDecoder for FixedSizeBinaryDecoder {
82    fn decode(&self, rows_to_skip: u64, num_rows: u64) -> Result<DataBlock> {
83        let rows_to_skip = rows_to_skip * self.byte_width;
84        let num_bytes = num_rows * self.byte_width;
85        let bytes = self.bytes_decoder.decode(rows_to_skip, num_bytes)?;
86        let bytes = bytes.as_fixed_width().unwrap();
87        debug_assert_eq!(bytes.bits_per_value, self.byte_width * 8);
88
89        let offsets_buffer = match self.bytes_per_offset {
90            8 => {
91                let offsets_vec = (0..(num_rows + 1))
92                    .map(|i| i * self.byte_width)
93                    .collect::<Vec<_>>();
94
95                ScalarBuffer::from(offsets_vec).into_inner()
96            }
97            4 => {
98                let offsets_vec = (0..(num_rows as u32 + 1))
99                    .map(|i| i * self.byte_width as u32)
100                    .collect::<Vec<_>>();
101
102                ScalarBuffer::from(offsets_vec).into_inner()
103            }
104            _ => panic!("Unsupported offsets type"),
105        };
106
107        let string_data = DataBlock::VariableWidth(VariableWidthBlock {
108            bits_per_offset: (self.bytes_per_offset * 8) as u8,
109            data: bytes.data,
110            num_values: num_rows,
111            offsets: LanceBuffer::from(offsets_buffer),
112            block_info: BlockInfo::new(),
113        });
114
115        Ok(string_data)
116    }
117}
118
119#[derive(Debug)]
120pub struct FixedSizeBinaryEncoder {
121    bytes_encoder: Box<dyn ArrayEncoder>,
122    byte_width: usize,
123}
124
125impl FixedSizeBinaryEncoder {
126    pub fn new(bytes_encoder: Box<dyn ArrayEncoder>, byte_width: usize) -> Self {
127        Self {
128            bytes_encoder,
129            byte_width,
130        }
131    }
132}
133
134impl ArrayEncoder for FixedSizeBinaryEncoder {
135    fn encode(
136        &self,
137        data: DataBlock,
138        _data_type: &DataType,
139        buffer_index: &mut u32,
140    ) -> Result<EncodedArray> {
141        let bytes_data = data.as_variable_width().unwrap();
142        let fixed_data = DataBlock::FixedWidth(FixedWidthDataBlock {
143            bits_per_value: 8 * self.byte_width as u64,
144            data: bytes_data.data,
145            num_values: bytes_data.num_values,
146            block_info: BlockInfo::new(),
147        });
148
149        let encoded_data = self.bytes_encoder.encode(
150            fixed_data,
151            &DataType::FixedSizeBinary(self.byte_width as i32),
152            buffer_index,
153        )?;
154        let encoding =
155            ProtobufUtils::fixed_size_binary(encoded_data.encoding, self.byte_width as u32);
156
157        Ok(EncodedArray {
158            data: encoded_data.data,
159            encoding,
160        })
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use std::{collections::HashMap, sync::Arc};
167
168    use arrow_array::{
169        builder::LargeStringBuilder, Array, ArrayRef, FixedSizeBinaryArray, LargeStringArray,
170        StringArray,
171    };
172    use arrow_buffer::Buffer;
173    use arrow_data::ArrayData;
174    use arrow_schema::{DataType, Field};
175
176    use crate::data::{DataBlock, FixedWidthDataBlock};
177    use crate::decoder::PrimitivePageDecoder;
178    use crate::previous::encodings::physical::fixed_size_binary::FixedSizeBinaryDecoder;
179    use crate::testing::{check_basic_random, check_round_trip_encoding_of_data, TestCases};
180
181    #[test_log::test(tokio::test)]
182    async fn test_fixed_size_utf8_binary() {
183        let field = Field::new("", DataType::Utf8, false);
184        // This test only generates fixed size binary arrays anyway
185        check_basic_random(field).await;
186    }
187
188    #[test_log::test(tokio::test)]
189    async fn test_fixed_size_binary() {
190        let field = Field::new("", DataType::Binary, false);
191        check_basic_random(field).await;
192    }
193
194    #[test_log::test(tokio::test)]
195    async fn test_fixed_size_large_binary() {
196        let field = Field::new("", DataType::LargeBinary, true);
197        check_basic_random(field).await;
198    }
199
200    #[test_log::test(tokio::test)]
201    async fn test_fixed_size_large_utf8() {
202        let field = Field::new("", DataType::LargeUtf8, true);
203        check_basic_random(field).await;
204    }
205
206    #[test_log::test(tokio::test)]
207    async fn test_simple_fixed_size_utf8() {
208        let string_array = StringArray::from(vec![
209            Some("abc"),
210            Some("def"),
211            Some("ghi"),
212            Some("jkl"),
213            Some("mno"),
214        ]);
215
216        let test_cases = TestCases::default()
217            .with_range(0..2)
218            .with_range(0..3)
219            .with_range(1..3)
220            .with_indices(vec![0, 1, 3, 4]);
221
222        check_round_trip_encoding_of_data(
223            vec![Arc::new(string_array)],
224            &test_cases,
225            HashMap::new(),
226        )
227        .await;
228    }
229
230    #[test_log::test(tokio::test)]
231    async fn test_simple_fixed_size_with_nulls_utf8() {
232        let string_array =
233            LargeStringArray::from(vec![Some("abc"), None, Some("ghi"), None, Some("mno")]);
234
235        let test_cases = TestCases::default()
236            .with_range(0..2)
237            .with_range(0..3)
238            .with_range(1..3)
239            .with_indices(vec![0, 1, 3, 4]);
240
241        check_round_trip_encoding_of_data(
242            vec![Arc::new(string_array)],
243            &test_cases,
244            HashMap::new(),
245        )
246        .await;
247    }
248
249    #[test_log::test(tokio::test)]
250    async fn test_fixed_size_sliced_utf8() {
251        let string_array = StringArray::from(vec![Some("abc"), Some("def"), None, Some("fgh")]);
252        let string_array = string_array.slice(1, 3);
253
254        let test_cases = TestCases::default()
255            .with_range(0..1)
256            .with_range(0..2)
257            .with_range(1..2);
258        check_round_trip_encoding_of_data(
259            vec![Arc::new(string_array)],
260            &test_cases,
261            HashMap::new(),
262        )
263        .await;
264    }
265
266    #[test_log::test(tokio::test)]
267    async fn test_fixed_size_empty_strings() {
268        // All strings are empty
269
270        // When encoding an array of empty strings there are no bytes to encode
271        // which is strange and we want to ensure we handle it
272        let string_array = Arc::new(StringArray::from(vec![Some(""), None, Some("")]));
273
274        let test_cases = TestCases::default().with_range(0..2).with_indices(vec![1]);
275        check_round_trip_encoding_of_data(vec![string_array.clone()], &test_cases, HashMap::new())
276            .await;
277        let test_cases = test_cases.with_batch_size(1);
278        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
279    }
280
281    #[test_log::test(tokio::test)]
282    #[ignore] // This test is quite slow in debug mode
283    async fn test_jumbo_string() {
284        // This is an overflow test.  We have a list of lists where each list
285        // has 1Mi items.  We encode 5000 of these lists and so we have over 4Gi in the
286        // offsets range
287        let mut string_builder = LargeStringBuilder::new();
288        // a 1 MiB string
289        let giant_string = String::from_iter((0..(1024 * 1024)).map(|_| '0'));
290        for _ in 0..5000 {
291            string_builder.append_option(Some(&giant_string));
292        }
293        let giant_array = Arc::new(string_builder.finish()) as ArrayRef;
294        let arrs = vec![giant_array];
295
296        // // We can't validate because our validation relies on concatenating all input arrays
297        let test_cases = TestCases::default().without_validation();
298        check_round_trip_encoding_of_data(arrs, &test_cases, HashMap::new()).await;
299    }
300
301    struct FixedWidthCloningPageDecoder {
302        data_block: FixedWidthDataBlock,
303    }
304
305    impl PrimitivePageDecoder for FixedWidthCloningPageDecoder {
306        // clone the given data block as decoded data block
307        fn decode(
308            &self,
309            _rows_to_skip: u64,
310            _num_rows: u64,
311        ) -> lance_core::error::Result<DataBlock> {
312            Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
313                data: self.data_block.data.deep_copy(),
314                bits_per_value: self.data_block.bits_per_value,
315                num_values: self.data_block.num_values,
316                block_info: self.data_block.block_info.clone(),
317            }))
318        }
319    }
320
321    #[test]
322    fn test_fixed_size_binary_decoder() {
323        let values: [u8; 6] = *b"aaabbb";
324        let num_values = 2u64;
325        let byte_width = 3;
326        let array_data = ArrayData::builder(DataType::FixedSizeBinary(byte_width))
327            .len(num_values as usize)
328            .add_buffer(Buffer::from(&values[..]))
329            .build()
330            .unwrap();
331        let fixed_size_binary_array = FixedSizeBinaryArray::from(array_data);
332        let arrays = vec![Arc::new(fixed_size_binary_array) as ArrayRef];
333        let fixed_width_data_block = DataBlock::from_arrays(&arrays, num_values);
334        assert_eq!(fixed_width_data_block.name(), "FixedWidth");
335
336        let bytes_decoder = FixedWidthCloningPageDecoder {
337            data_block: fixed_width_data_block.as_fixed_width().unwrap(),
338        };
339        let decoder = FixedSizeBinaryDecoder {
340            bytes_decoder: Box::new(bytes_decoder),
341            byte_width: byte_width as u64,
342            bytes_per_offset: 4, // 32-bits offset binary
343        };
344
345        let decoded_binary = decoder.decode(0, num_values).unwrap();
346        let maybe_data = decoded_binary.into_arrow(DataType::Utf8, true);
347        assert!(maybe_data.is_ok());
348        let data = maybe_data.unwrap();
349        let string_array = StringArray::from(data);
350        assert_eq!(string_array.len(), num_values as usize);
351        assert_eq!(string_array.value(0), "aaa");
352        assert_eq!(string_array.value(1), "bbb");
353    }
354}
lance_encoding/previous/encodings/physical/fixed_size_binary.rs

lance_encoding/previous/encodings/physical/
fixed_size_binary.rs
