lance_encoding/encodings/logical/
blob.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{collections::HashMap, sync::Arc};
5
6use arrow_array::{
7    Array, ArrayRef, StructArray, UInt64Array,
8    builder::{PrimitiveBuilder, StringBuilder},
9    cast::AsArray,
10    types::{UInt8Type, UInt32Type, UInt64Type},
11};
12use arrow_buffer::Buffer;
13use arrow_schema::{DataType, Field as ArrowField, Fields};
14use futures::{FutureExt, future::BoxFuture};
15use lance_core::{
16    Error, Result,
17    datatypes::{BLOB_V2_DESC_FIELDS, BlobV2Layout, Field},
18    error::LanceOptionExt,
19};
20
21use crate::{
22    buffer::LanceBuffer,
23    constants::PACKED_STRUCT_META_KEY,
24    decoder::PageEncoding,
25    encoder::{EncodeTask, EncodedColumn, EncodedPage, FieldEncoder, OutOfLineBuffers},
26    format::EncodingUtils,
27    repdef::{DefinitionInterpretation, RepDefBuilder},
28};
29use lance_core::datatypes::BlobKind;
30
31/// Blob structural encoder - stores large binary data in external buffers
32///
33/// This encoder takes large binary arrays and stores them outside the normal
34/// page structure. It creates a descriptor (position, size) for each blob
35/// that is stored inline in the page.
36pub struct BlobStructuralEncoder {
37    // Encoder for the descriptors (position/size struct)
38    descriptor_encoder: Box<dyn FieldEncoder>,
39    // Set when we first see data
40    def_meaning: Option<Arc<[DefinitionInterpretation]>>,
41}
42
43impl BlobStructuralEncoder {
44    pub fn new(
45        field: &Field,
46        make_descriptor_encoder: impl FnOnce(Field) -> Result<Box<dyn FieldEncoder>>,
47    ) -> Result<Self> {
48        // Create descriptor field: struct<position: u64, size: u64>
49        // Preserve the original field's metadata for packed struct
50        let mut descriptor_metadata = HashMap::with_capacity(1);
51        descriptor_metadata.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
52
53        let descriptor_data_type = DataType::Struct(Fields::from(vec![
54            ArrowField::new("position", DataType::UInt64, false),
55            ArrowField::new("size", DataType::UInt64, false),
56        ]));
57
58        // Use the original field's name for the descriptor
59        let descriptor_field = Field::try_from(
60            ArrowField::new(&field.name, descriptor_data_type, field.nullable)
61                .with_metadata(descriptor_metadata),
62        )?;
63
64        // Use PrimitiveStructuralEncoder to handle the descriptor
65        let descriptor_encoder = make_descriptor_encoder(descriptor_field)?;
66
67        Ok(Self {
68            descriptor_encoder,
69            def_meaning: None,
70        })
71    }
72
73    fn wrap_tasks(
74        tasks: Vec<EncodeTask>,
75        def_meaning: Arc<[DefinitionInterpretation]>,
76    ) -> Vec<EncodeTask> {
77        tasks
78            .into_iter()
79            .map(|task| {
80                let def_meaning = def_meaning.clone();
81                task.then(|encoded_page| async move {
82                    let encoded_page = encoded_page?;
83
84                    let PageEncoding::Structural(inner_layout) = encoded_page.description else {
85                        return Err(Error::internal(
86                            "Expected inner encoding to return structural layout".to_string(),
87                        ));
88                    };
89
90                    let wrapped = EncodingUtils::blob_layout(inner_layout, &def_meaning);
91                    Ok(EncodedPage {
92                        column_idx: encoded_page.column_idx,
93                        data: encoded_page.data,
94                        description: PageEncoding::Structural(wrapped),
95                        num_rows: encoded_page.num_rows,
96                        row_number: encoded_page.row_number,
97                    })
98                })
99                .boxed()
100            })
101            .collect::<Vec<_>>()
102    }
103}
104
105impl FieldEncoder for BlobStructuralEncoder {
106    fn maybe_encode(
107        &mut self,
108        array: ArrayRef,
109        external_buffers: &mut OutOfLineBuffers,
110        mut repdef: RepDefBuilder,
111        row_number: u64,
112        num_rows: u64,
113    ) -> Result<Vec<EncodeTask>> {
114        if let Some(validity) = array.nulls() {
115            repdef.add_validity_bitmap(validity.clone());
116        } else {
117            repdef.add_no_null(array.len());
118        }
119
120        // Convert input array to LargeBinary
121        let binary_array = array.as_binary_opt::<i64>().ok_or_else(|| {
122            Error::invalid_input_source(
123                format!("Expected LargeBinary array, got {}", array.data_type()).into(),
124            )
125        })?;
126
127        let repdef = RepDefBuilder::serialize(vec![repdef]);
128
129        let rep = repdef.repetition_levels.as_ref();
130        let def = repdef.definition_levels.as_ref();
131        let def_meaning: Arc<[DefinitionInterpretation]> = repdef.def_meaning.into();
132
133        // A blob page stores one definition interpretation for all of its rows.
134        // The descriptor encoder can buffer multiple input arrays, so finish the
135        // pending page before a later array changes from all-valid to nullable (or
136        // vice versa).
137        let mut encode_tasks = match self.def_meaning.as_ref() {
138            Some(existing) if existing != &def_meaning => {
139                let existing = existing.clone();
140                Self::wrap_tasks(self.descriptor_encoder.flush(external_buffers)?, existing)
141            }
142            _ => Vec::new(),
143        };
144        self.def_meaning = Some(def_meaning.clone());
145
146        // Collect positions and sizes
147        let mut positions = Vec::with_capacity(binary_array.len());
148        let mut sizes = Vec::with_capacity(binary_array.len());
149
150        for i in 0..binary_array.len() {
151            if binary_array.is_null(i) {
152                // Null values are smuggled into the positions array
153
154                // If we have null values we must have definition levels
155                let mut repdef = (def.expect_ok()?[i] as u64) << 16;
156                if let Some(rep) = rep {
157                    repdef += rep[i] as u64;
158                }
159
160                debug_assert_ne!(repdef, 0);
161                positions.push(repdef);
162                sizes.push(0);
163            } else {
164                let value = binary_array.value(i);
165                if value.is_empty() {
166                    // Empty values
167                    positions.push(0);
168                    sizes.push(0);
169                } else {
170                    // Add data to external buffers
171                    let position =
172                        external_buffers.add_buffer(LanceBuffer::from(Buffer::from(value)));
173                    positions.push(position);
174                    sizes.push(value.len() as u64);
175                }
176            }
177        }
178
179        // Create descriptor array
180        let position_array = Arc::new(UInt64Array::from(positions));
181        let size_array = Arc::new(UInt64Array::from(sizes));
182        let descriptor_array = Arc::new(StructArray::new(
183            Fields::from(vec![
184                ArrowField::new("position", DataType::UInt64, false),
185                ArrowField::new("size", DataType::UInt64, false),
186            ]),
187            vec![position_array as ArrayRef, size_array as ArrayRef],
188            None, // Descriptors are never null
189        ));
190
191        // Delegate to descriptor encoder
192        let descriptor_tasks = self.descriptor_encoder.maybe_encode(
193            descriptor_array,
194            external_buffers,
195            RepDefBuilder::default(),
196            row_number,
197            num_rows,
198        )?;
199        encode_tasks.extend(Self::wrap_tasks(descriptor_tasks, def_meaning));
200
201        Ok(encode_tasks)
202    }
203
204    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
205        let encode_tasks = self.descriptor_encoder.flush(external_buffers)?;
206
207        // Use the cached def meaning.  If we haven't seen any data yet then we can just use a dummy
208        // value (not clear there would be any encode tasks in that case)
209        let def_meaning = self
210            .def_meaning
211            .clone()
212            .unwrap_or_else(|| Arc::new([DefinitionInterpretation::AllValidItem]));
213
214        Ok(Self::wrap_tasks(encode_tasks, def_meaning))
215    }
216
217    fn finish(
218        &mut self,
219        external_buffers: &mut OutOfLineBuffers,
220    ) -> BoxFuture<'_, Result<Vec<EncodedColumn>>> {
221        self.descriptor_encoder.finish(external_buffers)
222    }
223
224    fn num_columns(&self) -> u32 {
225        self.descriptor_encoder.num_columns()
226    }
227}
228
229/// Blob v2 structural encoder
230pub struct BlobV2StructuralEncoder {
231    descriptor_encoder: Box<dyn FieldEncoder>,
232}
233
234impl BlobV2StructuralEncoder {
235    pub fn new(
236        field: &Field,
237        make_descriptor_encoder: impl FnOnce(Field) -> Result<Box<dyn FieldEncoder>>,
238    ) -> Result<Self> {
239        let mut descriptor_metadata = HashMap::with_capacity(1);
240        descriptor_metadata.insert(PACKED_STRUCT_META_KEY.to_string(), "true".to_string());
241
242        let descriptor_data_type = DataType::Struct(BLOB_V2_DESC_FIELDS.clone());
243
244        let descriptor_field = Field::try_from(
245            ArrowField::new(&field.name, descriptor_data_type, field.nullable)
246                .with_metadata(descriptor_metadata),
247        )?;
248
249        let descriptor_encoder = make_descriptor_encoder(descriptor_field)?;
250
251        Ok(Self { descriptor_encoder })
252    }
253}
254
255impl FieldEncoder for BlobV2StructuralEncoder {
256    fn maybe_encode(
257        &mut self,
258        array: ArrayRef,
259        external_buffers: &mut OutOfLineBuffers,
260        repdef: RepDefBuilder,
261        row_number: u64,
262        num_rows: u64,
263    ) -> Result<Vec<EncodeTask>> {
264        let struct_arr = array
265            .as_any()
266            .downcast_ref::<StructArray>()
267            .ok_or_else(|| {
268                Error::invalid_input_source(
269                    format!(
270                        "Blob v2 encoder expected StructArray, got {}",
271                        array.data_type()
272                    )
273                    .into(),
274                )
275            })?;
276        if BlobV2Layout::classify(struct_arr.fields()) != Some(BlobV2Layout::Prepared) {
277            let actual = BlobV2Layout::classify(struct_arr.fields())
278                .map(|layout| layout.to_string())
279                .unwrap_or_else(|| format!("unrecognized ({:?})", struct_arr.fields()));
280            return Err(Error::invalid_input_source(
281                format!("Blob v2 encoder expected prepared array layout, got {actual} layout")
282                    .into(),
283            ));
284        }
285
286        let kind_col = struct_arr
287            .column_by_name("kind")
288            .ok_or_else(|| {
289                Error::invalid_input_source("Blob v2 struct missing `kind` field".into())
290            })?
291            .as_primitive::<UInt8Type>();
292        let data_col = struct_arr
293            .column_by_name("data")
294            .ok_or_else(|| {
295                Error::invalid_input_source("Blob v2 struct missing `data` field".into())
296            })?
297            .as_binary::<i64>();
298        let uri_col = struct_arr
299            .column_by_name("uri")
300            .ok_or_else(|| {
301                Error::invalid_input_source("Blob v2 struct missing `uri` field".into())
302            })?
303            .as_string::<i32>();
304        let blob_id_col = struct_arr
305            .column_by_name("blob_id")
306            .ok_or_else(|| {
307                Error::invalid_input_source("Blob v2 struct missing `blob_id` field".into())
308            })?
309            .as_primitive::<UInt32Type>();
310        let blob_size_col = struct_arr
311            .column_by_name("blob_size")
312            .ok_or_else(|| {
313                Error::invalid_input_source("Blob v2 struct missing `blob_size` field".into())
314            })?
315            .as_primitive::<UInt64Type>();
316        let packed_position_col = struct_arr
317            .column_by_name("position")
318            .ok_or_else(|| {
319                Error::invalid_input_source("Blob v2 struct missing `position` field".into())
320            })?
321            .as_primitive::<UInt64Type>();
322
323        let row_count = struct_arr.len();
324
325        let mut kind_builder = PrimitiveBuilder::<UInt8Type>::with_capacity(row_count);
326        let mut position_builder = PrimitiveBuilder::<UInt64Type>::with_capacity(row_count);
327        let mut size_builder = PrimitiveBuilder::<UInt64Type>::with_capacity(row_count);
328        let mut blob_id_builder = PrimitiveBuilder::<UInt32Type>::with_capacity(row_count);
329        let mut uri_builder = StringBuilder::with_capacity(row_count, row_count * 16);
330
331        for i in 0..row_count {
332            let (kind_value, position_value, size_value, blob_id_value, uri_value) = if struct_arr
333                .is_null(i)
334                || kind_col.is_null(i)
335            {
336                (BlobKind::Inline as u8, 0, 0, 0, "".to_string())
337            } else {
338                let kind_val = BlobKind::try_from(kind_col.value(i))?;
339                match kind_val {
340                    BlobKind::Managed | BlobKind::ManagedWithBase => {
341                        if uri_col.is_null(i)
342                            || ((kind_val == BlobKind::ManagedWithBase) != blob_id_col.is_valid(i))
343                            || packed_position_col.is_null(i)
344                            || blob_size_col.is_null(i)
345                        {
346                            return Err(Error::invalid_input(format!(
347                                "Managed blob row {i} requires URI, position, and size, plus a base ID for an explicit base"
348                            )));
349                        }
350                        let uri = uri_col.value(i);
351                        let position = packed_position_col.value(i);
352                        let size = blob_size_col.value(i);
353                        lance_core::utils::blob::validate_managed_reference(uri, position, size)?;
354                        (
355                            kind_val as u8,
356                            position,
357                            size,
358                            if kind_val == BlobKind::ManagedWithBase {
359                                blob_id_col.value(i)
360                            } else {
361                                0
362                            },
363                            uri.to_string(),
364                        )
365                    }
366                    BlobKind::Dedicated => (
367                        BlobKind::Dedicated as u8,
368                        0,
369                        blob_size_col.value(i),
370                        blob_id_col.value(i),
371                        "".to_string(),
372                    ),
373                    BlobKind::External => {
374                        let uri = uri_col.value(i).to_string();
375                        let position = if packed_position_col.is_null(i) {
376                            0
377                        } else {
378                            packed_position_col.value(i)
379                        };
380                        let size = if blob_size_col.is_null(i) {
381                            0
382                        } else {
383                            blob_size_col.value(i)
384                        };
385                        let external_base_id = if blob_id_col.is_null(i) {
386                            0
387                        } else {
388                            blob_id_col.value(i)
389                        };
390                        (
391                            BlobKind::External as u8,
392                            position,
393                            size,
394                            external_base_id,
395                            uri,
396                        )
397                    }
398                    BlobKind::Packed => (
399                        BlobKind::Packed as u8,
400                        packed_position_col.value(i),
401                        blob_size_col.value(i),
402                        blob_id_col.value(i),
403                        "".to_string(),
404                    ),
405                    BlobKind::Inline => {
406                        let data_val = data_col.value(i);
407                        let blob_len = data_val.len() as u64;
408                        let position =
409                            external_buffers.add_buffer(LanceBuffer::from(Buffer::from(data_val)));
410
411                        (
412                            BlobKind::Inline as u8,
413                            position,
414                            blob_len,
415                            0,
416                            "".to_string(),
417                        )
418                    }
419                }
420            };
421
422            kind_builder.append_value(kind_value);
423            position_builder.append_value(position_value);
424            size_builder.append_value(size_value);
425            blob_id_builder.append_value(blob_id_value);
426            uri_builder.append_value(uri_value);
427        }
428        let children: Vec<ArrayRef> = vec![
429            Arc::new(kind_builder.finish()),
430            Arc::new(position_builder.finish()),
431            Arc::new(size_builder.finish()),
432            Arc::new(blob_id_builder.finish()),
433            Arc::new(uri_builder.finish()),
434        ];
435
436        let descriptor_array = Arc::new(StructArray::try_new(
437            BLOB_V2_DESC_FIELDS.clone(),
438            children,
439            struct_arr.nulls().cloned(),
440        )?) as ArrayRef;
441
442        self.descriptor_encoder.maybe_encode(
443            descriptor_array,
444            external_buffers,
445            repdef,
446            row_number,
447            num_rows,
448        )
449    }
450
451    fn flush(&mut self, external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
452        self.descriptor_encoder.flush(external_buffers)
453    }
454
455    fn finish(
456        &mut self,
457        external_buffers: &mut OutOfLineBuffers,
458    ) -> BoxFuture<'_, Result<Vec<EncodedColumn>>> {
459        self.descriptor_encoder.finish(external_buffers)
460    }
461
462    fn num_columns(&self) -> u32 {
463        self.descriptor_encoder.num_columns()
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470    use crate::{
471        encoder::{ColumnIndexSequence, EncodingOptions},
472        testing::{
473            TestCases, TestEncoding, check_round_trip_encoding_of_data,
474            check_round_trip_encoding_of_data_with_expected, create_test_field_encoder,
475            test_encoding_strategy,
476        },
477    };
478    use arrow_array::{
479        ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt8Array, UInt32Array, UInt64Array,
480    };
481    use arrow_schema::{DataType, Field as ArrowField};
482    use lance_core::datatypes::BLOB_V2_LOGICAL_MINIMAL_FIELDS;
483
484    #[test]
485    fn test_blob_encoder_creation() {
486        let field = Field::try_from(
487            ArrowField::new("blob_field", DataType::LargeBinary, true).with_metadata(
488                HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]),
489            ),
490        )
491        .unwrap();
492        let mut column_index = ColumnIndexSequence::default();
493        let options = EncodingOptions::default();
494        let strategy = test_encoding_strategy(TestEncoding::StructuralU16);
495
496        let encoder =
497            create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options);
498
499        assert!(encoder.is_ok());
500    }
501
502    #[test]
503    fn test_blob_v2_encoder_rejects_logical_array_layout() {
504        let field = Field::try_from(
505            ArrowField::new(
506                "blob_field",
507                DataType::Struct(BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone()),
508                true,
509            )
510            .with_metadata(HashMap::from([(
511                lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
512                lance_arrow::BLOB_V2_EXT_NAME.to_string(),
513            )])),
514        )
515        .unwrap();
516        let mut column_index = ColumnIndexSequence::default();
517        let options = EncodingOptions::default();
518        let strategy = test_encoding_strategy(TestEncoding::StructuralU32);
519        let mut encoder =
520            create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options)
521                .unwrap();
522        let array = Arc::new(
523            StructArray::try_new(
524                BLOB_V2_LOGICAL_MINIMAL_FIELDS.clone(),
525                vec![
526                    Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_ref())])) as ArrayRef,
527                    Arc::new(StringArray::from(vec![None::<&str>])) as ArrayRef,
528                ],
529                None,
530            )
531            .unwrap(),
532        ) as ArrayRef;
533        let mut external_buffers = OutOfLineBuffers::new(0, 8);
534        let Err(error) =
535            encoder.maybe_encode(array, &mut external_buffers, RepDefBuilder::default(), 0, 1)
536        else {
537            panic!("logical array layout unexpectedly reached the descriptor encoder");
538        };
539        assert!(matches!(error, Error::InvalidInput { .. }));
540        assert!(
541            error
542                .to_string()
543                .contains("expected prepared array layout, got logical layout")
544        );
545    }
546
547    #[tokio::test]
548    async fn test_blob_encoding_simple() {
549        let field = Field::try_from(
550            ArrowField::new("blob_field", DataType::LargeBinary, true).with_metadata(
551                HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]),
552            ),
553        )
554        .unwrap();
555        let mut column_index = ColumnIndexSequence::default();
556        let options = EncodingOptions::default();
557        let strategy = test_encoding_strategy(TestEncoding::StructuralU16);
558
559        let mut encoder =
560            create_test_field_encoder(strategy.as_ref(), &field, &mut column_index, &options)
561                .unwrap();
562
563        // Create test data with larger blobs
564        let large_data = vec![0u8; 1024 * 100]; // 100KB blob
565        let data: Vec<Option<&[u8]>> =
566            vec![Some(b"hello world"), None, Some(&large_data), Some(b"")];
567        let array = Arc::new(LargeBinaryArray::from(data));
568
569        // Test encoding
570        let mut external_buffers = OutOfLineBuffers::new(0, 8);
571        let repdef = RepDefBuilder::default();
572
573        let tasks = encoder
574            .maybe_encode(array, &mut external_buffers, repdef, 0, 4)
575            .unwrap();
576
577        // If no tasks yet, flush to force encoding
578        if tasks.is_empty() {
579            let _flush_tasks = encoder.flush(&mut external_buffers).unwrap();
580        }
581
582        // Should produce encode tasks for the descriptor (or we need more data)
583        // For now, just verify no errors occurred
584        assert!(encoder.num_columns() > 0);
585
586        // Verify external buffers were used for large data
587        let buffers = external_buffers.take_buffers();
588        assert!(
589            !buffers.is_empty(),
590            "Large blobs should be stored in external buffers"
591        );
592    }
593
594    #[tokio::test]
595    async fn test_blob_round_trip() {
596        // Test round-trip encoding with blob metadata
597        let blob_metadata =
598            HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]);
599
600        // Create test data
601        let val1: &[u8] = &vec![1u8; 1024]; // 1KB
602        let val2: &[u8] = &vec![2u8; 10240]; // 10KB
603        let val3: &[u8] = &vec![3u8; 102400]; // 100KB
604        let array = Arc::new(LargeBinaryArray::from(vec![
605            Some(val1),
606            None,
607            Some(val2),
608            Some(val3),
609        ]));
610
611        // Use the standard test harness
612        check_round_trip_encoding_of_data(
613            vec![array],
614            &TestCases::default().with_array_and_u16_encodings(),
615            blob_metadata,
616        )
617        .await;
618    }
619
620    #[tokio::test]
621    async fn test_blob_round_trip_empty_values() {
622        // Empty values share size == 0 with nulls in the descriptor layout
623        // and schedule no read; each must decode to zero-length bytes without
624        // consuming the read result of a following non-empty blob. Empties
625        // are placed before payloads so a misassignment corrupts the output
626        // instead of only exhausting the read iterator.
627        let blob_metadata =
628            HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]);
629
630        let val1: &[u8] = &vec![1u8; 1024];
631        let val2: &[u8] = &vec![2u8; 10240];
632        let empty: &[u8] = &[];
633        let array = Arc::new(LargeBinaryArray::from(vec![
634            Some(empty),
635            Some(val1),
636            None,
637            Some(empty),
638            Some(val2),
639            None,
640            Some(empty),
641        ]));
642
643        check_round_trip_encoding_of_data(
644            vec![array],
645            &TestCases::default().with_encodings([
646                TestEncoding::Array,
647                TestEncoding::StructuralU16,
648                TestEncoding::StructuralU32,
649            ]),
650            blob_metadata,
651        )
652        .await;
653    }
654
655    #[tokio::test]
656    async fn test_blob_round_trip_varying_chunk_nullability() {
657        let blob_metadata =
658            HashMap::from([(lance_arrow::BLOB_META_KEY.to_string(), "true".to_string())]);
659        let all_valid = Arc::new(LargeBinaryArray::from(vec![Some(b"first".as_ref())]));
660        let with_null = Arc::new(LargeBinaryArray::from(vec![
661            Some(b"second".as_ref()),
662            None,
663            Some(b"".as_ref()),
664        ]));
665        let all_valid_again = Arc::new(LargeBinaryArray::from(vec![Some(b"last".as_ref())]));
666
667        check_round_trip_encoding_of_data(
668            vec![all_valid, with_null, all_valid_again],
669            &TestCases::default().with_encoding(TestEncoding::StructuralU16),
670            blob_metadata,
671        )
672        .await;
673    }
674
675    #[tokio::test]
676    async fn test_blob_v2_external_round_trip() {
677        let blob_metadata = HashMap::from([(
678            lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
679            lance_arrow::BLOB_V2_EXT_NAME.to_string(),
680        )]);
681
682        let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true));
683        let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true));
684        let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true));
685        let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true));
686        let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true));
687        let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true));
688
689        let kind_array = UInt8Array::from(vec![
690            BlobKind::Inline as u8,
691            BlobKind::External as u8,
692            BlobKind::External as u8,
693        ]);
694        let data_array = LargeBinaryArray::from(vec![Some(b"inline".as_ref()), None, None]);
695        let uri_array = StringArray::from(vec![
696            None,
697            Some("file:///tmp/external.bin"),
698            Some("s3://bucket/blob"),
699        ]);
700        let blob_id_array = UInt32Array::from(vec![0, 0, 0]);
701        let blob_size_array = UInt64Array::from(vec![0, 0, 0]);
702        let position_array = UInt64Array::from(vec![0, 0, 0]);
703
704        let struct_array = StructArray::from(vec![
705            (kind_field, Arc::new(kind_array) as ArrayRef),
706            (data_field, Arc::new(data_array) as ArrayRef),
707            (uri_field, Arc::new(uri_array) as ArrayRef),
708            (blob_id_field, Arc::new(blob_id_array) as ArrayRef),
709            (blob_size_field, Arc::new(blob_size_array) as ArrayRef),
710            (position_field, Arc::new(position_array) as ArrayRef),
711        ]);
712
713        let expected_descriptor = StructArray::from(vec![
714            (
715                Arc::new(ArrowField::new("kind", DataType::UInt8, false)),
716                Arc::new(UInt8Array::from(vec![
717                    BlobKind::Inline as u8,
718                    BlobKind::External as u8,
719                    BlobKind::External as u8,
720                ])) as ArrayRef,
721            ),
722            (
723                Arc::new(ArrowField::new("position", DataType::UInt64, false)),
724                Arc::new(UInt64Array::from(vec![0, 0, 0])) as ArrayRef,
725            ),
726            (
727                Arc::new(ArrowField::new("size", DataType::UInt64, false)),
728                Arc::new(UInt64Array::from(vec![6, 0, 0])) as ArrayRef,
729            ),
730            (
731                Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)),
732                Arc::new(UInt32Array::from(vec![0, 0, 0])) as ArrayRef,
733            ),
734            (
735                Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)),
736                Arc::new(StringArray::from(vec![
737                    "",
738                    "file:///tmp/external.bin",
739                    "s3://bucket/blob",
740                ])) as ArrayRef,
741            ),
742        ]);
743
744        check_round_trip_encoding_of_data_with_expected(
745            vec![Arc::new(struct_array)],
746            Some(Arc::new(expected_descriptor)),
747            &TestCases::default().with_u32_structural_encodings(),
748            blob_metadata,
749        )
750        .await;
751    }
752
753    #[tokio::test]
754    async fn test_blob_v2_dedicated_round_trip() {
755        let blob_metadata = HashMap::from([(
756            lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
757            lance_arrow::BLOB_V2_EXT_NAME.to_string(),
758        )]);
759
760        let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true));
761        let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true));
762        let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true));
763        let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true));
764        let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true));
765        let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true));
766
767        let kind_array = UInt8Array::from(vec![BlobKind::Dedicated as u8, BlobKind::Inline as u8]);
768        let data_array = LargeBinaryArray::from(vec![None, Some(b"abc".as_ref())]);
769        let uri_array = StringArray::from(vec![Option::<&str>::None, None]);
770        let blob_id_array = UInt32Array::from(vec![42, 0]);
771        let blob_size_array = UInt64Array::from(vec![12, 0]);
772        let position_array = UInt64Array::from(vec![0, 0]);
773
774        let struct_array = StructArray::from(vec![
775            (kind_field, Arc::new(kind_array) as ArrayRef),
776            (data_field, Arc::new(data_array) as ArrayRef),
777            (uri_field, Arc::new(uri_array) as ArrayRef),
778            (blob_id_field, Arc::new(blob_id_array) as ArrayRef),
779            (blob_size_field, Arc::new(blob_size_array) as ArrayRef),
780            (position_field, Arc::new(position_array) as ArrayRef),
781        ]);
782
783        let expected_descriptor = StructArray::from(vec![
784            (
785                Arc::new(ArrowField::new("kind", DataType::UInt8, false)),
786                Arc::new(UInt8Array::from(vec![
787                    BlobKind::Dedicated as u8,
788                    BlobKind::Inline as u8,
789                ])) as ArrayRef,
790            ),
791            (
792                Arc::new(ArrowField::new("position", DataType::UInt64, false)),
793                Arc::new(UInt64Array::from(vec![0, 0])) as ArrayRef,
794            ),
795            (
796                Arc::new(ArrowField::new("size", DataType::UInt64, false)),
797                Arc::new(UInt64Array::from(vec![12, 3])) as ArrayRef,
798            ),
799            (
800                Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)),
801                Arc::new(UInt32Array::from(vec![42, 0])) as ArrayRef,
802            ),
803            (
804                Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)),
805                Arc::new(StringArray::from(vec!["", ""])) as ArrayRef,
806            ),
807        ]);
808
809        check_round_trip_encoding_of_data_with_expected(
810            vec![Arc::new(struct_array)],
811            Some(Arc::new(expected_descriptor)),
812            &TestCases::default().with_u32_structural_encodings(),
813            blob_metadata,
814        )
815        .await;
816    }
817
818    #[tokio::test]
819    async fn test_blob_v2_external_with_range_round_trip() {
820        let blob_metadata = HashMap::from([(
821            lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
822            lance_arrow::BLOB_V2_EXT_NAME.to_string(),
823        )]);
824
825        let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true));
826        let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true));
827        let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true));
828        let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true));
829        let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true));
830        let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true));
831
832        let kind_array = UInt8Array::from(vec![BlobKind::External as u8]);
833        let data_array = LargeBinaryArray::from(vec![None::<&[u8]>]);
834        let uri_array = StringArray::from(vec![Some("memory://container.pack")]);
835        let blob_id_array = UInt32Array::from(vec![0]);
836        let blob_size_array = UInt64Array::from(vec![42]);
837        let position_array = UInt64Array::from(vec![7]);
838
839        let struct_array = StructArray::from(vec![
840            (kind_field, Arc::new(kind_array) as ArrayRef),
841            (data_field, Arc::new(data_array) as ArrayRef),
842            (uri_field, Arc::new(uri_array) as ArrayRef),
843            (blob_id_field, Arc::new(blob_id_array) as ArrayRef),
844            (blob_size_field, Arc::new(blob_size_array) as ArrayRef),
845            (position_field, Arc::new(position_array) as ArrayRef),
846        ]);
847
848        let expected_descriptor = StructArray::from(vec![
849            (
850                Arc::new(ArrowField::new("kind", DataType::UInt8, false)),
851                Arc::new(UInt8Array::from(vec![BlobKind::External as u8])) as ArrayRef,
852            ),
853            (
854                Arc::new(ArrowField::new("position", DataType::UInt64, false)),
855                Arc::new(UInt64Array::from(vec![7])) as ArrayRef,
856            ),
857            (
858                Arc::new(ArrowField::new("size", DataType::UInt64, false)),
859                Arc::new(UInt64Array::from(vec![42])) as ArrayRef,
860            ),
861            (
862                Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)),
863                Arc::new(UInt32Array::from(vec![0])) as ArrayRef,
864            ),
865            (
866                Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)),
867                Arc::new(StringArray::from(vec!["memory://container.pack"])) as ArrayRef,
868            ),
869        ]);
870
871        check_round_trip_encoding_of_data_with_expected(
872            vec![Arc::new(struct_array)],
873            Some(Arc::new(expected_descriptor)),
874            &TestCases::default().with_u32_structural_encodings(),
875            blob_metadata,
876        )
877        .await;
878    }
879
880    #[tokio::test]
881    async fn test_blob_v2_packed_round_trip() {
882        let blob_metadata = HashMap::from([(
883            lance_arrow::ARROW_EXT_NAME_KEY.to_string(),
884            lance_arrow::BLOB_V2_EXT_NAME.to_string(),
885        )]);
886
887        let kind_field = Arc::new(ArrowField::new("kind", DataType::UInt8, true));
888        let data_field = Arc::new(ArrowField::new("data", DataType::LargeBinary, true));
889        let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, true));
890        let blob_id_field = Arc::new(ArrowField::new("blob_id", DataType::UInt32, true));
891        let blob_size_field = Arc::new(ArrowField::new("blob_size", DataType::UInt64, true));
892        let position_field = Arc::new(ArrowField::new("position", DataType::UInt64, true));
893
894        let kind_array = UInt8Array::from(vec![BlobKind::Packed as u8]);
895        let data_array = LargeBinaryArray::from(vec![None::<&[u8]>]);
896        let uri_array = StringArray::from(vec![None::<&str>]);
897        let blob_id_array = UInt32Array::from(vec![7]);
898        let blob_size_array = UInt64Array::from(vec![5]);
899        let position_array = UInt64Array::from(vec![10]);
900
901        let struct_array = StructArray::from(vec![
902            (kind_field, Arc::new(kind_array) as ArrayRef),
903            (data_field, Arc::new(data_array) as ArrayRef),
904            (uri_field, Arc::new(uri_array) as ArrayRef),
905            (blob_id_field, Arc::new(blob_id_array) as ArrayRef),
906            (blob_size_field, Arc::new(blob_size_array) as ArrayRef),
907            (position_field, Arc::new(position_array) as ArrayRef),
908        ]);
909
910        let expected_descriptor = StructArray::from(vec![
911            (
912                Arc::new(ArrowField::new("kind", DataType::UInt8, false)),
913                Arc::new(UInt8Array::from(vec![BlobKind::Packed as u8])) as ArrayRef,
914            ),
915            (
916                Arc::new(ArrowField::new("position", DataType::UInt64, false)),
917                Arc::new(UInt64Array::from(vec![10])) as ArrayRef,
918            ),
919            (
920                Arc::new(ArrowField::new("size", DataType::UInt64, false)),
921                Arc::new(UInt64Array::from(vec![5])) as ArrayRef,
922            ),
923            (
924                Arc::new(ArrowField::new("blob_id", DataType::UInt32, false)),
925                Arc::new(UInt32Array::from(vec![7])) as ArrayRef,
926            ),
927            (
928                Arc::new(ArrowField::new("blob_uri", DataType::Utf8, false)),
929                Arc::new(StringArray::from(vec![""])) as ArrayRef,
930            ),
931        ]);
932
933        check_round_trip_encoding_of_data_with_expected(
934            vec![Arc::new(struct_array)],
935            Some(Arc::new(expected_descriptor)),
936            &TestCases::default().with_u32_structural_encodings(),
937            blob_metadata,
938        )
939        .await;
940    }
941}
lance_encoding/encodings/logical/blob.rs

lance_encoding/encodings/logical/
blob.rs
