lance_file/versions/v1/encoding/
plain.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Plain encoding
5//!
6//! Plain encoding works with fixed stride types, i.e., `boolean`, `i8...i64`, `f16...f64`,
7//! it stores the array directly in the file. It offers O(1) read access.
8
9use std::ops::Range;
10use std::slice::from_raw_parts;
11use std::sync::Arc;
12
13use arrow_arith::numeric::sub;
14use arrow_array::{
15    Array, ArrayRef, BooleanArray, FixedSizeBinaryArray, FixedSizeListArray, UInt8Array,
16    UInt32Array, builder::BooleanBuilder, cast::AsArray, make_array, new_empty_array,
17};
18use arrow_buffer::{Buffer, bit_util};
19use arrow_data::{ArrayDataBuilder, BufferSpec, layout};
20use arrow_schema::{DataType, Field};
21use arrow_select::{concat::concat, take::take};
22use async_recursion::async_recursion;
23use futures::stream::{self, StreamExt, TryStreamExt};
24use lance_arrow::*;
25use lance_core::{Error, Result};
26use lance_io::{
27    ReadBatchParams,
28    traits::{Reader, Writer},
29};
30use tokio::io::AsyncWriteExt;
31
32/// Encoder for plain encoding.
33///
34pub struct PlainEncoder<'a> {
35    writer: &'a mut dyn Writer,
36    data_type: &'a DataType,
37}
38
39impl<'a> PlainEncoder<'a> {
40    pub fn new(writer: &'a mut dyn Writer, data_type: &'a DataType) -> Self {
41        PlainEncoder { writer, data_type }
42    }
43
44    /// Write an continuous plain-encoded array to the writer.
45    pub async fn write(writer: &'a mut dyn Writer, arrays: &[&'a dyn Array]) -> Result<usize> {
46        let pos = writer.tell().await?;
47        if !arrays.is_empty() {
48            let mut encoder = Self::new(writer, arrays[0].data_type());
49            encoder.encode(arrays).await?;
50        }
51        Ok(pos)
52    }
53
54    /// Encode an slice of an Array of a batch.
55    /// Returns the offset of the metadata
56    pub async fn encode(&mut self, arrays: &[&dyn Array]) -> Result<usize> {
57        self.encode_internal(arrays, self.data_type).await
58    }
59
60    #[async_recursion]
61    async fn encode_internal(
62        &mut self,
63        array: &[&dyn Array],
64        data_type: &DataType,
65    ) -> Result<usize> {
66        if let DataType::FixedSizeList(items, _) = data_type {
67            self.encode_fixed_size_list(array, items).await
68        } else {
69            self.encode_primitive(array).await
70        }
71    }
72
73    async fn encode_boolean(&mut self, arrays: &[&BooleanArray]) -> Result<()> {
74        let capacity: usize = arrays.iter().map(|a| a.len()).sum();
75        let mut builder = BooleanBuilder::with_capacity(capacity);
76
77        for array in arrays {
78            for val in array.iter() {
79                builder.append_value(val.unwrap_or_default());
80            }
81        }
82
83        let boolean_array = builder.finish();
84        self.writer
85            .write_all(boolean_array.into_data().buffers()[0].as_slice())
86            .await?;
87        Ok(())
88    }
89
90    /// Encode array of primitive values.
91    async fn encode_primitive(&mut self, arrays: &[&dyn Array]) -> Result<usize> {
92        assert!(!arrays.is_empty());
93        let data_type = arrays[0].data_type();
94        let offset = self.writer.tell().await?;
95
96        if matches!(data_type, DataType::Boolean) {
97            let boolean_arr = arrays
98                .iter()
99                .map(|a| a.as_boolean())
100                .collect::<Vec<&BooleanArray>>();
101            self.encode_boolean(boolean_arr.as_slice()).await?;
102        } else {
103            let byte_width = data_type.byte_width();
104            for a in arrays.iter() {
105                let data = a.to_data();
106                let slice = unsafe {
107                    from_raw_parts(
108                        data.buffers()[0].as_ptr().add(a.offset() * byte_width),
109                        a.len() * byte_width,
110                    )
111                };
112                self.writer.write_all(slice).await?;
113            }
114        }
115        Ok(offset)
116    }
117
118    /// Encode fixed size list.
119    async fn encode_fixed_size_list(
120        &mut self,
121        arrays: &[&dyn Array],
122        items: &Field,
123    ) -> Result<usize> {
124        let mut value_arrs: Vec<ArrayRef> = Vec::new();
125
126        for array in arrays {
127            let list_array = array
128                .as_any()
129                .downcast_ref::<FixedSizeListArray>()
130                .ok_or_else(|| {
131                    Error::schema(format!(
132                        "Needed a FixedSizeListArray but got {}",
133                        array.data_type()
134                    ))
135                })?;
136            let offset = list_array.value_offset(0) as usize;
137            let length = list_array.len();
138            let value_length = list_array.value_length() as usize;
139            let value_array = list_array.values().slice(offset, length * value_length);
140            value_arrs.push(value_array);
141        }
142
143        self.encode_internal(
144            value_arrs
145                .iter()
146                .map(|a| a.as_ref())
147                .collect::<Vec<_>>()
148                .as_slice(),
149            items.data_type(),
150        )
151        .await
152    }
153}
154
155/// Decoder for plain encoding.
156pub struct PlainDecoder<'a> {
157    reader: &'a dyn Reader,
158    data_type: &'a DataType,
159    /// The start position of the batch in the file.
160    position: usize,
161    /// Number of the rows in this batch.
162    length: usize,
163}
164
165/// Get byte range from the row offset range.
166#[inline]
167fn get_byte_range(data_type: &DataType, row_range: Range<usize>) -> Range<usize> {
168    match data_type {
169        DataType::Boolean => row_range.start / 8..bit_util::ceil(row_range.end, 8),
170        _ => row_range.start * data_type.byte_width()..row_range.end * data_type.byte_width(),
171    }
172}
173
174fn bytes_to_array(
175    data_type: &DataType,
176    bytes: bytes::Bytes,
177    len: usize,
178    offset: usize,
179) -> Result<ArrayRef> {
180    let layout = layout(data_type);
181    if layout.buffers.len() != 1 {
182        return Err(Error::internal(format!(
183            "v1 plain encoding requires one value buffer, found {data_type}"
184        )));
185    }
186
187    let buffer = if let BufferSpec::FixedWidth {
188        byte_width,
189        alignment,
190    } = &layout.buffers[0]
191    {
192        let min_buffer_size = (len + offset).saturating_mul(*byte_width);
193        if bytes.len() < min_buffer_size {
194            Buffer::copy_bytes_bytes(bytes, min_buffer_size)
195        } else {
196            Buffer::from_bytes_bytes(bytes, *alignment as u64)
197        }
198    } else {
199        Buffer::from_slice_ref(bytes)
200    };
201
202    let data = ArrayDataBuilder::new(data_type.clone())
203        .len(len)
204        .offset(offset)
205        .null_count(0)
206        .add_buffer(buffer)
207        .build()?;
208    Ok(make_array(data))
209}
210
211impl<'a> PlainDecoder<'a> {
212    pub fn new(
213        reader: &'a dyn Reader,
214        data_type: &'a DataType,
215        position: usize,
216        length: usize,
217    ) -> Result<Self> {
218        Ok(PlainDecoder {
219            reader,
220            data_type,
221            position,
222            length,
223        })
224    }
225
226    /// Decode primitive values, from "offset" to "offset + length".
227    ///
228    async fn decode_primitive(&self, start: usize, end: usize) -> Result<ArrayRef> {
229        if end > self.length {
230            return Err(Error::invalid_input(format!(
231                "PlainDecoder: request([{}..{}]) out of range: [0..{}]",
232                start, end, self.length
233            )));
234        }
235        let byte_range = get_byte_range(self.data_type, start..end);
236        let range = Range {
237            start: self.position + byte_range.start,
238            end: self.position + byte_range.end,
239        };
240
241        let data = self.reader.get_range(range).await?;
242        // booleans are bitpacked, so we need an offset to provide the exact
243        // requested range.
244        let offset = if self.data_type == &DataType::Boolean {
245            start % 8
246        } else {
247            0
248        };
249        bytes_to_array(self.data_type, data, end - start, offset)
250    }
251
252    async fn decode_fixed_size_list(
253        &self,
254        items: &Field,
255        list_size: i32,
256        start: usize,
257        end: usize,
258    ) -> Result<ArrayRef> {
259        if !items.data_type().is_fixed_stride() {
260            return Err(Error::schema(format!(
261                "Items for fixed size list should be primitives but found {}",
262                items.data_type()
263            )));
264        };
265        let item_decoder = PlainDecoder::new(
266            self.reader,
267            items.data_type(),
268            self.position,
269            self.length * list_size as usize,
270        )?;
271        let item_array = item_decoder
272            .get(start * list_size as usize..end * list_size as usize)
273            .await?;
274        Ok(Arc::new(FixedSizeListArray::new(
275            Arc::new(items.clone()),
276            list_size,
277            item_array,
278            None,
279        )) as ArrayRef)
280    }
281
282    async fn decode_fixed_size_binary(
283        &self,
284        stride: i32,
285        start: usize,
286        end: usize,
287    ) -> Result<ArrayRef> {
288        let bytes_decoder = PlainDecoder::new(
289            self.reader,
290            &DataType::UInt8,
291            self.position,
292            self.length * stride as usize,
293        )?;
294        let bytes_array = bytes_decoder
295            .get(start * stride as usize..end * stride as usize)
296            .await?;
297        let values = bytes_array
298            .as_any()
299            .downcast_ref::<UInt8Array>()
300            .ok_or_else(|| {
301                Error::schema("Could not cast to UInt8Array for FixedSizeBinary".to_string())
302            })?;
303        Ok(Arc::new(FixedSizeBinaryArray::try_new_from_values(values, stride)?) as ArrayRef)
304    }
305
306    async fn take_boolean(&self, indices: &UInt32Array) -> Result<ArrayRef> {
307        let block_size = self.reader.block_size() as u32;
308        let boolean_block_size = block_size * 8;
309
310        let mut chunk_ranges = vec![];
311        let mut start: u32 = 0;
312        for j in 0..(indices.len() - 1) as u32 {
313            if (indices.value(j as usize + 1) / boolean_block_size)
314                > (indices.value(start as usize) / boolean_block_size)
315            {
316                let next_start = j + 1;
317                chunk_ranges.push(start..next_start);
318                start = next_start;
319            }
320        }
321        // Remaining
322        chunk_ranges.push(start..indices.len() as u32);
323
324        let arrays = stream::iter(chunk_ranges)
325            .map(|cr| async move {
326                let request = indices.slice(cr.start as usize, cr.len());
327                // request contains the array indices we are retrieving in this chunk.
328
329                // Get the starting index
330                let start = request.value(0);
331                // Final index is the last value
332                let end = request.value(request.len() - 1);
333                let array = self.get(start as usize..end as usize + 1).await?;
334
335                let shifted_indices = sub(&request, &UInt32Array::new_scalar(start))?;
336                Ok::<ArrayRef, Error>(take(&array, &shifted_indices, None)?)
337            })
338            .buffered(self.reader.io_parallelism())
339            .try_collect::<Vec<_>>()
340            .await?;
341        let references = arrays.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
342        Ok(concat(&references)?)
343    }
344}
345
346fn make_chunked_requests(
347    indices: &[u32],
348    byte_width: usize,
349    block_size: usize,
350) -> Vec<Range<usize>> {
351    let mut chunked_ranges = vec![];
352    let mut start: usize = 0;
353    // Note: limit the I/O size to the block size.
354    //
355    // Another option could be checking whether `indices[i]` and `indices[i+1]` are not
356    // farther way than the block size:
357    //    indices[i] * byte_width + block_size < indices[i+1] * byte_width
358    // It might allow slightly larger sequential reads.
359    for i in 0..indices.len() - 1 {
360        // If contiguous, continue
361        if indices[i + 1] == indices[i] + 1 {
362            continue;
363        }
364        if indices[i + 1] as usize * byte_width > indices[start] as usize * byte_width + block_size
365        {
366            chunked_ranges.push(start..i + 1);
367            start = i + 1;
368        }
369    }
370    chunked_ranges.push(start..indices.len());
371    chunked_ranges
372}
373
374impl PlainDecoder<'_> {
375    pub async fn decode(&self) -> Result<ArrayRef> {
376        self.get(0..self.length).await
377    }
378
379    pub async fn take(&self, indices: &UInt32Array) -> Result<ArrayRef> {
380        if indices.is_empty() {
381            return Ok(new_empty_array(self.data_type));
382        }
383
384        if matches!(self.data_type, DataType::Boolean) {
385            return self.take_boolean(indices).await;
386        }
387
388        let block_size = self.reader.block_size();
389        let byte_width = self.data_type.byte_width();
390
391        let chunked_ranges = make_chunked_requests(indices.values(), byte_width, block_size);
392
393        let arrays = stream::iter(chunked_ranges)
394            .map(|cr| async move {
395                let request = indices.slice(cr.start, cr.len());
396
397                let start = request.value(0);
398                let end = request.value(request.len() - 1);
399                let array = self.get(start as usize..end as usize + 1).await?;
400                let adjusted_offsets = sub(&request, &UInt32Array::new_scalar(start))?;
401                Ok::<ArrayRef, Error>(take(&array, &adjusted_offsets, None)?)
402            })
403            .buffered(self.reader.io_parallelism())
404            .try_collect::<Vec<_>>()
405            .await?;
406        let references = arrays.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
407        Ok(concat(&references)?)
408    }
409
410    #[async_recursion]
411    async fn decode_range(&self, index: Range<usize>) -> Result<ArrayRef> {
412        if index.is_empty() {
413            return Ok(new_empty_array(self.data_type));
414        }
415        match self.data_type {
416            DataType::FixedSizeList(items, list_size) => {
417                self.decode_fixed_size_list(items, *list_size, index.start, index.end)
418                    .await
419            }
420            DataType::FixedSizeBinary(stride) => {
421                self.decode_fixed_size_binary(*stride, index.start, index.end)
422                    .await
423            }
424            _ => self.decode_primitive(index.start, index.end).await,
425        }
426    }
427
428    pub async fn get(&self, params: impl Into<ReadBatchParams>) -> Result<ArrayRef> {
429        match params.into() {
430            ReadBatchParams::Range(range) => self.decode_range(range).await,
431            ReadBatchParams::Ranges(_) => Err(Error::invalid_input(
432                "multiple ranges are not supported by v1 plain encoding",
433            )),
434            ReadBatchParams::RangeFull => self.decode_range(0..self.length).await,
435            ReadBatchParams::RangeTo(range) => self.decode_range(0..range.end).await,
436            ReadBatchParams::RangeFrom(range) => self.decode_range(range.start..self.length).await,
437            ReadBatchParams::Indices(indices) => self.take(&indices).await,
438        }
439    }
440}
441
442#[cfg(test)]
443mod tests {
444    use std::ops::Deref;
445
446    use arrow_array::*;
447    use arrow_buffer::Buffer;
448    use arrow_data::ArrayDataBuilder;
449    use bytes::Bytes;
450    use lance_core::utils::tempfile::TempStdFile;
451    use lance_io::local::LocalObjectReader;
452    use rand::prelude::*;
453
454    use super::*;
455
456    #[tokio::test]
457    async fn test_encode_decode_primitive_array() {
458        let int_types = vec![
459            DataType::Int8,
460            DataType::Int16,
461            DataType::Int32,
462            DataType::Int64,
463            DataType::UInt8,
464            DataType::UInt16,
465            DataType::UInt32,
466            DataType::UInt64,
467        ];
468        let input: Vec<i64> = Vec::from_iter(1..127_i64);
469        for t in int_types {
470            let buffer = Buffer::from_slice_ref(input.as_slice());
471            let mut arrs: Vec<ArrayRef> = Vec::new();
472            for _ in 0..10 {
473                arrs.push(Arc::new(make_array_(&t, &buffer).await));
474            }
475            test_round_trip(arrs.as_slice(), t).await;
476        }
477
478        let float_types = vec![DataType::Float16, DataType::Float32, DataType::Float64];
479        let mut rng = rand::rng();
480        let input: Vec<f64> = (1..127).map(|_| rng.random()).collect();
481        for t in float_types {
482            let buffer = Buffer::from_slice_ref(input.as_slice());
483            let mut arrs: Vec<ArrayRef> = Vec::new();
484
485            for _ in 0..10 {
486                arrs.push(Arc::new(make_array_(&t, &buffer).await));
487            }
488            test_round_trip(arrs.as_slice(), t).await;
489        }
490    }
491
492    async fn test_round_trip(expected: &[ArrayRef], data_type: DataType) {
493        let path = TempStdFile::default();
494
495        let expected_as_array = expected
496            .iter()
497            .map(|e| e.as_ref())
498            .collect::<Vec<&dyn Array>>();
499        {
500            let mut writer = tokio::fs::File::create(&path).await.unwrap();
501            let mut encoder = PlainEncoder::new(&mut writer, &data_type);
502            assert_eq!(
503                encoder.encode(expected_as_array.as_slice()).await.unwrap(),
504                0
505            );
506            writer.flush().await.unwrap();
507        }
508
509        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
510            .await
511            .unwrap();
512        assert!(reader.size().await.unwrap() > 0);
513        // Expected size is the total of all arrays
514        let expected_size = expected.iter().map(|e| e.len()).sum();
515        let decoder = PlainDecoder::new(reader.as_ref(), &data_type, 0, expected_size).unwrap();
516        let arr = decoder.decode().await.unwrap();
517        let actual = arr.as_ref();
518        let expected_merged = concat(expected_as_array.as_slice()).unwrap();
519        assert_eq!(expected_merged.deref(), actual);
520        assert_eq!(expected_size, actual.len());
521    }
522
523    #[tokio::test]
524    async fn test_encode_decode_bool_array() {
525        let mut arrs: Vec<ArrayRef> = Vec::new();
526
527        for _ in 0..10 {
528            // It is important that the boolean array length is < 8 so we can test if the Arrays are merged correctly
529            arrs.push(Arc::new(BooleanArray::from(vec![true, true, true])) as ArrayRef);
530        }
531        test_round_trip(arrs.as_slice(), DataType::Boolean).await;
532    }
533
534    #[tokio::test]
535    async fn test_encode_decode_fixed_size_list_array() {
536        let int_types = vec![
537            DataType::Int8,
538            DataType::Int16,
539            DataType::Int32,
540            DataType::Int64,
541            DataType::UInt8,
542            DataType::UInt16,
543            DataType::UInt32,
544            DataType::UInt64,
545        ];
546        let input = Vec::from_iter(1..127_i64);
547        for t in int_types {
548            let buffer = Buffer::from_slice_ref(input.as_slice());
549            let list_type =
550                DataType::FixedSizeList(Arc::new(Field::new("item", t.clone(), true)), 3);
551            let mut arrs: Vec<ArrayRef> = Vec::new();
552
553            for _ in 0..10 {
554                let items = make_array_(&t.clone(), &buffer).await;
555                let arr = FixedSizeListArray::try_new_from_values(items, 3).unwrap();
556                arrs.push(Arc::new(arr) as ArrayRef);
557            }
558            test_round_trip(arrs.as_slice(), list_type).await;
559        }
560    }
561
562    #[tokio::test]
563    async fn test_encode_decode_fixed_size_binary_array() {
564        let t = DataType::FixedSizeBinary(3);
565        let mut arrs: Vec<ArrayRef> = Vec::new();
566
567        for _ in 0..10 {
568            let values = UInt8Array::from(Vec::from_iter(1..127_u8));
569            let arr = FixedSizeBinaryArray::try_new_from_values(&values, 3).unwrap();
570            arrs.push(Arc::new(arr) as ArrayRef);
571        }
572        test_round_trip(arrs.as_slice(), t).await;
573    }
574
575    #[tokio::test]
576    async fn test_bytes_to_array_padding() {
577        let bytes = Bytes::from_static(&[0x01, 0x00, 0x02, 0x00, 0x03]);
578        let arr = bytes_to_array(&DataType::UInt16, bytes, 3, 0).unwrap();
579
580        let expected = UInt16Array::from(vec![1, 2, 3]);
581        assert_eq!(arr.as_ref(), &expected);
582
583        // Underlying data is padded to the nearest multiple of two bytes (for u16).
584        let data = arr.to_data();
585        let buf = &data.buffers()[0];
586        let repr = format!("{:?}", buf);
587        assert!(
588            repr.contains("[1, 0, 2, 0, 3, 0]"),
589            "Underlying buffer contains unexpected data: {}",
590            repr
591        );
592    }
593
594    #[tokio::test]
595    async fn test_encode_decode_nested_fixed_size_list() {
596        // FixedSizeList of FixedSizeList
597        let inner = DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 2);
598        let t = DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2);
599        let mut arrs: Vec<ArrayRef> = Vec::new();
600
601        for _ in 0..10 {
602            let values = Int64Array::from_iter_values(1..=120_i64);
603            let arr = FixedSizeListArray::try_new_from_values(
604                FixedSizeListArray::try_new_from_values(values, 2).unwrap(),
605                2,
606            )
607            .unwrap();
608            arrs.push(Arc::new(arr) as ArrayRef);
609        }
610        test_round_trip(arrs.as_slice(), t).await;
611
612        // FixedSizeList of FixedSizeBinary
613        let inner = DataType::FixedSizeBinary(2);
614        let t = DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2);
615        let mut arrs: Vec<ArrayRef> = Vec::new();
616
617        for _ in 0..10 {
618            let values = UInt8Array::from_iter_values(1..=120_u8);
619            let arr = FixedSizeListArray::try_new_from_values(
620                FixedSizeBinaryArray::try_new_from_values(&values, 2).unwrap(),
621                2,
622            )
623            .unwrap();
624            arrs.push(Arc::new(arr) as ArrayRef);
625        }
626        test_round_trip(arrs.as_slice(), t).await;
627    }
628
629    async fn make_array_(data_type: &DataType, buffer: &Buffer) -> ArrayRef {
630        arrow_array::make_array(
631            ArrayDataBuilder::new(data_type.clone())
632                .len(126)
633                .add_buffer(buffer.clone())
634                .build()
635                .unwrap(),
636        )
637    }
638
639    #[tokio::test]
640    async fn test_decode_by_range() {
641        let path = TempStdFile::default();
642
643        let array = Int32Array::from_iter_values([0, 1, 2, 3, 4, 5]);
644        {
645            let mut writer = tokio::fs::File::create(&path).await.unwrap();
646            let mut encoder = PlainEncoder::new(&mut writer, array.data_type());
647            assert_eq!(encoder.encode(&[&array]).await.unwrap(), 0);
648            writer.flush().await.unwrap();
649        }
650
651        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
652            .await
653            .unwrap();
654        assert!(reader.size().await.unwrap() > 0);
655        let decoder =
656            PlainDecoder::new(reader.as_ref(), array.data_type(), 0, array.len()).unwrap();
657        assert_eq!(
658            decoder.get(2..4).await.unwrap().as_ref(),
659            &Int32Array::from_iter_values([2, 3])
660        );
661
662        assert_eq!(
663            decoder.get(..4).await.unwrap().as_ref(),
664            &Int32Array::from_iter_values([0, 1, 2, 3])
665        );
666
667        assert_eq!(
668            decoder.get(2..).await.unwrap().as_ref(),
669            &Int32Array::from_iter_values([2, 3, 4, 5])
670        );
671
672        assert_eq!(
673            &decoder.get(2..2).await.unwrap(),
674            &new_empty_array(&DataType::Int32)
675        );
676
677        assert_eq!(
678            &decoder.get(5..5).await.unwrap(),
679            &new_empty_array(&DataType::Int32)
680        );
681
682        assert!(decoder.get(3..1000).await.is_err());
683    }
684
685    #[tokio::test]
686    async fn test_take() {
687        let path = TempStdFile::default();
688
689        let array = Int32Array::from_iter_values(0..100);
690
691        {
692            let mut writer = tokio::fs::File::create(&path).await.unwrap();
693            let mut encoder = PlainEncoder::new(&mut writer, array.data_type());
694            assert_eq!(encoder.encode(&[&array]).await.unwrap(), 0);
695            AsyncWriteExt::shutdown(&mut writer).await.unwrap();
696        }
697
698        let reader = LocalObjectReader::open_local_path(&path, 2048, None)
699            .await
700            .unwrap();
701        assert!(reader.size().await.unwrap() > 0);
702        let decoder =
703            PlainDecoder::new(reader.as_ref(), array.data_type(), 0, array.len()).unwrap();
704
705        let results = decoder
706            .take(&UInt32Array::from_iter(
707                [2, 4, 5, 20, 30, 55, 60].iter().map(|i| *i as u32),
708            ))
709            .await
710            .unwrap();
711        assert_eq!(
712            results.as_ref(),
713            &Int32Array::from_iter_values([2, 4, 5, 20, 30, 55, 60])
714        );
715    }
716
717    // Re-enable the following tests once the Lance FileReader / FileWrite is migrated.
718
719    // #[tokio::test]
720    // async fn test_boolean_slice() {
721    //     let store = ObjectStore::memory();
722    //     let path = Path::from("/bool_slice");
723    //
724    //     let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
725    //         "b",
726    //         DataType::Boolean,
727    //         true,
728    //     )]));
729    //     let schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
730    //     let mut file_writer =
731    //         FileWriter::try_new(&store, &path, schema.clone(), &Default::default())
732    //             .await
733    //             .unwrap();
734    //
735    //     let array = BooleanArray::from((0..120).map(|v| v % 5 == 0).collect::<Vec<_>>());
736    //     for i in 0..10 {
737    //         let data = array.slice(i * 12, 12); // one and half byte
738    //         file_writer
739    //             .write(&[RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(data)]).unwrap()])
740    //             .await
741    //             .unwrap();
742    //     }
743    //     file_writer.finish().await.unwrap();
744    //
745    //     let batch = read_file_as_one_batch(&store, &path).await;
746    //     assert_eq!(batch.column_by_name("b").unwrap().as_ref(), &array);
747    //
748    //     let array = BooleanArray::from(vec![Some(true), Some(false), None]);
749    //     let mut file_writer = FileWriter::try_new(&store, &path, schema, &Default::default())
750    //         .await
751    //         .unwrap();
752    //     file_writer
753    //         .write(&[RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(array)]).unwrap()])
754    //         .await
755    //         .unwrap();
756    //     file_writer.finish().await.unwrap();
757    //     let batch = read_file_as_one_batch(&store, &path).await;
758    //
759    //     // None default to Some(false), since we don't support null values yet.
760    //     let expected = BooleanArray::from(vec![Some(true), Some(false), Some(false)]);
761    //     assert_eq!(batch.column_by_name("b").unwrap().as_ref(), &expected);
762    // }
763    //
764    // #[tokio::test]
765    // async fn test_encode_fixed_size_list_slice() {
766    //     let store = ObjectStore::memory();
767    //     let path = Path::from("/shared_slice");
768    //
769    //     let array = Int32Array::from_iter_values(0..1600);
770    //     let fixed_size_list = FixedSizeListArray::try_new_from_values(array, 16).unwrap();
771    //     let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
772    //         "fl",
773    //         fixed_size_list.data_type().clone(),
774    //         false,
775    //     )]));
776    //     let schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
777    //     let mut file_writer = FileWriter::try_new(&store, &path, schema, &Default::default())
778    //         .await
779    //         .unwrap();
780    //
781    //     for i in (0..100).step_by(4) {
782    //         let slice: FixedSizeListArray = fixed_size_list.slice(i, 4);
783    //         file_writer
784    //             .write(&[
785    //                 RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(slice)]).unwrap(),
786    //             ])
787    //             .await
788    //             .unwrap();
789    //     }
790    //     file_writer.finish().await.unwrap();
791    //
792    //     let batch = read_file_as_one_batch(&store, &path).await;
793    //     assert_eq!(
794    //         batch.column_by_name("fl").unwrap().as_ref(),
795    //         &fixed_size_list
796    //     );
797    // }
798    //
799    // #[tokio::test]
800    // async fn test_take_boolean() {
801    //     let temp_dir = tempfile::tempdir().unwrap();
802    //     let path = temp_dir.join("/bool_take");
803    //
804    //     let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
805    //         "b",
806    //         DataType::Boolean,
807    //         false,
808    //     )]));
809    //     let schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
810    //     let mut file_writer =
811    //         FileWriter::try_new(&store, &path, schema.clone(), &Default::default())
812    //             .await
813    //             .unwrap();
814    //
815    //     let array = BooleanArray::from((0..120).map(|v| v % 5 == 0).collect::<Vec<_>>());
816    //     let batch =
817    //         RecordBatch::try_new(arrow_schema.clone(), vec![Arc::new(array.clone())]).unwrap();
818    //     file_writer.write(&[batch]).await.unwrap();
819    //     file_writer.finish().await.unwrap();
820    //
821    //     let reader = FileReader::try_new(&store, &path).await.unwrap();
822    //     let actual = reader
823    //         .take(&[2, 4, 5, 8, 63, 64, 65], &schema)
824    //         .await
825    //         .unwrap();
826    //
827    //     assert_eq!(
828    //         actual.column_by_name("b").unwrap().as_ref(),
829    //         &BooleanArray::from(vec![false, false, true, false, false, false, true])
830    //     );
831    // }
832
833    #[test]
834    fn test_make_chunked_request() {
835        let byte_width: usize = 4096; // 4K
836        let prefetch_size: usize = 64 * 1024; // 64KB.
837        let u32_overflow: usize = u32::MAX as usize + 10;
838
839        let indices: Vec<u32> = vec![
840            1,
841            10,
842            20,
843            100,
844            120,
845            (u32_overflow / byte_width) as u32, // Two overflow offsets
846            (u32_overflow / byte_width) as u32 + 100,
847        ];
848        let chunks = make_chunked_requests(&indices, byte_width, prefetch_size);
849        assert_eq!(chunks.len(), 6, "got chunks: {:?}", chunks);
850        assert_eq!(chunks, vec![(0..2), (2..3), (3..4), (4..5), (5..6), (6..7)])
851    }
852}
lance_file/versions/v1/encoding/plain.rs

lance_file/versions/v1/encoding/
plain.rs
