lance_file/versions/v1/encoding/
binary.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Var-length binary encoding.
5//!
6
7use std::marker::PhantomData;
8use std::ops::Range;
9use std::sync::Arc;
10
11use arrow_arith::numeric::sub;
12use arrow_array::{
13    Array, ArrayRef, GenericByteArray, Int64Array, OffsetSizeTrait, UInt32Array,
14    builder::{ArrayBuilder, PrimitiveBuilder},
15    cast::AsArray,
16    cast::as_primitive_array,
17    new_empty_array,
18    types::{
19        BinaryType, ByteArrayType, Int64Type, LargeBinaryType, LargeUtf8Type, UInt32Type, Utf8Type,
20    },
21};
22use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ScalarBuffer, bit_util};
23use arrow_cast::cast::cast;
24use arrow_data::ArrayDataBuilder;
25use arrow_schema::DataType;
26use bytes::Bytes;
27use futures::{StreamExt, TryStreamExt};
28use lance_arrow::BufferExt;
29use tokio::io::AsyncWriteExt;
30
31use super::plain::PlainDecoder;
32use lance_core::Result;
33use lance_io::{
34    ReadBatchParams,
35    traits::{Reader, Writer},
36};
37
38/// Encoder for Var-binary encoding.
39pub struct BinaryEncoder<'a> {
40    writer: &'a mut dyn Writer,
41}
42
43impl<'a> BinaryEncoder<'a> {
44    pub fn new(writer: &'a mut dyn Writer) -> Self {
45        Self { writer }
46    }
47
48    async fn encode_typed_arr<T: ByteArrayType>(&mut self, arrs: &[&dyn Array]) -> Result<usize> {
49        let capacity: usize = arrs.iter().map(|a| a.len()).sum();
50        let mut pos_builder: PrimitiveBuilder<Int64Type> =
51            PrimitiveBuilder::with_capacity(capacity + 1);
52
53        let mut last_offset: usize = self.writer.tell().await?;
54        pos_builder.append_value(last_offset as i64);
55        for array in arrs.iter() {
56            let arr = array
57                .as_any()
58                .downcast_ref::<GenericByteArray<T>>()
59                .unwrap();
60
61            let offsets = arr.value_offsets();
62
63            let start = offsets[0].as_usize();
64            let end = offsets[offsets.len() - 1].as_usize();
65            let b = unsafe {
66                std::slice::from_raw_parts(
67                    arr.to_data().buffers()[1].as_ptr().add(start),
68                    end - start,
69                )
70            };
71            self.writer.write_all(b).await?;
72
73            let start_offset = offsets[0].as_usize();
74            offsets
75                .iter()
76                .skip(1)
77                .map(|b| b.as_usize() - start_offset + last_offset)
78                .for_each(|o| pos_builder.append_value(o as i64));
79            last_offset = pos_builder.values_slice()[pos_builder.len() - 1] as usize;
80        }
81
82        let positions_offset = self.writer.tell().await?;
83        let pos_array = pos_builder.finish();
84        self.writer
85            .write_all(pos_array.to_data().buffers()[0].as_slice())
86            .await?;
87        Ok(positions_offset)
88    }
89}
90
91impl BinaryEncoder<'_> {
92    pub async fn encode(&mut self, arrs: &[&dyn Array]) -> Result<usize> {
93        assert!(!arrs.is_empty());
94        let data_type = arrs[0].data_type();
95        match data_type {
96            DataType::Utf8 => self.encode_typed_arr::<Utf8Type>(arrs).await,
97            DataType::Binary => self.encode_typed_arr::<BinaryType>(arrs).await,
98            DataType::LargeUtf8 => self.encode_typed_arr::<LargeUtf8Type>(arrs).await,
99            DataType::LargeBinary => self.encode_typed_arr::<LargeBinaryType>(arrs).await,
100            _ => Err(lance_core::Error::invalid_input(format!(
101                "Unsupported data type for binary encoding: {}",
102                data_type
103            ))),
104        }
105    }
106}
107
108/// Var-binary encoding decoder.
109pub struct BinaryDecoder<'a, T: ByteArrayType> {
110    reader: &'a dyn Reader,
111
112    position: usize,
113
114    length: usize,
115
116    nullable: bool,
117
118    phantom: PhantomData<T>,
119}
120
121/// Var-length Binary Decoder
122///
123impl<'a, T: ByteArrayType> BinaryDecoder<'a, T> {
124    /// Create a [BinaryEncoder] to decode one batch.
125    ///
126    ///  - `position`, file position where this batch starts.
127    ///  - `length`, the number of records in this batch.
128    ///  - `nullable`, whether this batch contains nullable value.
129    ///
130    /// ## Example
131    ///
132    /// ```rust
133    /// use arrow_array::types::Utf8Type;
134    /// use object_store::path::Path;
135    /// use lance_file::versions::v1::encoding::binary::BinaryDecoder;
136    /// use lance_io::{local::LocalObjectReader, traits::Reader};
137    ///
138    /// async {
139    ///     let reader = LocalObjectReader::open_local_path("/tmp/foo.lance", 2048, None).await.unwrap();
140    ///     let string_decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), 100, 1024, true);
141    /// };
142    /// ```
143    pub fn new(reader: &'a dyn Reader, position: usize, length: usize, nullable: bool) -> Self {
144        Self {
145            reader,
146            position,
147            length,
148            nullable,
149            phantom: PhantomData,
150        }
151    }
152
153    /// Get the position array for the batch.
154    async fn get_positions(&self, index: Range<usize>) -> Result<Arc<Int64Array>> {
155        let position_decoder = PlainDecoder::new(
156            self.reader,
157            &DataType::Int64,
158            self.position,
159            self.length + 1,
160        )?;
161        let values = position_decoder.get(index.start..index.end + 1).await?;
162        Ok(Arc::new(as_primitive_array(&values).clone()))
163    }
164
165    fn count_nulls<O: OffsetSizeTrait>(offsets: &ScalarBuffer<O>) -> (usize, Option<Buffer>) {
166        let mut null_count = 0;
167        let mut null_buf = MutableBuffer::new_null(offsets.len() - 1);
168        offsets.windows(2).enumerate().for_each(|(idx, w)| {
169            if w[0] == w[1] {
170                bit_util::unset_bit(null_buf.as_mut(), idx);
171                null_count += 1;
172            } else {
173                bit_util::set_bit(null_buf.as_mut(), idx);
174            }
175        });
176        let null_buf = if null_count > 0 {
177            Some(null_buf.into())
178        } else {
179            None
180        };
181        (null_count, null_buf)
182    }
183
184    /// Read the array with batch positions and range.
185    ///
186    /// Parameters
187    ///
188    ///  - *positions*: position array for the batch.
189    ///  - *range*: range of rows to read.
190    async fn get_range(&self, positions: &Int64Array, range: Range<usize>) -> Result<ArrayRef> {
191        assert!(positions.len() >= range.end);
192        let start = positions.value(range.start);
193        let end = positions.value(range.end);
194
195        let start_scalar = Int64Array::new_scalar(start);
196
197        let slice = positions.slice(range.start, range.len() + 1);
198        let offset_data = if T::Offset::IS_LARGE {
199            sub(&slice, &start_scalar)?.into_data()
200        } else {
201            cast(
202                &(Arc::new(sub(&slice, &start_scalar)?) as ArrayRef),
203                &DataType::Int32,
204            )?
205            .into_data()
206        };
207
208        let bytes: Bytes = if start >= end {
209            Bytes::new()
210        } else {
211            self.reader.get_range(start as usize..end as usize).await?
212        };
213
214        let mut data_builder = ArrayDataBuilder::new(T::DATA_TYPE)
215            .len(range.len())
216            .null_count(0);
217
218        // Count nulls
219        if self.nullable {
220            let (null_count, null_buf) = Self::count_nulls(slice.values());
221            data_builder = data_builder
222                .null_count(null_count)
223                .null_bit_buffer(null_buf);
224        }
225
226        let buf = Buffer::from_bytes_bytes(bytes, /*bytes_per_value=*/ 1);
227        let array_data = data_builder
228            .add_buffer(offset_data.buffers()[0].clone())
229            .add_buffer(buf)
230            .build()?;
231
232        Ok(Arc::new(GenericByteArray::<T>::from(array_data)))
233    }
234}
235
236#[derive(Debug)]
237struct TakeChunksPlan {
238    indices: UInt32Array,
239    is_contiguous: bool,
240}
241
242/// Group the indices into chunks, such that either:
243/// 1. the indices are contiguous (and non-repeating)
244/// 2. the values are within `min_io_size` of each other (and thus are worth
245///    grabbing in a single request)
246fn plan_take_chunks(
247    positions: &Int64Array,
248    indices: &UInt32Array,
249    min_io_size: i64,
250) -> Result<Vec<TakeChunksPlan>> {
251    let start = indices.value(0);
252    let indices = sub(indices, &UInt32Array::new_scalar(start))?;
253    let indices_ref = indices.as_primitive::<UInt32Type>();
254
255    let mut chunks: Vec<TakeChunksPlan> = vec![];
256    let mut start_idx = 0;
257    let mut last_idx: i64 = -1;
258    let mut is_contiguous = true;
259    for i in 0..indices.len() {
260        let current = indices_ref.value(i) as usize;
261        let curr_contiguous = current == start_idx || current as i64 - last_idx == 1;
262
263        if !curr_contiguous
264            && positions.value(current) - positions.value(indices_ref.value(start_idx) as usize)
265                > min_io_size
266        {
267            chunks.push(TakeChunksPlan {
268                indices: as_primitive_array(&indices.slice(start_idx, i - start_idx)).clone(),
269                is_contiguous,
270            });
271            start_idx = i;
272            is_contiguous = true;
273        } else {
274            is_contiguous &= curr_contiguous;
275        }
276
277        last_idx = current as i64;
278    }
279    chunks.push(TakeChunksPlan {
280        indices: as_primitive_array(&indices.slice(start_idx, indices.len() - start_idx)).clone(),
281        is_contiguous,
282    });
283
284    Ok(chunks)
285}
286
287impl<T: ByteArrayType> BinaryDecoder<'_, T> {
288    pub async fn decode(&self) -> Result<ArrayRef> {
289        self.get(..).await
290    }
291
292    /// Take the values at the given indices.
293    ///
294    /// This function assumes indices are sorted.
295    pub async fn take(&self, indices: &UInt32Array) -> Result<ArrayRef> {
296        if indices.is_empty() {
297            return Ok(new_empty_array(&T::DATA_TYPE));
298        }
299
300        let start = indices.value(0);
301        let end = indices.value(indices.len() - 1);
302
303        // TODO: make min batch size configurable.
304        // TODO: make reading positions in chunks too.
305        const MIN_IO_SIZE: i64 = 64 * 1024; // 64KB
306        let positions = self
307            .get_positions(start as usize..(end + 1) as usize)
308            .await?;
309        // Use indices and positions to pre-allocate an exact-size buffer
310        let capacity = indices
311            .iter()
312            .map(|i| {
313                let relative_index = (i.unwrap() - start) as usize;
314                let start = positions.value(relative_index) as usize;
315                let end = positions.value(relative_index + 1) as usize;
316                end - start
317            })
318            .sum();
319        let mut buffer = MutableBuffer::with_capacity(capacity);
320
321        let offsets_capacity = std::mem::size_of::<T::Offset>() * (indices.len() + 1);
322        let mut offsets = MutableBuffer::with_capacity(offsets_capacity);
323        let mut offset = T::Offset::from_usize(0).unwrap();
324        // Safety: We allocated appropriate capacity just above.
325        unsafe {
326            offsets.push_unchecked(offset);
327        }
328
329        let chunks = plan_take_chunks(&positions, indices, MIN_IO_SIZE)?;
330
331        let positions_ref = positions.as_ref();
332        futures::stream::iter(chunks)
333            .map(|chunk| async move {
334                let chunk_offset = chunk.indices.value(0);
335                let chunk_end = chunk.indices.value(chunk.indices.len() - 1);
336                let array = self
337                    .get_range(positions_ref, chunk_offset as usize..chunk_end as usize + 1)
338                    .await?;
339                Result::Ok((chunk, chunk_offset, array))
340            })
341            .buffered(self.reader.io_parallelism())
342            .try_for_each(|(chunk, chunk_offset, array)| {
343                let array: &GenericByteArray<T> = array.as_bytes();
344
345                // Faster to do one large memcpy than O(n) small ones.
346                if chunk.is_contiguous {
347                    buffer.extend_from_slice(array.value_data());
348                }
349
350                // Append each value to the buffer in the correct order
351                for index in chunk.indices.values() {
352                    if !chunk.is_contiguous {
353                        let value = array.value((index - chunk_offset) as usize);
354                        let value_ref: &[u8] = value.as_ref();
355                        buffer.extend_from_slice(value_ref);
356                    }
357
358                    offset += array.value_length((index - chunk_offset) as usize);
359                    // Append next offset
360                    // Safety: We allocated appropriate capacity on initialization
361                    unsafe {
362                        offsets.push_unchecked(offset);
363                    }
364                }
365                futures::future::ready(Ok(()))
366            })
367            .await?;
368
369        let mut data_builder = ArrayDataBuilder::new(T::DATA_TYPE)
370            .len(indices.len())
371            .null_count(0);
372
373        let offsets: ScalarBuffer<T::Offset> = ScalarBuffer::from(Buffer::from(offsets));
374
375        // We should have pre-sized perfectly.
376        debug_assert_eq!(buffer.len(), capacity);
377
378        if self.nullable {
379            let (null_count, null_buf) = Self::count_nulls(&offsets);
380            data_builder = data_builder
381                .null_count(null_count)
382                .null_bit_buffer(null_buf);
383        }
384
385        let array_data = data_builder
386            .add_buffer(offsets.into_inner())
387            .add_buffer(buffer.into())
388            .build()?;
389
390        Ok(Arc::new(GenericByteArray::<T>::from(array_data)))
391    }
392}
393
394impl<T: ByteArrayType> BinaryDecoder<'_, T> {
395    async fn decode_range(&self, index: Range<usize>) -> Result<ArrayRef> {
396        if index.end > self.length {
397            return Err(lance_core::Error::invalid_input(format!(
398                "v1 binary row range {}..{} exceeds length {}",
399                index.start, index.end, self.length
400            )));
401        }
402        if index.is_empty() {
403            return Ok(new_empty_array(&T::DATA_TYPE));
404        }
405        let position_decoder = PlainDecoder::new(
406            self.reader,
407            &DataType::Int64,
408            self.position,
409            self.length + 1,
410        )?;
411        let positions = position_decoder.get(index.start..index.end + 1).await?;
412        let int64_positions: &Int64Array = as_primitive_array(&positions);
413
414        self.get_range(int64_positions, 0..index.len()).await
415    }
416
417    pub async fn get(&self, params: impl Into<ReadBatchParams>) -> Result<ArrayRef> {
418        match params.into() {
419            ReadBatchParams::Range(range) => self.decode_range(range).await,
420            ReadBatchParams::Ranges(_) => Err(lance_core::Error::invalid_input(
421                "multiple ranges are not supported by v1 binary encoding",
422            )),
423            ReadBatchParams::RangeFull => self.decode_range(0..self.length).await,
424            ReadBatchParams::RangeTo(range) => self.decode_range(0..range.end).await,
425            ReadBatchParams::RangeFrom(range) => self.decode_range(range.start..self.length).await,
426            ReadBatchParams::Indices(indices) => self.take(&indices).await,
427        }
428    }
429}
430
431#[cfg(test)]
432mod tests {
433    use super::*;
434
435    use arrow_array::{
436        BinaryArray, GenericStringArray, LargeStringArray, StringArray, types::GenericStringType,
437    };
438    use arrow_select::concat::concat;
439    use lance_core::utils::tempfile::TempStdFile;
440
441    use lance_io::local::LocalObjectReader;
442
443    async fn write_test_data<O: OffsetSizeTrait>(
444        path: impl AsRef<std::path::Path>,
445        arr: &[&GenericStringArray<O>],
446    ) -> Result<usize> {
447        let mut writer = tokio::fs::File::create(path).await?;
448        // Write some garbage to reset "tell()".
449        writer.write_all(b"1234").await.unwrap();
450        let mut encoder = BinaryEncoder::new(&mut writer);
451
452        let arrs = arr.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
453        let pos = encoder.encode(arrs.as_slice()).await.unwrap();
454        AsyncWriteExt::shutdown(&mut writer).await.unwrap();
455        Ok(pos)
456    }
457
458    async fn test_round_trips<O: OffsetSizeTrait>(arrs: &[&GenericStringArray<O>]) {
459        let path = TempStdFile::default();
460
461        let pos = write_test_data(&path, arrs).await.unwrap();
462
463        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
464            .await
465            .unwrap();
466        let read_len = arrs.iter().map(|a| a.len()).sum();
467        let decoder =
468            BinaryDecoder::<GenericStringType<O>>::new(reader.as_ref(), pos, read_len, true);
469        let actual_arr = decoder.decode().await.unwrap();
470
471        let arrs_ref = arrs.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
472        let expected = concat(arrs_ref.as_slice()).unwrap();
473        assert_eq!(
474            actual_arr
475                .as_any()
476                .downcast_ref::<GenericStringArray<O>>()
477                .unwrap(),
478            expected
479                .as_any()
480                .downcast_ref::<GenericStringArray<O>>()
481                .unwrap(),
482        );
483    }
484
485    #[tokio::test]
486    async fn test_write_binary_data() {
487        test_round_trips(&[&StringArray::from(vec!["a", "b", "cd", "efg"])]).await;
488        test_round_trips(&[&StringArray::from(vec![Some("a"), None, Some("cd"), None])]).await;
489        test_round_trips(&[
490            &StringArray::from(vec![Some("a"), None, Some("cd"), None]),
491            &StringArray::from(vec![Some("f"), None, Some("gh"), None]),
492            &StringArray::from(vec![Some("t"), None, Some("uv"), None]),
493        ])
494        .await;
495        test_round_trips(&[&LargeStringArray::from(vec!["a", "b", "cd", "efg"])]).await;
496        test_round_trips(&[&LargeStringArray::from(vec![
497            Some("a"),
498            None,
499            Some("cd"),
500            None,
501        ])])
502        .await;
503        test_round_trips(&[
504            &LargeStringArray::from(vec![Some("a"), Some("b")]),
505            &LargeStringArray::from(vec![Some("c")]),
506            &LargeStringArray::from(vec![Some("d"), Some("e")]),
507        ])
508        .await;
509    }
510
511    #[tokio::test]
512    async fn test_write_binary_data_with_offset() {
513        let array: StringArray = StringArray::from(vec![Some("d"), Some("e")]).slice(1, 1);
514        test_round_trips(&[&array]).await;
515    }
516
517    #[tokio::test]
518    async fn test_range_query() {
519        let data = StringArray::from_iter_values(["a", "b", "c", "d", "e", "f", "g"]);
520
521        let path = TempStdFile::default();
522        let mut object_writer = tokio::fs::File::create(&path).await.unwrap();
523
524        // Write some garbage to reset "tell()".
525        object_writer.write_all(b"1234").await.unwrap();
526        let mut encoder = BinaryEncoder::new(&mut object_writer);
527        let pos = encoder.encode(&[&data]).await.unwrap();
528        AsyncWriteExt::shutdown(&mut object_writer).await.unwrap();
529
530        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
531            .await
532            .unwrap();
533        let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
534        assert_eq!(
535            decoder.decode().await.unwrap().as_ref(),
536            &StringArray::from_iter_values(["a", "b", "c", "d", "e", "f", "g"])
537        );
538
539        assert_eq!(
540            decoder.get(..).await.unwrap().as_ref(),
541            &StringArray::from_iter_values(["a", "b", "c", "d", "e", "f", "g"])
542        );
543
544        assert_eq!(
545            decoder.get(2..5).await.unwrap().as_ref(),
546            &StringArray::from_iter_values(["c", "d", "e"])
547        );
548
549        assert_eq!(
550            decoder.get(..5).await.unwrap().as_ref(),
551            &StringArray::from_iter_values(["a", "b", "c", "d", "e"])
552        );
553
554        assert_eq!(
555            decoder.get(4..).await.unwrap().as_ref(),
556            &StringArray::from_iter_values(["e", "f", "g"])
557        );
558        assert_eq!(
559            decoder.get(2..2).await.unwrap().as_ref(),
560            &new_empty_array(&DataType::Utf8)
561        );
562        assert!(decoder.get(100..100).await.is_err());
563    }
564
565    #[tokio::test]
566    async fn test_take() {
567        let data = StringArray::from_iter_values(["a", "b", "c", "d", "e", "f", "g"]);
568
569        let path = TempStdFile::default();
570
571        let pos = write_test_data(&path, &[&data]).await.unwrap();
572        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
573            .await
574            .unwrap();
575        let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
576
577        let actual = decoder
578            .take(&UInt32Array::from_iter_values([1, 2, 5]))
579            .await
580            .unwrap();
581        assert_eq!(
582            actual.as_ref(),
583            &StringArray::from_iter_values(["b", "c", "f"])
584        );
585    }
586
587    #[tokio::test]
588    async fn test_take_sparse_indices() {
589        let data = StringArray::from_iter_values((0..1000000).map(|v| format!("string-{v}")));
590
591        let path = TempStdFile::default();
592        let pos = write_test_data(&path, &[&data]).await.unwrap();
593        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
594            .await
595            .unwrap();
596        let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
597
598        let positions = decoder.get_positions(1..999998).await.unwrap();
599        let indices = UInt32Array::from_iter_values([1, 999998]);
600        let chunks = plan_take_chunks(positions.as_ref(), &indices, 64 * 1024).unwrap();
601        // Relative offset within the positions.
602        assert_eq!(chunks.len(), 2);
603        assert_eq!(chunks[0].indices, UInt32Array::from_iter_values([0]),);
604        assert_eq!(chunks[1].indices, UInt32Array::from_iter_values([999997]),);
605
606        let actual = decoder
607            .take(&UInt32Array::from_iter_values([1, 999998]))
608            .await
609            .unwrap();
610        assert_eq!(
611            actual.as_ref(),
612            &StringArray::from_iter_values(["string-1", "string-999998"])
613        );
614    }
615
616    #[tokio::test]
617    async fn test_take_dense_indices() {
618        let data = StringArray::from_iter_values((0..1000000).map(|v| format!("string-{v}")));
619
620        let path = TempStdFile::default();
621        let pos = write_test_data(&path, &[&data]).await.unwrap();
622
623        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
624            .await
625            .unwrap();
626        let decoder = BinaryDecoder::<Utf8Type>::new(reader.as_ref(), pos, data.len(), false);
627
628        let positions = decoder.get_positions(1..999998).await.unwrap();
629        let indices = UInt32Array::from_iter_values([
630            2, 3, 4, 1001, 1001, 1002, 2001, 2002, 2004, 3004, 3005,
631        ]);
632
633        let chunks = plan_take_chunks(positions.as_ref(), &indices, 1024).unwrap();
634        assert_eq!(chunks.len(), 4);
635        // A contiguous range.
636        assert_eq!(chunks[0].indices, UInt32Array::from_iter_values(0..3));
637        assert!(chunks[0].is_contiguous);
638        // Not contiguous because of repeats
639        assert_eq!(
640            chunks[1].indices,
641            UInt32Array::from_iter_values([999, 999, 1000])
642        );
643        assert!(!chunks[1].is_contiguous);
644        // Not contiguous because of gaps
645        assert_eq!(
646            chunks[2].indices,
647            UInt32Array::from_iter_values([1999, 2000, 2002])
648        );
649        assert!(!chunks[2].is_contiguous);
650        // Another contiguous range, this time after non-contiguous ones
651        assert_eq!(
652            chunks[3].indices,
653            UInt32Array::from_iter_values([3002, 3003])
654        );
655        assert!(chunks[3].is_contiguous);
656
657        let actual = decoder.take(&indices).await.unwrap();
658        assert_eq!(
659            actual.as_ref(),
660            &StringArray::from_iter_values(indices.values().iter().map(|v| format!("string-{v}")))
661        );
662    }
663
664    #[tokio::test]
665    async fn test_write_slice() {
666        let path = TempStdFile::default();
667        let data = StringArray::from_iter_values((0..100).map(|v| format!("abcdef-{v:#03}")));
668
669        let mut object_writer = tokio::fs::File::create(&path).await.unwrap();
670        let mut encoder = BinaryEncoder::new(&mut object_writer);
671        for i in 0..10 {
672            let pos = encoder.encode(&[&data.slice(i * 10, 10)]).await.unwrap();
673            assert_eq!(pos, (i * (8 * 11) /* offset array */ + (i + 1) * (10 * 10)));
674        }
675    }
676
677    #[tokio::test]
678    async fn test_write_binary_with_nulls() {
679        let data = BinaryArray::from_iter((0..60000).map(|v| {
680            if v % 4 != 0 {
681                Some::<&[u8]>(b"abcdefgh")
682            } else {
683                None
684            }
685        }));
686        let path = TempStdFile::default();
687
688        let pos = {
689            let mut object_writer = tokio::fs::File::create(&path).await.unwrap();
690
691            // Write some garbage to reset "tell()".
692            object_writer.write_all(b"1234").await.unwrap();
693            let mut encoder = BinaryEncoder::new(&mut object_writer);
694
695            // let arrs = arr.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
696            let pos = encoder.encode(&[&data]).await.unwrap();
697            AsyncWriteExt::shutdown(&mut object_writer).await.unwrap();
698            pos
699        };
700
701        let reader = LocalObjectReader::open_local_path(&path, 1024, None)
702            .await
703            .unwrap();
704        let decoder = BinaryDecoder::<BinaryType>::new(reader.as_ref(), pos, data.len(), true);
705        let idx = UInt32Array::from(vec![0_u32, 5_u32, 59996_u32]);
706        let actual = decoder.take(&idx).await.unwrap();
707        let values: Vec<Option<&[u8]>> = vec![None, Some(b"abcdefgh"), None];
708        assert_eq!(actual.as_binary::<i32>(), &BinaryArray::from(values));
709    }
710}
lance_file/versions/v1/encoding/binary.rs

lance_file/versions/v1/encoding/
binary.rs
