1use std::{
5 any::Any,
6 collections::{HashMap, VecDeque},
7 env,
8 fmt::Debug,
9 iter,
10 ops::Range,
11 sync::Arc,
12 vec,
13};
14
15use crate::{
16 constants::{
17 STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
18 },
19 data::DictionaryDataBlock,
20 encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler},
21 format::{
22 ProtobufUtils21,
23 pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression},
24 },
25};
26use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type};
27use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
28use arrow_schema::{DataType, Field as ArrowField};
29use bytes::Bytes;
30use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
31use itertools::Itertools;
32use lance_arrow::DataTypeExt;
33use lance_arrow::deepcopy::deep_copy_nulls;
34use lance_core::{
35 cache::{CacheKey, Context, DeepSizeOf},
36 error::{Error, LanceOptionExt},
37 utils::bit::pad_bytes,
38};
39use log::trace;
40
41use crate::{
42 compression::{
43 BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
44 },
45 data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
46 utils::bytepack::BytepackedIntegerEncoder,
47};
48use crate::{
49 compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
50 encodings::logical::primitive::fullzip::PerValueDataBlock,
51};
52use crate::{
53 encodings::logical::primitive::miniblock::MiniBlockChunk, utils::bytepack::ByteUnpacker,
54};
55use crate::{
56 encodings::logical::primitive::miniblock::MiniBlockCompressed,
57 statistics::{ComputeStat, GetStat, Stat},
58};
59use crate::{
60 repdef::{
61 CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
62 RepDefSlicer, build_control_word_iterator,
63 },
64 utils::accumulation::AccumulationQueue,
65};
66use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
67
68use crate::constants::{
69 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
70 DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
71 DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
72 DICT_VALUES_COMPRESSION_META_KEY,
73};
74use crate::version::LanceFileVersion;
75use crate::{
76 EncodingsIo,
77 buffer::LanceBuffer,
78 data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
79 decoder::{
80 ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
81 MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
82 StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
83 StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
84 },
85 encoder::{
86 EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
87 },
88 repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
89};
90
91pub mod blob;
92pub mod constant;
93pub mod dict;
94pub mod fullzip;
95pub mod miniblock;
96
97const FILL_BYTE: u8 = 0xFE;
98const DEFAULT_DICT_DIVISOR: u64 = 2;
99const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
100const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
101const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
102
103struct PageLoadTask {
104 decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
105 num_rows: u64,
106}
107
108trait StructuralPageScheduler: std::fmt::Debug + Send {
111 fn initialize<'a>(
113 &'a mut self,
114 io: &Arc<dyn EncodingsIo>,
115 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
116 fn load(&mut self, data: &Arc<dyn CachedPageData>);
118 fn schedule_ranges(
127 &self,
128 ranges: &[Range<u64>],
129 io: &Arc<dyn EncodingsIo>,
130 ) -> Result<Vec<PageLoadTask>>;
131}
132
133#[derive(Debug)]
135struct ChunkMeta {
136 num_values: u64,
137 chunk_size_bytes: u64,
138 offset_bytes: u64,
139}
140
141#[derive(Debug, Clone)]
143struct DecodedMiniBlockChunk {
144 rep: Option<ScalarBuffer<u16>>,
145 def: Option<ScalarBuffer<u16>>,
146 values: DataBlock,
147}
148
149#[derive(Debug)]
157struct DecodeMiniBlockTask {
158 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
159 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
160 value_decompressor: Arc<dyn MiniBlockDecompressor>,
161 dictionary_data: Option<Arc<DataBlock>>,
162 def_meaning: Arc<[DefinitionInterpretation]>,
163 num_buffers: u64,
164 max_visible_level: u16,
165 instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
166 has_large_chunk: bool,
167}
168
169impl DecodeMiniBlockTask {
170 fn decode_levels(
171 rep_decompressor: &dyn BlockDecompressor,
172 levels: LanceBuffer,
173 num_levels: u16,
174 ) -> Result<ScalarBuffer<u16>> {
175 let rep = rep_decompressor.decompress(levels, num_levels as u64)?;
176 let rep = rep.as_fixed_width().unwrap();
177 debug_assert_eq!(rep.num_values, num_levels as u64);
178 debug_assert_eq!(rep.bits_per_value, 16);
179 Ok(rep.data.borrow_to_typed_slice::<u16>())
180 }
181
182 fn extend_levels(
189 range: Range<u64>,
190 levels: &mut Option<LevelBuffer>,
191 level_buf: &Option<impl AsRef<[u16]>>,
192 dest_offset: usize,
193 ) {
194 if let Some(level_buf) = level_buf {
195 if levels.is_none() {
196 let mut new_levels_vec =
199 LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
200 new_levels_vec.extend(iter::repeat_n(0, dest_offset));
201 *levels = Some(new_levels_vec);
202 }
203 levels.as_mut().unwrap().extend(
204 level_buf.as_ref()[range.start as usize..range.end as usize]
205 .iter()
206 .copied(),
207 );
208 } else if let Some(levels) = levels {
209 let num_values = (range.end - range.start) as usize;
210 levels.extend(iter::repeat_n(0, num_values));
213 }
214 }
215
216 fn map_range(
253 range: Range<u64>,
254 rep: Option<&impl AsRef<[u16]>>,
255 def: Option<&impl AsRef<[u16]>>,
256 max_rep: u16,
257 max_visible_def: u16,
258 total_items: u64,
261 preamble_action: PreambleAction,
262 ) -> (Range<u64>, Range<u64>) {
263 if let Some(rep) = rep {
264 let mut rep = rep.as_ref();
265 let mut items_in_preamble = 0_u64;
268 let first_row_start = match preamble_action {
269 PreambleAction::Skip | PreambleAction::Take => {
270 let first_row_start = if let Some(def) = def.as_ref() {
271 let mut first_row_start = None;
272 for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
273 if *rep == max_rep {
274 first_row_start = Some(idx as u64);
275 break;
276 }
277 if *def <= max_visible_def {
278 items_in_preamble += 1;
279 }
280 }
281 first_row_start
282 } else {
283 let first_row_start =
284 rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
285 items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
286 first_row_start
287 };
288 if first_row_start.is_none() {
291 assert!(preamble_action == PreambleAction::Take);
292 return (0..total_items, 0..rep.len() as u64);
293 }
294 let first_row_start = first_row_start.unwrap();
295 rep = &rep[first_row_start as usize..];
296 first_row_start
297 }
298 PreambleAction::Absent => {
299 debug_assert!(rep[0] == max_rep);
300 0
301 }
302 };
303
304 if range.start == range.end {
306 debug_assert!(preamble_action == PreambleAction::Take);
307 debug_assert!(items_in_preamble <= total_items);
308 return (0..items_in_preamble, 0..first_row_start);
309 }
310 assert!(range.start < range.end);
311
312 let mut rows_seen = 0;
313 let mut new_start = 0;
314 let mut new_levels_start = 0;
315
316 if let Some(def) = def {
317 let def = &def.as_ref()[first_row_start as usize..];
318
319 let mut lead_invis_seen = 0;
321
322 if range.start > 0 {
323 if def[0] > max_visible_def {
324 lead_invis_seen += 1;
325 }
326 for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
327 if *rep == max_rep {
328 rows_seen += 1;
329 if rows_seen == range.start {
330 new_start = idx as u64 + 1 - lead_invis_seen;
331 new_levels_start = idx as u64 + 1;
332 break;
333 }
334 }
335 if *def > max_visible_def {
336 lead_invis_seen += 1;
337 }
338 }
339 }
340
341 rows_seen += 1;
342
343 let mut new_end = u64::MAX;
344 let mut new_levels_end = rep.len() as u64;
345 let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
346 let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
347 for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
348 .iter()
349 .zip(&def[(new_levels_start + 1) as usize..])
350 .enumerate()
351 {
352 if *rep == max_rep {
353 rows_seen += 1;
354 if rows_seen == range.end + 1 {
355 new_end = idx as u64 + new_start + 1 - tail_invis_seen;
356 new_levels_end = idx as u64 + new_levels_start + 1;
357 break;
358 }
359 }
360 if *def > max_visible_def {
361 tail_invis_seen += 1;
362 }
363 }
364
365 if new_end == u64::MAX {
366 new_levels_end = rep.len() as u64;
367 let total_invis_seen = lead_invis_seen + tail_invis_seen;
368 new_end = rep.len() as u64 - total_invis_seen;
369 }
370
371 assert_ne!(new_end, u64::MAX);
372
373 if preamble_action == PreambleAction::Skip {
375 new_start += items_in_preamble;
376 new_end += items_in_preamble;
377 new_levels_start += first_row_start;
378 new_levels_end += first_row_start;
379 } else if preamble_action == PreambleAction::Take {
380 debug_assert_eq!(new_start, 0);
381 debug_assert_eq!(new_levels_start, 0);
382 new_end += items_in_preamble;
383 new_levels_end += first_row_start;
384 }
385
386 debug_assert!(new_end <= total_items);
387 (new_start..new_end, new_levels_start..new_levels_end)
388 } else {
389 if range.start > 0 {
395 for (idx, rep) in rep.iter().skip(1).enumerate() {
396 if *rep == max_rep {
397 rows_seen += 1;
398 if rows_seen == range.start {
399 new_start = idx as u64 + 1;
400 break;
401 }
402 }
403 }
404 }
405 let mut new_end = rep.len() as u64;
406 if range.end < total_items {
408 for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
409 if *rep == max_rep {
410 rows_seen += 1;
411 if rows_seen == range.end {
412 new_end = idx as u64 + new_start + 1;
413 break;
414 }
415 }
416 }
417 }
418
419 if preamble_action == PreambleAction::Skip {
421 new_start += first_row_start;
422 new_end += first_row_start;
423 } else if preamble_action == PreambleAction::Take {
424 debug_assert_eq!(new_start, 0);
425 new_end += first_row_start;
426 }
427
428 debug_assert!(new_end <= total_items);
429 (new_start..new_end, new_start..new_end)
430 }
431 } else {
432 (range.clone(), range)
435 }
436 }
437
438 fn read_buffer_sizes<const LARGE: bool>(
440 buf: &[u8],
441 offset: &mut usize,
442 num_buffers: u64,
443 ) -> Vec<u32> {
444 let read_size = if LARGE { 4 } else { 2 };
445 (0..num_buffers)
446 .map(|_| {
447 let bytes = &buf[*offset..*offset + read_size];
448 let size = if LARGE {
449 u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
450 } else {
451 u16::from_le_bytes([bytes[0], bytes[1]]) as u32
453 };
454 *offset += read_size;
455 size
456 })
457 .collect()
458 }
459
460 fn decode_miniblock_chunk(
462 &self,
463 buf: &LanceBuffer,
464 items_in_chunk: u64,
465 ) -> Result<DecodedMiniBlockChunk> {
466 let mut offset = 0;
467 let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
468 offset += 2;
469
470 let rep_size = if self.rep_decompressor.is_some() {
471 let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
472 offset += 2;
473 Some(rep_size)
474 } else {
475 None
476 };
477 let def_size = if self.def_decompressor.is_some() {
478 let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
479 offset += 2;
480 Some(def_size)
481 } else {
482 None
483 };
484
485 let buffer_sizes = if self.has_large_chunk {
486 Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
487 } else {
488 Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
489 };
490
491 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
492
493 let rep = rep_size.map(|rep_size| {
494 let rep = buf.slice_with_length(offset, rep_size as usize);
495 offset += rep_size as usize;
496 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
497 rep
498 });
499
500 let def = def_size.map(|def_size| {
501 let def = buf.slice_with_length(offset, def_size as usize);
502 offset += def_size as usize;
503 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
504 def
505 });
506
507 let buffers = buffer_sizes
508 .into_iter()
509 .map(|buf_size| {
510 let buf = buf.slice_with_length(offset, buf_size as usize);
511 offset += buf_size as usize;
512 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
513 buf
514 })
515 .collect::<Vec<_>>();
516
517 let values = self
518 .value_decompressor
519 .decompress(buffers, items_in_chunk)?;
520
521 let rep = rep
522 .map(|rep| {
523 Self::decode_levels(
524 self.rep_decompressor.as_ref().unwrap().as_ref(),
525 rep,
526 num_levels,
527 )
528 })
529 .transpose()?;
530 let def = def
531 .map(|def| {
532 Self::decode_levels(
533 self.def_decompressor.as_ref().unwrap().as_ref(),
534 def,
535 num_levels,
536 )
537 })
538 .transpose()?;
539
540 Ok(DecodedMiniBlockChunk { rep, def, values })
541 }
542}
543
544impl DecodePageTask for DecodeMiniBlockTask {
545 fn decode(self: Box<Self>) -> Result<DecodedPage> {
546 let mut repbuf: Option<LevelBuffer> = None;
548 let mut defbuf: Option<LevelBuffer> = None;
549
550 let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
551
552 let estimated_size_bytes = self
554 .instructions
555 .iter()
556 .map(|(_, chunk)| chunk.data.len())
557 .sum::<usize>()
558 * 2;
559 let mut data_builder =
560 DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
561
562 let mut level_offset = 0;
564
565 let needs_caching: Vec<bool> = self
567 .instructions
568 .windows(2)
569 .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
570 .chain(std::iter::once(false)) .collect();
572
573 let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
575
576 for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
578 let should_cache_this_chunk = needs_caching[idx];
579
580 let decoded_chunk = match &chunk_cache {
581 Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
582 cached_chunk.clone()
584 }
585 _ => {
586 let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
588
589 if should_cache_this_chunk {
591 chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
592 }
593 decoded
594 }
595 };
596
597 let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
598
599 let row_range_start =
601 instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
602 let row_range_end = row_range_start + instructions.rows_to_take;
603
604 let (item_range, level_range) = Self::map_range(
606 row_range_start..row_range_end,
607 rep.as_ref(),
608 def.as_ref(),
609 max_rep,
610 self.max_visible_level,
611 chunk.items_in_chunk,
612 instructions.preamble_action,
613 );
614 if item_range.end - item_range.start > chunk.items_in_chunk {
615 return Err(lance_core::Error::internal(format!(
616 "Item range {:?} is greater than chunk items in chunk {:?}",
617 item_range, chunk.items_in_chunk
618 )));
619 }
620
621 Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
623 Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
624 level_offset += (level_range.end - level_range.start) as usize;
625 data_builder.append(&values, item_range);
626 }
627
628 let mut data = data_builder.finish();
629
630 let unraveler =
631 RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
632
633 if let Some(dictionary) = &self.dictionary_data {
634 let DataBlock::FixedWidth(indices) = data else {
636 return Err(lance_core::Error::internal(format!(
637 "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
638 data
639 )));
640 };
641 data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
642 indices,
643 dictionary.as_ref().clone(),
644 ));
645 }
646
647 Ok(DecodedPage {
648 data,
649 repdef: unraveler,
650 })
651 }
652}
653
654#[derive(Debug)]
657struct LoadedChunk {
658 data: LanceBuffer,
659 items_in_chunk: u64,
660 byte_range: Range<u64>,
661 chunk_idx: usize,
662}
663
664impl Clone for LoadedChunk {
665 fn clone(&self) -> Self {
666 Self {
667 data: self.data.clone(),
669 items_in_chunk: self.items_in_chunk,
670 byte_range: self.byte_range.clone(),
671 chunk_idx: self.chunk_idx,
672 }
673 }
674}
675
676#[derive(Debug)]
679struct MiniBlockDecoder {
680 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
681 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
682 value_decompressor: Arc<dyn MiniBlockDecompressor>,
683 def_meaning: Arc<[DefinitionInterpretation]>,
684 loaded_chunks: VecDeque<LoadedChunk>,
685 instructions: VecDeque<ChunkInstructions>,
686 offset_in_current_chunk: u64,
687 num_rows: u64,
688 num_buffers: u64,
689 dictionary: Option<Arc<DataBlock>>,
690 has_large_chunk: bool,
691}
692
693impl StructuralPageDecoder for MiniBlockDecoder {
696 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
697 let mut items_desired = num_rows;
698 let mut need_preamble = false;
699 let mut skip_in_chunk = self.offset_in_current_chunk;
700 let mut drain_instructions = Vec::new();
701 while items_desired > 0 || need_preamble {
702 let (instructions, consumed) = self
703 .instructions
704 .front()
705 .unwrap()
706 .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
707
708 while self.loaded_chunks.front().unwrap().chunk_idx
709 != instructions.chunk_instructions.chunk_idx
710 {
711 self.loaded_chunks.pop_front();
712 }
713 drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
714 if consumed {
715 self.instructions.pop_front();
716 }
717 }
718 self.offset_in_current_chunk = skip_in_chunk;
721
722 let max_visible_level = self
723 .def_meaning
724 .iter()
725 .take_while(|l| !l.is_list())
726 .map(|l| l.num_def_levels())
727 .sum::<u16>();
728
729 Ok(Box::new(DecodeMiniBlockTask {
730 instructions: drain_instructions,
731 def_decompressor: self.def_decompressor.clone(),
732 rep_decompressor: self.rep_decompressor.clone(),
733 value_decompressor: self.value_decompressor.clone(),
734 dictionary_data: self.dictionary.clone(),
735 def_meaning: self.def_meaning.clone(),
736 num_buffers: self.num_buffers,
737 max_visible_level,
738 has_large_chunk: self.has_large_chunk,
739 }))
740 }
741
742 fn num_rows(&self) -> u64 {
743 self.num_rows
744 }
745}
746
747#[derive(Debug)]
748struct CachedComplexAllNullState {
749 rep: Option<ScalarBuffer<u16>>,
750 def: Option<ScalarBuffer<u16>>,
751}
752
753impl DeepSizeOf for CachedComplexAllNullState {
754 fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
755 self.rep.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
756 + self.def.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
757 }
758}
759
760impl CachedPageData for CachedComplexAllNullState {
761 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
762 self
763 }
764}
765
766#[derive(Debug)]
775pub struct ComplexAllNullScheduler {
776 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
778 def_meaning: Arc<[DefinitionInterpretation]>,
779 repdef: Option<Arc<CachedComplexAllNullState>>,
780 max_rep: u16,
781 max_visible_level: u16,
782 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
783 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
784 num_rep_values: u64,
785 num_def_values: u64,
786}
787
788impl ComplexAllNullScheduler {
789 pub fn new(
790 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
791 def_meaning: Arc<[DefinitionInterpretation]>,
792 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
793 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
794 num_rep_values: u64,
795 num_def_values: u64,
796 ) -> Self {
797 let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
798 let max_visible_level = def_meaning
799 .iter()
800 .take_while(|l| !l.is_list())
801 .map(|l| l.num_def_levels())
802 .sum::<u16>();
803 Self {
804 buffer_offsets_and_sizes,
805 def_meaning,
806 repdef: None,
807 max_rep,
808 max_visible_level,
809 rep_decompressor,
810 def_decompressor,
811 num_rep_values,
812 num_def_values,
813 }
814 }
815}
816
817impl StructuralPageScheduler for ComplexAllNullScheduler {
818 fn initialize<'a>(
819 &'a mut self,
820 io: &Arc<dyn EncodingsIo>,
821 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
822 let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
824 let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
825 let has_rep = rep_size > 0;
826 let has_def = def_size > 0;
827
828 let mut reads = Vec::with_capacity(2);
829 if has_rep {
830 reads.push(rep_pos..rep_pos + rep_size);
831 }
832 if has_def {
833 reads.push(def_pos..def_pos + def_size);
834 }
835
836 let data = io.submit_request(reads, 0);
837 let rep_decompressor = self.rep_decompressor.clone();
838 let def_decompressor = self.def_decompressor.clone();
839 let num_rep_values = self.num_rep_values;
840 let num_def_values = self.num_def_values;
841
842 async move {
843 let data = data.await?;
844 let mut data_iter = data.into_iter();
845
846 let decompress_levels = |compressed_bytes: Bytes,
847 decompressor: &Arc<dyn BlockDecompressor>,
848 num_values: u64,
849 level_type: &str|
850 -> Result<ScalarBuffer<u16>> {
851 let compressed_buffer = LanceBuffer::from_bytes(compressed_bytes, 1);
852 let decompressed = decompressor.decompress(compressed_buffer, num_values)?;
853 match decompressed {
854 DataBlock::FixedWidth(block) => {
855 if block.num_values != num_values {
856 return Err(Error::invalid_input_source(format!(
857 "Unexpected {} level count after decompression: expected {}, got {}",
858 level_type, num_values, block.num_values
859 )
860 .into()));
861 }
862 if block.bits_per_value != 16 {
863 return Err(Error::invalid_input_source(format!(
864 "Unexpected {} level bit width after decompression: expected 16, got {}",
865 level_type, block.bits_per_value
866 )
867 .into()));
868 }
869 Ok(block.data.borrow_to_typed_slice::<u16>())
870 }
871 _ => Err(Error::invalid_input_source(format!(
872 "Expected fixed-width data block for {} levels",
873 level_type
874 )
875 .into())),
876 }
877 };
878
879 let rep = if has_rep {
880 let rep = data_iter.next().unwrap();
881 if let Some(rep_decompressor) = rep_decompressor.as_ref() {
882 Some(decompress_levels(
883 rep,
884 rep_decompressor,
885 num_rep_values,
886 "repetition",
887 )?)
888 } else {
889 let rep = LanceBuffer::from_bytes(rep, 2);
890 let rep = rep.borrow_to_typed_slice::<u16>();
891 Some(rep)
892 }
893 } else {
894 None
895 };
896
897 let def = if has_def {
898 let def = data_iter.next().unwrap();
899 if let Some(def_decompressor) = def_decompressor.as_ref() {
900 Some(decompress_levels(
901 def,
902 def_decompressor,
903 num_def_values,
904 "definition",
905 )?)
906 } else {
907 let def = LanceBuffer::from_bytes(def, 2);
908 let def = def.borrow_to_typed_slice::<u16>();
909 Some(def)
910 }
911 } else {
912 None
913 };
914
915 let repdef = Arc::new(CachedComplexAllNullState { rep, def });
916
917 self.repdef = Some(repdef.clone());
918
919 Ok(repdef as Arc<dyn CachedPageData>)
920 }
921 .boxed()
922 }
923
924 fn load(&mut self, data: &Arc<dyn CachedPageData>) {
925 self.repdef = Some(
926 data.clone()
927 .as_arc_any()
928 .downcast::<CachedComplexAllNullState>()
929 .unwrap(),
930 );
931 }
932
933 fn schedule_ranges(
934 &self,
935 ranges: &[Range<u64>],
936 _io: &Arc<dyn EncodingsIo>,
937 ) -> Result<Vec<PageLoadTask>> {
938 let ranges = VecDeque::from_iter(ranges.iter().cloned());
939 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
940 let decoder = Box::new(ComplexAllNullPageDecoder {
941 ranges,
942 rep: self.repdef.as_ref().unwrap().rep.clone(),
943 def: self.repdef.as_ref().unwrap().def.clone(),
944 num_rows,
945 def_meaning: self.def_meaning.clone(),
946 max_rep: self.max_rep,
947 max_visible_level: self.max_visible_level,
948 cursor_row: 0,
949 cursor_level: 0,
950 }) as Box<dyn StructuralPageDecoder>;
951 let page_load_task = PageLoadTask {
952 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
953 num_rows,
954 };
955 Ok(vec![page_load_task])
956 }
957}
958
959#[derive(Debug)]
960pub struct ComplexAllNullPageDecoder {
961 ranges: VecDeque<Range<u64>>,
962 rep: Option<ScalarBuffer<u16>>,
963 def: Option<ScalarBuffer<u16>>,
964 num_rows: u64,
965 def_meaning: Arc<[DefinitionInterpretation]>,
966 max_rep: u16,
967 max_visible_level: u16,
968 cursor_row: u64,
969 cursor_level: usize,
970}
971
972impl ComplexAllNullPageDecoder {
973 fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
974 let mut rows_desired = num_rows;
975 let mut ranges = Vec::with_capacity(self.ranges.len());
976 while rows_desired > 0 {
977 let front = self.ranges.front_mut().unwrap();
978 let avail = front.end - front.start;
979 if avail > rows_desired {
980 ranges.push(front.start..front.start + rows_desired);
981 front.start += rows_desired;
982 rows_desired = 0;
983 } else {
984 ranges.push(self.ranges.pop_front().unwrap());
985 rows_desired -= avail;
986 }
987 }
988 ranges
989 }
990
991 fn take_row(&mut self) -> Result<(Range<usize>, u64)> {
992 let start = self.cursor_level;
993 let end = if let Some(rep) = &self.rep {
994 if start >= rep.len() {
995 return Err(Error::internal(
996 "Invalid complex all-null layout: repetition buffer too short",
997 ));
998 }
999 if rep[start] != self.max_rep {
1000 return Err(Error::internal(
1001 "Invalid complex all-null layout: row did not start at max repetition level",
1002 ));
1003 }
1004 let mut end = start + 1;
1005 while end < rep.len() && rep[end] != self.max_rep {
1006 end += 1;
1007 }
1008 end
1009 } else {
1010 start + 1
1011 };
1012
1013 let visible = if let Some(def) = &self.def {
1014 if end > def.len() {
1015 return Err(Error::internal(
1016 "Invalid complex all-null layout: definition buffer too short",
1017 ));
1018 }
1019 def[start..end]
1020 .iter()
1021 .filter(|d| **d <= self.max_visible_level)
1022 .count() as u64
1023 } else {
1024 (end - start) as u64
1025 };
1026
1027 self.cursor_level = end;
1028 self.cursor_row += 1;
1029 Ok((start..end, visible))
1030 }
1031
1032 fn skip_to_row(&mut self, target_row: u64) -> Result<()> {
1033 while self.cursor_row < target_row {
1034 self.take_row()?;
1035 }
1036 Ok(())
1037 }
1038}
1039
1040impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1041 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1042 let drained_ranges = self.drain_ranges(num_rows);
1043 let mut level_slices: Vec<Range<usize>> = Vec::new();
1044 let mut visible_items_total = 0;
1045
1046 for range in drained_ranges {
1047 self.skip_to_row(range.start)?;
1048 for _ in range.start..range.end {
1049 let (level_range, visible) = self.take_row()?;
1050 visible_items_total += visible;
1051 if let Some(last) = level_slices.last_mut()
1052 && last.end == level_range.start
1053 {
1054 last.end = level_range.end;
1055 continue;
1056 }
1057 level_slices.push(level_range);
1058 }
1059 }
1060
1061 Ok(Box::new(DecodeComplexAllNullTask {
1062 level_slices,
1063 visible_items_total,
1064 rep: self.rep.clone(),
1065 def: self.def.clone(),
1066 def_meaning: self.def_meaning.clone(),
1067 max_visible_level: self.max_visible_level,
1068 }))
1069 }
1070
1071 fn num_rows(&self) -> u64 {
1072 self.num_rows
1073 }
1074}
1075
1076#[derive(Debug)]
1079pub struct DecodeComplexAllNullTask {
1080 level_slices: Vec<Range<usize>>,
1081 visible_items_total: u64,
1082 rep: Option<ScalarBuffer<u16>>,
1083 def: Option<ScalarBuffer<u16>>,
1084 def_meaning: Arc<[DefinitionInterpretation]>,
1085 max_visible_level: u16,
1086}
1087
1088impl DecodeComplexAllNullTask {
1089 fn decode_level(&self, levels: &Option<ScalarBuffer<u16>>) -> Option<Vec<u16>> {
1090 levels.as_ref().map(|levels| {
1091 let num_levels = self
1092 .level_slices
1093 .iter()
1094 .map(|range| range.end - range.start)
1095 .sum();
1096 let mut referenced_levels = Vec::with_capacity(num_levels);
1097 for range in &self.level_slices {
1098 referenced_levels.extend(levels[range.start..range.end].iter().copied());
1099 }
1100 referenced_levels
1101 })
1102 }
1103}
1104
1105impl DecodePageTask for DecodeComplexAllNullTask {
1106 fn decode(self: Box<Self>) -> Result<DecodedPage> {
1107 let rep = self.decode_level(&self.rep);
1108 let def = self.decode_level(&self.def);
1109
1110 let num_values = if let Some(def) = &def {
1114 def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
1115 } else {
1116 self.visible_items_total
1117 };
1118
1119 let data = DataBlock::AllNull(AllNullDataBlock { num_values });
1120 let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
1121 Ok(DecodedPage {
1122 data,
1123 repdef: unraveler,
1124 })
1125 }
1126}
1127
1128#[derive(Debug, Default)]
1133pub struct SimpleAllNullScheduler {}
1134
1135impl StructuralPageScheduler for SimpleAllNullScheduler {
1136 fn initialize<'a>(
1137 &'a mut self,
1138 _io: &Arc<dyn EncodingsIo>,
1139 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1140 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
1141 }
1142
1143 fn load(&mut self, _cache: &Arc<dyn CachedPageData>) {}
1144
1145 fn schedule_ranges(
1146 &self,
1147 ranges: &[Range<u64>],
1148 _io: &Arc<dyn EncodingsIo>,
1149 ) -> Result<Vec<PageLoadTask>> {
1150 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1151 let decoder =
1152 Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
1153 let page_load_task = PageLoadTask {
1154 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1155 num_rows,
1156 };
1157 Ok(vec![page_load_task])
1158 }
1159}
1160
1161#[derive(Debug)]
1164struct SimpleAllNullDecodePageTask {
1165 num_values: u64,
1166}
1167impl DecodePageTask for SimpleAllNullDecodePageTask {
1168 fn decode(self: Box<Self>) -> Result<DecodedPage> {
1169 let unraveler = RepDefUnraveler::new(
1170 None,
1171 Some(vec![1; self.num_values as usize]),
1172 Arc::new([DefinitionInterpretation::NullableItem]),
1173 self.num_values,
1174 );
1175 Ok(DecodedPage {
1176 data: DataBlock::AllNull(AllNullDataBlock {
1177 num_values: self.num_values,
1178 }),
1179 repdef: unraveler,
1180 })
1181 }
1182}
1183
1184#[derive(Debug)]
1185pub struct SimpleAllNullPageDecoder {
1186 num_rows: u64,
1187}
1188
1189impl StructuralPageDecoder for SimpleAllNullPageDecoder {
1190 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1191 Ok(Box::new(SimpleAllNullDecodePageTask {
1192 num_values: num_rows,
1193 }))
1194 }
1195
1196 fn num_rows(&self) -> u64 {
1197 self.num_rows
1198 }
1199}
1200
1201#[derive(Debug, Clone)]
1202struct MiniBlockSchedulerDictionary {
1203 dictionary_decompressor: Arc<dyn BlockDecompressor>,
1205 dictionary_buf_position_and_size: (u64, u64),
1206 dictionary_data_alignment: u64,
1207 num_dictionary_items: u64,
1208}
1209
1210#[derive(Debug)]
1212struct MiniBlockRepIndexBlock {
1213 first_row: u64,
1217 starts_including_trailer: u64,
1220 has_preamble: bool,
1222 has_trailer: bool,
1224}
1225
1226impl DeepSizeOf for MiniBlockRepIndexBlock {
1227 fn deep_size_of_children(&self, _context: &mut Context) -> usize {
1228 0
1229 }
1230}
1231
1232#[derive(Debug)]
1237struct MiniBlockRepIndex {
1238 blocks: Vec<MiniBlockRepIndexBlock>,
1239}
1240
1241impl DeepSizeOf for MiniBlockRepIndex {
1242 fn deep_size_of_children(&self, context: &mut Context) -> usize {
1243 self.blocks.deep_size_of_children(context)
1244 }
1245}
1246
1247impl MiniBlockRepIndex {
1248 pub fn default_from_chunks(chunks: &[ChunkMeta]) -> Self {
1253 let mut blocks = Vec::with_capacity(chunks.len());
1254 let mut offset: u64 = 0;
1255
1256 for c in chunks {
1257 blocks.push(MiniBlockRepIndexBlock {
1258 first_row: offset,
1259 starts_including_trailer: c.num_values,
1260 has_preamble: false,
1261 has_trailer: false,
1262 });
1263
1264 offset += c.num_values;
1265 }
1266
1267 Self { blocks }
1268 }
1269
1270 pub fn decode_from_bytes(rep_bytes: &[u8], stride: usize) -> Self {
1276 let buffer = crate::buffer::LanceBuffer::from(rep_bytes.to_vec());
1278 let u64_slice = buffer.borrow_to_typed_slice::<u64>();
1279 let n = u64_slice.len() / stride;
1280
1281 let mut blocks = Vec::with_capacity(n);
1282 let mut chunk_has_preamble = false;
1283 let mut offset: u64 = 0;
1284
1285 for i in 0..n {
1287 let base_idx = i * stride;
1288 let ends = u64_slice[base_idx];
1289 let partial = u64_slice[base_idx + 1];
1290
1291 let has_trailer = partial > 0;
1292 let starts_including_trailer =
1294 ends + (has_trailer as u64) - (chunk_has_preamble as u64);
1295
1296 blocks.push(MiniBlockRepIndexBlock {
1297 first_row: offset,
1298 starts_including_trailer,
1299 has_preamble: chunk_has_preamble,
1300 has_trailer,
1301 });
1302
1303 chunk_has_preamble = has_trailer;
1304 offset += starts_including_trailer;
1305 }
1306
1307 Self { blocks }
1308 }
1309}
1310
1311#[derive(Debug)]
1313struct MiniBlockCacheableState {
1314 chunk_meta: Vec<ChunkMeta>,
1316 rep_index: MiniBlockRepIndex,
1318 dictionary: Option<Arc<DataBlock>>,
1320}
1321
1322impl DeepSizeOf for MiniBlockCacheableState {
1323 fn deep_size_of_children(&self, context: &mut Context) -> usize {
1324 self.rep_index.deep_size_of_children(context)
1325 + self
1326 .dictionary
1327 .as_ref()
1328 .map(|dict| dict.data_size() as usize)
1329 .unwrap_or(0)
1330 }
1331}
1332
1333impl CachedPageData for MiniBlockCacheableState {
1334 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1335 self
1336 }
1337}
1338
1339#[derive(Debug)]
1366pub struct MiniBlockScheduler {
1367 buffer_offsets_and_sizes: Vec<(u64, u64)>,
1369 priority: u64,
1370 items_in_page: u64,
1371 repetition_index_depth: u16,
1372 num_buffers: u64,
1373 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
1374 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
1375 value_decompressor: Arc<dyn MiniBlockDecompressor>,
1376 def_meaning: Arc<[DefinitionInterpretation]>,
1377 dictionary: Option<MiniBlockSchedulerDictionary>,
1378 page_meta: Option<Arc<MiniBlockCacheableState>>,
1380 has_large_chunk: bool,
1381}
1382
1383impl MiniBlockScheduler {
1384 fn try_new(
1385 buffer_offsets_and_sizes: &[(u64, u64)],
1386 priority: u64,
1387 items_in_page: u64,
1388 layout: &pb21::MiniBlockLayout,
1389 decompressors: &dyn DecompressionStrategy,
1390 ) -> Result<Self> {
1391 let rep_decompressor = layout
1392 .rep_compression
1393 .as_ref()
1394 .map(|rep_compression| {
1395 decompressors
1396 .create_block_decompressor(rep_compression)
1397 .map(Arc::from)
1398 })
1399 .transpose()?;
1400 let def_decompressor = layout
1401 .def_compression
1402 .as_ref()
1403 .map(|def_compression| {
1404 decompressors
1405 .create_block_decompressor(def_compression)
1406 .map(Arc::from)
1407 })
1408 .transpose()?;
1409 let def_meaning = layout
1410 .layers
1411 .iter()
1412 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
1413 .collect::<Vec<_>>();
1414 let value_decompressor = decompressors.create_miniblock_decompressor(
1415 layout.value_compression.as_ref().unwrap(),
1416 decompressors,
1417 )?;
1418
1419 let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
1420 let num_dictionary_items = layout.num_dictionary_items;
1421 let dictionary_decompressor = decompressors
1422 .create_block_decompressor(dictionary_encoding)?
1423 .into();
1424 let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
1425 {
1426 Compression::Variable(_) => 4,
1427 Compression::Flat(_) => 16,
1428 Compression::General(_) => 1,
1429 Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
1430 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
1431 }
1432 _ => {
1433 return Err(Error::invalid_input_source(
1434 format!(
1435 "Unsupported mini-block dictionary encoding: {:?}",
1436 dictionary_encoding.compression.as_ref().unwrap()
1437 )
1438 .into(),
1439 ));
1440 }
1441 };
1442 Some(MiniBlockSchedulerDictionary {
1443 dictionary_decompressor,
1444 dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
1445 dictionary_data_alignment,
1446 num_dictionary_items,
1447 })
1448 } else {
1449 None
1450 };
1451
1452 Ok(Self {
1453 buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
1454 rep_decompressor,
1455 def_decompressor,
1456 value_decompressor: value_decompressor.into(),
1457 repetition_index_depth: layout.repetition_index_depth as u16,
1458 num_buffers: layout.num_buffers,
1459 priority,
1460 items_in_page,
1461 dictionary,
1462 def_meaning: def_meaning.into(),
1463 page_meta: None,
1464 has_large_chunk: layout.has_large_chunk,
1465 })
1466 }
1467
1468 fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
1469 let page_meta = self.page_meta.as_ref().unwrap();
1470 chunk_indices
1471 .iter()
1472 .map(|&chunk_idx| {
1473 let chunk_meta = &page_meta.chunk_meta[chunk_idx];
1474 let bytes_start = chunk_meta.offset_bytes;
1475 let bytes_end = bytes_start + chunk_meta.chunk_size_bytes;
1476 LoadedChunk {
1477 byte_range: bytes_start..bytes_end,
1478 items_in_chunk: chunk_meta.num_values,
1479 chunk_idx,
1480 data: LanceBuffer::empty(),
1481 }
1482 })
1483 .collect()
1484 }
1485}
1486
1487#[derive(Debug, PartialEq, Eq, Clone, Copy)]
1488enum PreambleAction {
1489 Take,
1490 Skip,
1491 Absent,
1492}
1493
1494#[derive(Clone, Debug, PartialEq, Eq)]
1529struct ChunkInstructions {
1530 chunk_idx: usize,
1532 preamble: PreambleAction,
1538 rows_to_skip: u64,
1542 rows_to_take: u64,
1545 take_trailer: bool,
1552}
1553
1554#[derive(Debug, PartialEq, Eq)]
1572struct ChunkDrainInstructions {
1573 chunk_instructions: ChunkInstructions,
1574 rows_to_skip: u64,
1575 rows_to_take: u64,
1576 preamble_action: PreambleAction,
1577}
1578
1579impl ChunkInstructions {
1580 fn schedule_instructions(
1586 rep_index: &MiniBlockRepIndex,
1587 user_ranges: &[Range<u64>],
1588 ) -> Vec<Self> {
1589 let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
1593
1594 for user_range in user_ranges {
1595 let mut rows_needed = user_range.end - user_range.start;
1596 let mut need_preamble = false;
1597
1598 let mut block_index = match rep_index
1601 .blocks
1602 .binary_search_by_key(&user_range.start, |block| block.first_row)
1603 {
1604 Ok(idx) => {
1605 let mut idx = idx;
1608 while idx > 0 && rep_index.blocks[idx - 1].first_row == user_range.start {
1609 idx -= 1;
1610 }
1611 idx
1612 }
1613 Err(idx) => idx - 1,
1615 };
1616
1617 let mut to_skip = user_range.start - rep_index.blocks[block_index].first_row;
1618
1619 while rows_needed > 0 || need_preamble {
1620 if block_index >= rep_index.blocks.len() {
1622 log::warn!(
1623 "schedule_instructions inconsistency: block_index >= rep_index.blocks.len(), exiting early"
1624 );
1625 break;
1626 }
1627
1628 let chunk = &rep_index.blocks[block_index];
1629 let rows_avail = chunk.starts_including_trailer.saturating_sub(to_skip);
1630
1631 if rows_avail == 0 && to_skip == 0 {
1635 if chunk.has_preamble && need_preamble {
1637 chunk_instructions.push(Self {
1638 chunk_idx: block_index,
1639 preamble: PreambleAction::Take,
1640 rows_to_skip: 0,
1641 rows_to_take: 0,
1642 take_trailer: chunk.has_trailer,
1646 });
1647 if chunk.starts_including_trailer > 0
1651 || block_index == rep_index.blocks.len() - 1
1652 {
1653 need_preamble = false;
1654 }
1655 }
1656 block_index += 1;
1658 continue;
1659 }
1660
1661 if rows_avail == 0 && to_skip > 0 {
1665 to_skip -= chunk.starts_including_trailer;
1668 block_index += 1;
1669 continue;
1670 }
1671
1672 let rows_to_take = rows_avail.min(rows_needed);
1673 rows_needed -= rows_to_take;
1674
1675 let mut take_trailer = false;
1676 let preamble = if chunk.has_preamble {
1677 if need_preamble {
1678 PreambleAction::Take
1679 } else {
1680 PreambleAction::Skip
1681 }
1682 } else {
1683 PreambleAction::Absent
1684 };
1685
1686 if rows_to_take == rows_avail && chunk.has_trailer {
1688 take_trailer = true;
1689 need_preamble = true;
1690 } else {
1691 need_preamble = false;
1692 };
1693
1694 chunk_instructions.push(Self {
1695 preamble,
1696 chunk_idx: block_index,
1697 rows_to_skip: to_skip,
1698 rows_to_take,
1699 take_trailer,
1700 });
1701
1702 to_skip = 0;
1703 block_index += 1;
1704 }
1705 }
1706
1707 if user_ranges.len() > 1 {
1711 let mut merged_instructions = Vec::with_capacity(chunk_instructions.len());
1713 let mut instructions_iter = chunk_instructions.into_iter();
1714 merged_instructions.push(instructions_iter.next().unwrap());
1715 for instruction in instructions_iter {
1716 let last = merged_instructions.last_mut().unwrap();
1717 if last.chunk_idx == instruction.chunk_idx
1718 && last.rows_to_take + last.rows_to_skip == instruction.rows_to_skip
1719 {
1720 last.rows_to_take += instruction.rows_to_take;
1721 last.take_trailer |= instruction.take_trailer;
1722 } else {
1723 merged_instructions.push(instruction);
1724 }
1725 }
1726 merged_instructions
1727 } else {
1728 chunk_instructions
1729 }
1730 }
1731
1732 fn drain_from_instruction(
1733 &self,
1734 rows_desired: &mut u64,
1735 need_preamble: &mut bool,
1736 skip_in_chunk: &mut u64,
1737 ) -> (ChunkDrainInstructions, bool) {
1738 debug_assert!(!*need_preamble || *skip_in_chunk == 0);
1740 let rows_avail = self.rows_to_take - *skip_in_chunk;
1741 let has_preamble = self.preamble != PreambleAction::Absent;
1742 let preamble_action = match (*need_preamble, has_preamble) {
1743 (true, true) => PreambleAction::Take,
1744 (true, false) => panic!("Need preamble but there isn't one"),
1745 (false, true) => PreambleAction::Skip,
1746 (false, false) => PreambleAction::Absent,
1747 };
1748
1749 let rows_taking = if *rows_desired >= rows_avail {
1752 *need_preamble = self.take_trailer;
1760 rows_avail
1761 } else {
1762 *need_preamble = false;
1765 *rows_desired
1766 };
1767 let rows_skipped = *skip_in_chunk;
1768
1769 let consumed_chunk = if *rows_desired >= rows_avail {
1771 *rows_desired -= rows_avail;
1772 *skip_in_chunk = 0;
1773 true
1774 } else {
1775 *skip_in_chunk += *rows_desired;
1776 *rows_desired = 0;
1777 false
1778 };
1779
1780 (
1781 ChunkDrainInstructions {
1782 chunk_instructions: self.clone(),
1783 rows_to_skip: rows_skipped,
1784 rows_to_take: rows_taking,
1785 preamble_action,
1786 },
1787 consumed_chunk,
1788 )
1789 }
1790}
1791
1792enum Words {
1793 U16(ScalarBuffer<u16>),
1794 U32(ScalarBuffer<u32>),
1795}
1796
1797struct WordsIter<'a> {
1798 iter: Box<dyn Iterator<Item = u32> + 'a>,
1799}
1800
1801impl Words {
1802 pub fn len(&self) -> usize {
1803 match self {
1804 Self::U16(b) => b.len(),
1805 Self::U32(b) => b.len(),
1806 }
1807 }
1808
1809 pub fn iter(&self) -> WordsIter<'_> {
1810 match self {
1811 Self::U16(buf) => WordsIter {
1812 iter: Box::new(buf.iter().map(|&x| x as u32)),
1813 },
1814 Self::U32(buf) => WordsIter {
1815 iter: Box::new(buf.iter().copied()),
1816 },
1817 }
1818 }
1819
1820 pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
1821 let bytes_per_value = if has_large_chunk { 4 } else { 2 };
1822 assert_eq!(bytes.len() % bytes_per_value, 0);
1823 let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
1824 if has_large_chunk {
1825 Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
1826 } else {
1827 Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
1828 }
1829 }
1830}
1831
1832impl<'a> Iterator for WordsIter<'a> {
1833 type Item = u32;
1834
1835 fn next(&mut self) -> Option<Self::Item> {
1836 self.iter.next()
1837 }
1838}
1839
1840impl StructuralPageScheduler for MiniBlockScheduler {
1841 fn initialize<'a>(
1842 &'a mut self,
1843 io: &Arc<dyn EncodingsIo>,
1844 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1845 let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
1849 let value_buf_position = self.buffer_offsets_and_sizes[1].0;
1850 let mut bufs_needed = 1;
1851 if self.dictionary.is_some() {
1852 bufs_needed += 1;
1853 }
1854 if self.repetition_index_depth > 0 {
1855 bufs_needed += 1;
1856 }
1857 let mut required_ranges = Vec::with_capacity(bufs_needed);
1858 required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size);
1859 if let Some(ref dictionary) = self.dictionary {
1860 required_ranges.push(
1861 dictionary.dictionary_buf_position_and_size.0
1862 ..dictionary.dictionary_buf_position_and_size.0
1863 + dictionary.dictionary_buf_position_and_size.1,
1864 );
1865 }
1866 if self.repetition_index_depth > 0 {
1867 let (rep_index_pos, rep_index_size) = self.buffer_offsets_and_sizes.last().unwrap();
1868 required_ranges.push(*rep_index_pos..*rep_index_pos + *rep_index_size);
1869 }
1870 let io_req = io.submit_request(required_ranges, 0);
1871
1872 async move {
1873 let mut buffers = io_req.await?.into_iter().fuse();
1874 let meta_bytes = buffers.next().unwrap();
1875 let dictionary_bytes = self.dictionary.as_ref().and_then(|_| buffers.next());
1876 let rep_index_bytes = buffers.next();
1877
1878 let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
1880 let mut chunk_meta = Vec::with_capacity(words.len());
1881
1882 let mut rows_counter = 0;
1883 let mut offset_bytes = value_buf_position;
1884 for (word_idx, word) in words.iter().enumerate() {
1885 let log_num_values = word & 0x0F;
1886 let divided_bytes = word >> 4;
1887 let num_bytes = (divided_bytes as usize + 1) * MINIBLOCK_ALIGNMENT;
1888 debug_assert!(num_bytes > 0);
1889 let num_values = if word_idx < words.len() - 1 {
1890 debug_assert!(log_num_values > 0);
1891 1 << log_num_values
1892 } else {
1893 debug_assert!(
1894 log_num_values == 0
1895 || (1 << log_num_values) == (self.items_in_page - rows_counter)
1896 );
1897 self.items_in_page - rows_counter
1898 };
1899 rows_counter += num_values;
1900
1901 chunk_meta.push(ChunkMeta {
1902 num_values,
1903 chunk_size_bytes: num_bytes as u64,
1904 offset_bytes,
1905 });
1906 offset_bytes += num_bytes as u64;
1907 }
1908
1909 let rep_index = if let Some(rep_index_data) = rep_index_bytes {
1911 assert!(rep_index_data.len() % 8 == 0);
1912 let stride = self.repetition_index_depth as usize + 1;
1913 MiniBlockRepIndex::decode_from_bytes(&rep_index_data, stride)
1914 } else {
1915 MiniBlockRepIndex::default_from_chunks(&chunk_meta)
1916 };
1917
1918 let mut page_meta = MiniBlockCacheableState {
1919 chunk_meta,
1920 rep_index,
1921 dictionary: None,
1922 };
1923
1924 if let Some(ref mut dictionary) = self.dictionary {
1926 let dictionary_data = dictionary_bytes.unwrap();
1927 page_meta.dictionary =
1928 Some(Arc::new(dictionary.dictionary_decompressor.decompress(
1929 LanceBuffer::from_bytes(
1930 dictionary_data,
1931 dictionary.dictionary_data_alignment,
1932 ),
1933 dictionary.num_dictionary_items,
1934 )?));
1935 };
1936 let page_meta = Arc::new(page_meta);
1937 self.page_meta = Some(page_meta.clone());
1938 Ok(page_meta as Arc<dyn CachedPageData>)
1939 }
1940 .boxed()
1941 }
1942
1943 fn load(&mut self, data: &Arc<dyn CachedPageData>) {
1944 self.page_meta = Some(
1945 data.clone()
1946 .as_arc_any()
1947 .downcast::<MiniBlockCacheableState>()
1948 .unwrap(),
1949 );
1950 }
1951
1952 fn schedule_ranges(
1953 &self,
1954 ranges: &[Range<u64>],
1955 io: &Arc<dyn EncodingsIo>,
1956 ) -> Result<Vec<PageLoadTask>> {
1957 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
1958
1959 let page_meta = self.page_meta.as_ref().unwrap();
1960
1961 let chunk_instructions =
1962 ChunkInstructions::schedule_instructions(&page_meta.rep_index, ranges);
1963
1964 debug_assert_eq!(
1965 num_rows,
1966 chunk_instructions
1967 .iter()
1968 .map(|ci| ci.rows_to_take)
1969 .sum::<u64>()
1970 );
1971
1972 let chunks_needed = chunk_instructions
1973 .iter()
1974 .map(|ci| ci.chunk_idx)
1975 .unique()
1976 .collect::<Vec<_>>();
1977
1978 let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
1979 let chunk_ranges = loaded_chunks
1980 .iter()
1981 .map(|c| c.byte_range.clone())
1982 .collect::<Vec<_>>();
1983 let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
1984
1985 let rep_decompressor = self.rep_decompressor.clone();
1986 let def_decompressor = self.def_decompressor.clone();
1987 let value_decompressor = self.value_decompressor.clone();
1988 let num_buffers = self.num_buffers;
1989 let has_large_chunk = self.has_large_chunk;
1990 let dictionary = page_meta
1991 .dictionary
1992 .as_ref()
1993 .map(|dictionary| dictionary.clone());
1994 let def_meaning = self.def_meaning.clone();
1995
1996 let res = async move {
1997 let loaded_chunk_data = loaded_chunk_data.await?;
1998 for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
1999 loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
2000 }
2001
2002 Ok(Box::new(MiniBlockDecoder {
2003 rep_decompressor,
2004 def_decompressor,
2005 value_decompressor,
2006 def_meaning,
2007 loaded_chunks: VecDeque::from_iter(loaded_chunks),
2008 instructions: VecDeque::from(chunk_instructions),
2009 offset_in_current_chunk: 0,
2010 dictionary,
2011 num_rows,
2012 num_buffers,
2013 has_large_chunk,
2014 }) as Box<dyn StructuralPageDecoder>)
2015 }
2016 .boxed();
2017 let page_load_task = PageLoadTask {
2018 decoder_fut: res,
2019 num_rows,
2020 };
2021 Ok(vec![page_load_task])
2022 }
2023}
2024
2025#[derive(Debug, Clone, Copy)]
2026struct FullZipRepIndexDetails {
2027 buf_position: u64,
2028 bytes_per_value: u64, }
2030
2031#[derive(Debug)]
2032enum PerValueDecompressor {
2033 Fixed(Arc<dyn FixedPerValueDecompressor>),
2034 Variable(Arc<dyn VariablePerValueDecompressor>),
2035}
2036
2037#[derive(Debug)]
2038struct FullZipDecodeDetails {
2039 value_decompressor: PerValueDecompressor,
2040 def_meaning: Arc<[DefinitionInterpretation]>,
2041 ctrl_word_parser: ControlWordParser,
2042 max_rep: u16,
2043 max_visible_def: u16,
2044}
2045
2046#[derive(Debug, Clone)]
2058enum FullZipReadSource {
2059 Remote(Arc<dyn EncodingsIo>),
2061 PrefetchedPage { base_offset: u64, data: LanceBuffer },
2063}
2064
2065impl FullZipReadSource {
2066 fn fetch(
2070 &self,
2071 ranges: &[Range<u64>],
2072 priority: u64,
2073 ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
2074 match self {
2075 Self::Remote(io) => {
2076 let io = io.clone();
2077 let ranges = ranges.to_vec();
2078 async move {
2079 let data = io.submit_request(ranges, priority).await?;
2080 Ok(data
2081 .into_iter()
2082 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2083 .collect::<VecDeque<_>>())
2084 }
2085 .boxed()
2086 }
2087 Self::PrefetchedPage { base_offset, data } => {
2088 let base_offset = *base_offset;
2089 let data = data.clone();
2090 let page_end = base_offset + data.len() as u64;
2091 std::future::ready(
2092 ranges
2093 .iter()
2094 .map(|range| {
2095 if range.start > range.end
2096 || range.start < base_offset
2097 || range.end > page_end
2098 {
2099 return Err(Error::internal(format!(
2100 "Requested range {:?} is outside page range {}..{}",
2101 range, base_offset, page_end
2102 )));
2103 }
2104 let start = (range.start - base_offset) as usize;
2105 let len = (range.end - range.start) as usize;
2106 Ok(data.slice_with_length(start, len))
2107 })
2108 .collect::<Result<VecDeque<_>>>(),
2109 )
2110 .boxed()
2111 }
2112 }
2113 }
2114}
2115
2116#[derive(Debug)]
2124pub struct FullZipScheduler {
2125 data_buf_position: u64,
2126 data_buf_size: u64,
2127 rep_index: Option<FullZipRepIndexDetails>,
2128 priority: u64,
2129 rows_in_page: u64,
2130 bits_per_offset: u8,
2131 details: Arc<FullZipDecodeDetails>,
2132 cached_state: Option<Arc<FullZipCacheableState>>,
2134 enable_cache: bool,
2136}
2137
2138impl FullZipScheduler {
2139 fn try_new(
2140 buffer_offsets_and_sizes: &[(u64, u64)],
2141 priority: u64,
2142 rows_in_page: u64,
2143 layout: &pb21::FullZipLayout,
2144 decompressors: &dyn DecompressionStrategy,
2145 ) -> Result<Self> {
2146 let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
2147 let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
2148 let num_reps = rows_in_page + 1;
2149 let bytes_per_rep = len / num_reps;
2150 debug_assert_eq!(len % num_reps, 0);
2151 debug_assert!(
2152 bytes_per_rep == 1
2153 || bytes_per_rep == 2
2154 || bytes_per_rep == 4
2155 || bytes_per_rep == 8
2156 );
2157 FullZipRepIndexDetails {
2158 buf_position: *pos,
2159 bytes_per_value: bytes_per_rep,
2160 }
2161 });
2162
2163 let value_decompressor = match layout.details {
2164 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
2165 let decompressor = decompressors.create_fixed_per_value_decompressor(
2166 layout.value_compression.as_ref().unwrap(),
2167 )?;
2168 PerValueDecompressor::Fixed(decompressor.into())
2169 }
2170 Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
2171 let decompressor = decompressors.create_variable_per_value_decompressor(
2172 layout.value_compression.as_ref().unwrap(),
2173 )?;
2174 PerValueDecompressor::Variable(decompressor.into())
2175 }
2176 None => {
2177 panic!("Full-zip layout must have a `details` field");
2178 }
2179 };
2180 let ctrl_word_parser = ControlWordParser::new(
2181 layout.bits_rep.try_into().unwrap(),
2182 layout.bits_def.try_into().unwrap(),
2183 );
2184 let def_meaning = layout
2185 .layers
2186 .iter()
2187 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2188 .collect::<Vec<_>>();
2189
2190 let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
2191 let max_visible_def = def_meaning
2192 .iter()
2193 .filter(|d| !d.is_list())
2194 .map(|d| d.num_def_levels())
2195 .sum();
2196
2197 let bits_per_offset = match layout.details {
2198 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
2199 Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
2200 bits_per_offset as u8
2201 }
2202 None => panic!("Full-zip layout must have a `details` field"),
2203 };
2204
2205 let details = Arc::new(FullZipDecodeDetails {
2206 value_decompressor,
2207 def_meaning: def_meaning.into(),
2208 ctrl_word_parser,
2209 max_rep,
2210 max_visible_def,
2211 });
2212 Ok(Self {
2213 data_buf_position,
2214 data_buf_size,
2215 rep_index,
2216 details,
2217 priority,
2218 rows_in_page,
2219 bits_per_offset,
2220 cached_state: None,
2221 enable_cache: false,
2222 })
2223 }
2224
2225 fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
2226 if ranges.is_empty() {
2227 return false;
2228 }
2229 let mut expected_start = 0;
2230 for range in ranges {
2231 if range.start != expected_start || range.end > rows_in_page || range.end < range.start
2232 {
2233 return false;
2234 }
2235 expected_start = range.end;
2236 }
2237 expected_start == rows_in_page
2238 }
2239
2240 fn create_page_load_task(
2241 io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
2242 num_rows: u64,
2243 details: Arc<FullZipDecodeDetails>,
2244 bits_per_offset: u8,
2245 ) -> PageLoadTask {
2246 let load_task = async move {
2247 let buffers = io_future.await?;
2248 let data = buffers
2249 .into_iter()
2250 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2251 .collect::<VecDeque<_>>();
2252 Self::create_decoder(details, data, num_rows, bits_per_offset)
2253 }
2254 .boxed();
2255 PageLoadTask {
2256 decoder_fut: load_task,
2257 num_rows,
2258 }
2259 }
2260
2261 fn create_decoder(
2263 details: Arc<FullZipDecodeDetails>,
2264 data: VecDeque<LanceBuffer>,
2265 num_rows: u64,
2266 bits_per_offset: u8,
2267 ) -> Result<Box<dyn StructuralPageDecoder>> {
2268 match &details.value_decompressor {
2269 PerValueDecompressor::Fixed(decompressor) => {
2270 let bits_per_value = decompressor.bits_per_value();
2271 if bits_per_value % 8 != 0 {
2272 return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
2273 }
2274 let bytes_per_value = bits_per_value / 8;
2275 let total_bytes_per_value =
2276 bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
2277 if total_bytes_per_value == 0 {
2278 return Err(lance_core::Error::internal(
2279 "Invalid encoding: per-row byte width must be greater than 0",
2280 ));
2281 }
2282 Ok(Box::new(FixedFullZipDecoder {
2283 details,
2284 data,
2285 num_rows,
2286 offset_in_current: 0,
2287 bytes_per_value: bytes_per_value as usize,
2288 total_bytes_per_value,
2289 }) as Box<dyn StructuralPageDecoder>)
2290 }
2291 PerValueDecompressor::Variable(_decompressor) => {
2292 Ok(Box::new(VariableFullZipDecoder::new(
2293 details,
2294 data,
2295 num_rows,
2296 bits_per_offset,
2297 bits_per_offset,
2298 )?))
2299 }
2300 }
2301 }
2302
2303 fn extract_byte_ranges_from_pairs(
2306 buffer: LanceBuffer,
2307 bytes_per_value: u64,
2308 data_buf_position: u64,
2309 ) -> Vec<Range<u64>> {
2310 ByteUnpacker::new(buffer, bytes_per_value as usize)
2311 .chunks(2)
2312 .into_iter()
2313 .map(|mut c| {
2314 let start = c.next().unwrap() + data_buf_position;
2315 let end = c.next().unwrap() + data_buf_position;
2316 start..end
2317 })
2318 .collect::<Vec<_>>()
2319 }
2320
2321 fn extract_byte_ranges_from_cached(
2324 buffer: &LanceBuffer,
2325 ranges: &[Range<u64>],
2326 bytes_per_value: u64,
2327 data_buf_position: u64,
2328 ) -> Vec<Range<u64>> {
2329 ranges
2330 .iter()
2331 .map(|r| {
2332 let start_offset = (r.start * bytes_per_value) as usize;
2333 let end_offset = (r.end * bytes_per_value) as usize;
2334
2335 let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
2336 let start_val =
2337 ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
2338 .next()
2339 .unwrap();
2340
2341 let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
2342 let end_val =
2343 ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
2344 .next()
2345 .unwrap();
2346
2347 (data_buf_position + start_val)..(data_buf_position + end_val)
2348 })
2349 .collect()
2350 }
2351
2352 fn compute_rep_index_ranges(
2354 ranges: &[Range<u64>],
2355 rep_index: &FullZipRepIndexDetails,
2356 ) -> Vec<Range<u64>> {
2357 ranges
2358 .iter()
2359 .flat_map(|r| {
2360 let first_val_start =
2361 rep_index.buf_position + (r.start * rep_index.bytes_per_value);
2362 let first_val_end = first_val_start + rep_index.bytes_per_value;
2363 let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
2364 let last_val_end = last_val_start + rep_index.bytes_per_value;
2365 [first_val_start..first_val_end, last_val_start..last_val_end]
2366 })
2367 .collect()
2368 }
2369
2370 fn schedule_ranges_rep(
2372 &self,
2373 ranges: &[Range<u64>],
2374 io: &Arc<dyn EncodingsIo>,
2375 rep_index: FullZipRepIndexDetails,
2376 ) -> Result<Vec<PageLoadTask>> {
2377 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2378 let data_buf_position = self.data_buf_position;
2379 let priority = self.priority;
2380 let details = self.details.clone();
2381 let bits_per_offset = self.bits_per_offset;
2382
2383 if Self::covers_entire_page(ranges, self.rows_in_page) {
2384 let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
2385 let page_data = io.submit_single(full_range.clone(), priority);
2386 let load_task = async move {
2387 let page_data = page_data.await?;
2388 let source = FullZipReadSource::PrefetchedPage {
2389 base_offset: full_range.start,
2390 data: LanceBuffer::from_bytes(page_data, 1),
2391 };
2392 let read_ranges = vec![full_range];
2393 let data = source.fetch(&read_ranges, priority).await?;
2394 Self::create_decoder(details, data, num_rows, bits_per_offset)
2395 }
2396 .boxed();
2397 let page_load_task = PageLoadTask {
2398 decoder_fut: load_task,
2399 num_rows,
2400 };
2401 return Ok(vec![page_load_task]);
2402 }
2403
2404 if let Some(cached_state) = &self.cached_state {
2405 let byte_ranges = Self::extract_byte_ranges_from_cached(
2406 &cached_state.rep_index_buffer,
2407 ranges,
2408 rep_index.bytes_per_value,
2409 data_buf_position,
2410 );
2411 let io_future = io.submit_request(byte_ranges, priority);
2412 let page_load_task =
2413 Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
2414 return Ok(vec![page_load_task]);
2415 }
2416
2417 let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
2418 let rep_data = io.submit_request(rep_ranges, priority);
2419 let io_clone = io.clone();
2420 let load_task = async move {
2421 let rep_data = rep_data.await?;
2422 let rep_buffer = LanceBuffer::concat(
2423 &rep_data
2424 .into_iter()
2425 .map(|d| LanceBuffer::from_bytes(d, 1))
2426 .collect::<Vec<_>>(),
2427 );
2428 let byte_ranges = Self::extract_byte_ranges_from_pairs(
2429 rep_buffer,
2430 rep_index.bytes_per_value,
2431 data_buf_position,
2432 );
2433 let source = FullZipReadSource::Remote(io_clone);
2434 let data = source.fetch(&byte_ranges, priority).await?;
2435 Self::create_decoder(details, data, num_rows, bits_per_offset)
2436 }
2437 .boxed();
2438 let page_load_task = PageLoadTask {
2439 decoder_fut: load_task,
2440 num_rows,
2441 };
2442 Ok(vec![page_load_task])
2443 }
2444
2445 fn schedule_ranges_simple(
2449 &self,
2450 ranges: &[Range<u64>],
2451 io: &Arc<dyn EncodingsIo>,
2452 ) -> Result<Vec<PageLoadTask>> {
2453 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2455
2456 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
2457 unreachable!()
2458 };
2459
2460 let bits_per_value = decompressor.bits_per_value();
2462 assert_eq!(bits_per_value % 8, 0);
2463 let bytes_per_value = bits_per_value / 8;
2464 let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
2465 let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
2466 let byte_ranges = ranges
2467 .iter()
2468 .map(|r| {
2469 debug_assert!(r.end <= self.rows_in_page);
2470 let start = self.data_buf_position + r.start * total_bytes_per_value;
2471 let end = self.data_buf_position + r.end * total_bytes_per_value;
2472 start..end
2473 })
2474 .collect::<Vec<_>>();
2475
2476 let io_future = io.submit_request(byte_ranges, self.priority);
2477 let page_load_task = Self::create_page_load_task(
2478 io_future,
2479 num_rows,
2480 self.details.clone(),
2481 self.bits_per_offset,
2482 );
2483 Ok(vec![page_load_task])
2484 }
2485}
2486
2487#[derive(Debug)]
2489struct FullZipCacheableState {
2490 rep_index_buffer: LanceBuffer,
2492}
2493
2494impl DeepSizeOf for FullZipCacheableState {
2495 fn deep_size_of_children(&self, _context: &mut Context) -> usize {
2496 self.rep_index_buffer.len()
2497 }
2498}
2499
2500impl CachedPageData for FullZipCacheableState {
2501 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
2502 self
2503 }
2504}
2505
2506impl StructuralPageScheduler for FullZipScheduler {
2507 fn initialize<'a>(
2508 &'a mut self,
2509 io: &Arc<dyn EncodingsIo>,
2510 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2511 if self.enable_cache
2512 && let Some(rep_index) = self.rep_index
2513 {
2514 let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
2515 let rep_index_range = rep_index.buf_position..(rep_index.buf_position + total_size);
2516 let io_clone = io.clone();
2517 return async move {
2518 let rep_index_data = io_clone.submit_request(vec![rep_index_range], 0).await?;
2519 let state = Arc::new(FullZipCacheableState {
2520 rep_index_buffer: LanceBuffer::from_bytes(rep_index_data[0].clone(), 1),
2521 });
2522 self.cached_state = Some(state.clone());
2523 Ok(state as Arc<dyn CachedPageData>)
2524 }
2525 .boxed();
2526 }
2527 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
2528 }
2529
2530 fn load(&mut self, cache: &Arc<dyn CachedPageData>) {
2534 if let Ok(cached_state) = cache
2536 .clone()
2537 .as_arc_any()
2538 .downcast::<FullZipCacheableState>()
2539 {
2540 self.cached_state = Some(cached_state);
2542 }
2543 }
2544
2545 fn schedule_ranges(
2546 &self,
2547 ranges: &[Range<u64>],
2548 io: &Arc<dyn EncodingsIo>,
2549 ) -> Result<Vec<PageLoadTask>> {
2550 if let Some(rep_index) = self.rep_index {
2551 self.schedule_ranges_rep(ranges, io, rep_index)
2552 } else {
2553 self.schedule_ranges_simple(ranges, io)
2554 }
2555 }
2556}
2557
2558#[derive(Debug)]
2566struct FixedFullZipDecoder {
2567 details: Arc<FullZipDecodeDetails>,
2568 data: VecDeque<LanceBuffer>,
2569 offset_in_current: usize,
2570 bytes_per_value: usize,
2571 total_bytes_per_value: usize,
2572 num_rows: u64,
2573}
2574
2575impl FixedFullZipDecoder {
2576 fn slice_next_task(&mut self, num_rows: u64) -> FullZipDecodeTaskItem {
2577 debug_assert!(num_rows > 0);
2578 let cur_buf = self.data.front_mut().unwrap();
2579 let start = self.offset_in_current;
2580 if self.details.ctrl_word_parser.has_rep() {
2581 let mut rows_started = 0;
2584 let mut num_items = 0;
2587 while self.offset_in_current < cur_buf.len() {
2588 let control = self.details.ctrl_word_parser.parse_desc(
2589 &cur_buf[self.offset_in_current..],
2590 self.details.max_rep,
2591 self.details.max_visible_def,
2592 );
2593 if control.is_new_row {
2594 if rows_started == num_rows {
2595 break;
2596 }
2597 rows_started += 1;
2598 }
2599 num_items += 1;
2600 if control.is_visible {
2601 self.offset_in_current += self.total_bytes_per_value;
2602 } else {
2603 self.offset_in_current += self.details.ctrl_word_parser.bytes_per_word();
2604 }
2605 }
2606
2607 let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
2608 if self.offset_in_current == cur_buf.len() {
2609 self.data.pop_front();
2610 self.offset_in_current = 0;
2611 }
2612
2613 FullZipDecodeTaskItem {
2614 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2615 data: task_slice,
2616 bits_per_value: self.bytes_per_value as u64 * 8,
2617 num_values: num_items,
2618 block_info: BlockInfo::new(),
2619 }),
2620 rows_in_buf: rows_started,
2621 }
2622 } else {
2623 let cur_buf = self.data.front_mut().unwrap();
2626 let bytes_avail = cur_buf.len() - self.offset_in_current;
2627 let offset_in_cur = self.offset_in_current;
2628
2629 let bytes_needed = num_rows as usize * self.total_bytes_per_value;
2630 let mut rows_taken = num_rows;
2631 let task_slice = if bytes_needed >= bytes_avail {
2632 self.offset_in_current = 0;
2633 rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
2634 self.data
2635 .pop_front()
2636 .unwrap()
2637 .slice_with_length(offset_in_cur, bytes_avail)
2638 } else {
2639 self.offset_in_current += bytes_needed;
2640 cur_buf.slice_with_length(offset_in_cur, bytes_needed)
2641 };
2642 FullZipDecodeTaskItem {
2643 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2644 data: task_slice,
2645 bits_per_value: self.bytes_per_value as u64 * 8,
2646 num_values: rows_taken,
2647 block_info: BlockInfo::new(),
2648 }),
2649 rows_in_buf: rows_taken,
2650 }
2651 }
2652 }
2653}
2654
2655impl StructuralPageDecoder for FixedFullZipDecoder {
2656 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2657 let mut task_data = Vec::with_capacity(self.data.len());
2658 let mut remaining = num_rows;
2659 while remaining > 0 {
2660 let task_item = self.slice_next_task(remaining);
2661 remaining -= task_item.rows_in_buf;
2662 task_data.push(task_item);
2663 }
2664 Ok(Box::new(FixedFullZipDecodeTask {
2665 details: self.details.clone(),
2666 data: task_data,
2667 bytes_per_value: self.bytes_per_value,
2668 num_rows: num_rows as usize,
2669 }))
2670 }
2671
2672 fn num_rows(&self) -> u64 {
2673 self.num_rows
2674 }
2675}
2676
2677#[derive(Debug)]
2682struct VariableFullZipDecoder {
2683 details: Arc<FullZipDecodeDetails>,
2684 decompressor: Arc<dyn VariablePerValueDecompressor>,
2685 data: LanceBuffer,
2686 offsets: LanceBuffer,
2687 rep: ScalarBuffer<u16>,
2688 def: ScalarBuffer<u16>,
2689 repdef_starts: Vec<usize>,
2690 data_starts: Vec<usize>,
2691 offset_starts: Vec<usize>,
2692 visible_item_counts: Vec<u64>,
2693 bits_per_offset: u8,
2694 current_idx: usize,
2695 num_rows: u64,
2696}
2697
2698fn corrupt_file_named(name: &str, message: impl Into<String>) -> Error {
2699 Error::corrupt_file(name.into(), message)
2700}
2701
2702impl VariableFullZipDecoder {
2703 fn new(
2704 details: Arc<FullZipDecodeDetails>,
2705 data: VecDeque<LanceBuffer>,
2706 num_rows: u64,
2707 in_bits_per_length: u8,
2708 out_bits_per_offset: u8,
2709 ) -> Result<Self> {
2710 let decompressor = match details.value_decompressor {
2711 PerValueDecompressor::Variable(ref d) => d.clone(),
2712 _ => unreachable!(),
2713 };
2714
2715 assert_eq!(in_bits_per_length % 8, 0);
2716 assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
2717
2718 let mut decoder = Self {
2719 details,
2720 decompressor,
2721 data: LanceBuffer::empty(),
2722 offsets: LanceBuffer::empty(),
2723 rep: LanceBuffer::empty().borrow_to_typed_slice(),
2724 def: LanceBuffer::empty().borrow_to_typed_slice(),
2725 bits_per_offset: out_bits_per_offset,
2726 repdef_starts: Vec::with_capacity(num_rows as usize + 1),
2727 data_starts: Vec::with_capacity(num_rows as usize + 1),
2728 offset_starts: Vec::with_capacity(num_rows as usize + 1),
2729 visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
2730 current_idx: 0,
2731 num_rows,
2732 };
2733
2734 decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
2755
2756 Ok(decoder)
2757 }
2758
2759 fn slice_batch_data_and_rebase_offsets_typed<T>(
2760 data: &LanceBuffer,
2761 offsets: &LanceBuffer,
2762 ) -> Result<(LanceBuffer, LanceBuffer)>
2763 where
2764 T: arrow_buffer::ArrowNativeType
2765 + Copy
2766 + PartialOrd
2767 + std::ops::Sub<Output = T>
2768 + std::fmt::Display
2769 + TryInto<usize>,
2770 {
2771 let offsets_slice = offsets.borrow_to_typed_slice::<T>();
2772 let offsets_slice = offsets_slice.as_ref();
2773 if offsets_slice.is_empty() {
2774 return Err(Error::internal(
2775 "Variable offsets cannot be empty".to_string(),
2776 ));
2777 }
2778
2779 let base = offsets_slice[0];
2780 let end = *offsets_slice.last().unwrap();
2781 if end < base {
2782 return Err(Error::internal(format!(
2783 "Invalid variable offsets: end ({end}) is less than base ({base})"
2784 )));
2785 }
2786
2787 let data_start = base.try_into().map_err(|_| {
2788 Error::internal(format!("Variable offset ({base}) does not fit into usize"))
2789 })?;
2790 let data_end = end.try_into().map_err(|_| {
2791 Error::internal(format!("Variable offset ({end}) does not fit into usize"))
2792 })?;
2793 if data_end > data.len() {
2794 return Err(Error::internal(format!(
2795 "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
2796 data.len()
2797 )));
2798 }
2799
2800 let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
2801 for &offset in offsets_slice {
2802 if offset < base {
2803 return Err(Error::internal(format!(
2804 "Invalid variable offsets: offset ({offset}) is less than base ({base})"
2805 )));
2806 }
2807 rebased_offsets.push(offset - base);
2808 }
2809
2810 let sliced_data = data.slice_with_length(data_start, data_end - data_start);
2811 let sliced_data = LanceBuffer::copy_slice(&sliced_data);
2813 let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
2814 Ok((sliced_data, rebased_offsets))
2815 }
2816
2817 fn slice_batch_data_and_rebase_offsets(
2818 data: &LanceBuffer,
2819 offsets: &LanceBuffer,
2820 bits_per_offset: u8,
2821 ) -> Result<(LanceBuffer, LanceBuffer)> {
2822 match bits_per_offset {
2823 32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
2824 64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
2825 _ => Err(Error::internal(format!(
2826 "Unsupported bits_per_offset={bits_per_offset}"
2827 ))),
2828 }
2829 }
2830
2831 fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
2838 let width = bits_per_offset as usize / 8;
2839 if data.len() < width {
2840 return Err(corrupt_file_named(
2841 "variable_full_zip",
2842 format!(
2843 "truncated length prefix: {} byte(s) remain in the page buffer but a \
2844 {}-bit length prefix requires {}",
2845 data.len(),
2846 bits_per_offset,
2847 width
2848 ),
2849 ));
2850 }
2851 Ok(match bits_per_offset {
2852 8 => data[0] as u64,
2853 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
2854 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
2855 64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
2856 _ => unreachable!(),
2857 })
2858 }
2859
2860 fn unzip(
2861 &mut self,
2862 data: VecDeque<LanceBuffer>,
2863 in_bits_per_length: u8,
2864 out_bits_per_offset: u8,
2865 num_rows: u64,
2866 ) -> Result<()> {
2867 let mut rep = Vec::with_capacity(num_rows as usize);
2869 let mut def = Vec::with_capacity(num_rows as usize);
2870 let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
2871
2872 let bytes_per_offset = out_bits_per_offset as usize / 8;
2875 let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
2876 let mut offsets_data = Vec::with_capacity(bytes_offsets);
2877
2878 let bytes_per_length = in_bits_per_length as usize / 8;
2879 let bytes_lengths = bytes_per_length * num_rows as usize;
2880
2881 let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
2882 let mut unzipped_data =
2885 Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
2886
2887 let mut current_offset = 0_u64;
2888 let mut visible_item_count = 0_u64;
2889 for databuf in data.into_iter() {
2890 let mut databuf = databuf.as_ref();
2891 while !databuf.is_empty() {
2892 let data_start = unzipped_data.len();
2893 let offset_start = offsets_data.len();
2894 let repdef_start = rep.len().max(def.len());
2897 let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
2899 databuf,
2900 self.details.max_rep,
2901 self.details.max_visible_def,
2902 );
2903 self.details
2904 .ctrl_word_parser
2905 .parse(databuf, &mut rep, &mut def);
2906 databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
2907
2908 if ctrl_desc.is_new_row {
2909 self.repdef_starts.push(repdef_start);
2910 self.data_starts.push(data_start);
2911 self.offset_starts.push(offset_start);
2912 self.visible_item_counts.push(visible_item_count);
2913 }
2914 if ctrl_desc.is_visible {
2915 visible_item_count += 1;
2916 if ctrl_desc.is_valid_item {
2917 let length = Self::parse_length(databuf, in_bits_per_length)?;
2918 match out_bits_per_offset {
2919 32 => offsets_data
2920 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2921 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2922 _ => unreachable!(),
2923 };
2924 databuf = &databuf[bytes_per_offset..];
2925 unzipped_data.extend_from_slice(&databuf[..length as usize]);
2926 databuf = &databuf[length as usize..];
2927 current_offset += length;
2928 } else {
2929 match out_bits_per_offset {
2931 32 => offsets_data
2932 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2933 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2934 _ => unreachable!(),
2935 }
2936 }
2937 }
2938 }
2939 }
2940 self.repdef_starts.push(rep.len().max(def.len()));
2941 self.data_starts.push(unzipped_data.len());
2942 self.offset_starts.push(offsets_data.len());
2943 self.visible_item_counts.push(visible_item_count);
2944 match out_bits_per_offset {
2945 32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
2946 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2947 _ => unreachable!(),
2948 };
2949 self.rep = ScalarBuffer::from(rep);
2950 self.def = ScalarBuffer::from(def);
2951 self.data = LanceBuffer::from(unzipped_data);
2952 self.offsets = LanceBuffer::from(offsets_data);
2953 Ok(())
2954 }
2955}
2956
2957impl StructuralPageDecoder for VariableFullZipDecoder {
2958 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2959 let start = self.current_idx;
2960 let end = start + num_rows as usize;
2961
2962 let offset_start = self.offset_starts[start];
2963 let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
2964 let offsets = self
2965 .offsets
2966 .slice_with_length(offset_start, offset_end - offset_start);
2967 let (data, offsets) =
2969 Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
2970
2971 let repdef_start = self.repdef_starts[start];
2972 let repdef_end = self.repdef_starts[end];
2973 let rep = if self.rep.is_empty() {
2974 self.rep.clone()
2975 } else {
2976 self.rep.slice(repdef_start, repdef_end - repdef_start)
2977 };
2978 let def = if self.def.is_empty() {
2979 self.def.clone()
2980 } else {
2981 self.def.slice(repdef_start, repdef_end - repdef_start)
2982 };
2983
2984 let visible_item_counts_start = self.visible_item_counts[start];
2985 let visible_item_counts_end = self.visible_item_counts[end];
2986 let num_visible_items = visible_item_counts_end - visible_item_counts_start;
2987
2988 self.current_idx += num_rows as usize;
2989
2990 Ok(Box::new(VariableFullZipDecodeTask {
2991 details: self.details.clone(),
2992 decompressor: self.decompressor.clone(),
2993 data,
2994 offsets,
2995 bits_per_offset: self.bits_per_offset,
2996 num_visible_items,
2997 rep,
2998 def,
2999 }))
3000 }
3001
3002 fn num_rows(&self) -> u64 {
3003 self.num_rows
3004 }
3005}
3006
3007#[derive(Debug)]
3008struct VariableFullZipDecodeTask {
3009 details: Arc<FullZipDecodeDetails>,
3010 decompressor: Arc<dyn VariablePerValueDecompressor>,
3011 data: LanceBuffer,
3012 offsets: LanceBuffer,
3013 bits_per_offset: u8,
3014 num_visible_items: u64,
3015 rep: ScalarBuffer<u16>,
3016 def: ScalarBuffer<u16>,
3017}
3018
3019impl DecodePageTask for VariableFullZipDecodeTask {
3020 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3021 let block = VariableWidthBlock {
3022 data: self.data,
3023 offsets: self.offsets,
3024 bits_per_offset: self.bits_per_offset,
3025 num_values: self.num_visible_items,
3026 block_info: BlockInfo::new(),
3027 };
3028 let decomopressed = self.decompressor.decompress(block)?;
3029 let rep = if self.rep.is_empty() {
3030 None
3031 } else {
3032 Some(self.rep.to_vec())
3033 };
3034 let def = if self.def.is_empty() {
3035 None
3036 } else {
3037 Some(self.def.to_vec())
3038 };
3039 let unraveler = RepDefUnraveler::new(
3040 rep,
3041 def,
3042 self.details.def_meaning.clone(),
3043 self.num_visible_items,
3044 );
3045 Ok(DecodedPage {
3046 data: decomopressed,
3047 repdef: unraveler,
3048 })
3049 }
3050}
3051
3052#[derive(Debug)]
3053struct FullZipDecodeTaskItem {
3054 data: PerValueDataBlock,
3055 rows_in_buf: u64,
3056}
3057
3058#[derive(Debug)]
3061struct FixedFullZipDecodeTask {
3062 details: Arc<FullZipDecodeDetails>,
3063 data: Vec<FullZipDecodeTaskItem>,
3064 num_rows: usize,
3065 bytes_per_value: usize,
3066}
3067
3068impl DecodePageTask for FixedFullZipDecodeTask {
3069 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3070 let estimated_size_bytes = self
3072 .data
3073 .iter()
3074 .map(|task_item| task_item.data.data_size() as usize)
3075 .sum::<usize>()
3076 * 2;
3077 let mut data_builder =
3078 DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3079
3080 if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3081 for task_item in self.data.into_iter() {
3085 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3086 unreachable!()
3087 };
3088 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3089 else {
3090 unreachable!()
3091 };
3092 debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3093 let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3094 data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3095 }
3096
3097 let unraveler = RepDefUnraveler::new(
3098 None,
3099 None,
3100 self.details.def_meaning.clone(),
3101 self.num_rows as u64,
3102 );
3103
3104 Ok(DecodedPage {
3105 data: data_builder.finish(),
3106 repdef: unraveler,
3107 })
3108 } else {
3109 let mut rep = Vec::with_capacity(self.num_rows);
3111 let mut def = Vec::with_capacity(self.num_rows);
3112
3113 for task_item in self.data.into_iter() {
3114 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3115 unreachable!()
3116 };
3117 let mut buf_slice = fixed_data.data.as_ref();
3118 let num_values = fixed_data.num_values as usize;
3119 let mut values = Vec::with_capacity(
3122 fixed_data.data.len()
3123 - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3124 );
3125 let mut visible_items = 0;
3126 for _ in 0..num_values {
3127 self.details
3129 .ctrl_word_parser
3130 .parse(buf_slice, &mut rep, &mut def);
3131 buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3132
3133 let is_visible = def
3134 .last()
3135 .map(|d| *d <= self.details.max_visible_def)
3136 .unwrap_or(true);
3137 if is_visible {
3138 values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3140 buf_slice = &buf_slice[self.bytes_per_value..];
3141 visible_items += 1;
3142 }
3143 }
3144
3145 let values_buf = LanceBuffer::from(values);
3147 let fixed_data = FixedWidthDataBlock {
3148 bits_per_value: self.bytes_per_value as u64 * 8,
3149 block_info: BlockInfo::new(),
3150 data: values_buf,
3151 num_values: visible_items,
3152 };
3153 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3154 else {
3155 unreachable!()
3156 };
3157 let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3158 data_builder.append(&decompressed, 0..visible_items);
3159 }
3160
3161 let repetition = if rep.is_empty() { None } else { Some(rep) };
3162 let definition = if def.is_empty() { None } else { Some(def) };
3163
3164 let unraveler = RepDefUnraveler::new(
3165 repetition,
3166 definition,
3167 self.details.def_meaning.clone(),
3168 self.num_rows as u64,
3169 );
3170 let data = data_builder.finish();
3171
3172 Ok(DecodedPage {
3173 data,
3174 repdef: unraveler,
3175 })
3176 }
3177 }
3178}
3179
3180#[derive(Debug)]
3181struct StructuralPrimitiveFieldSchedulingJob<'a> {
3182 scheduler: &'a StructuralPrimitiveFieldScheduler,
3183 ranges: Vec<Range<u64>>,
3184 page_idx: usize,
3185 range_idx: usize,
3186 global_row_offset: u64,
3187}
3188
3189impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3190 pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3191 Self {
3192 scheduler,
3193 ranges,
3194 page_idx: 0,
3195 range_idx: 0,
3196 global_row_offset: 0,
3197 }
3198 }
3199}
3200
3201impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3202 fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3203 if self.range_idx >= self.ranges.len() {
3204 return Ok(Vec::new());
3205 }
3206 let mut range = self.ranges[self.range_idx].clone();
3208 let priority = range.start;
3209
3210 let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3211 trace!(
3212 "Current range is {:?} and current page has {} rows",
3213 range, cur_page.num_rows
3214 );
3215 while cur_page.num_rows + self.global_row_offset <= range.start {
3217 self.global_row_offset += cur_page.num_rows;
3218 self.page_idx += 1;
3219 trace!("Skipping entire page of {} rows", cur_page.num_rows);
3220 cur_page = &self.scheduler.page_schedulers[self.page_idx];
3221 }
3222
3223 let mut ranges_in_page = Vec::new();
3227 while cur_page.num_rows + self.global_row_offset > range.start {
3228 range.start = range.start.max(self.global_row_offset);
3229 let start_in_page = range.start - self.global_row_offset;
3230 let end_in_page = start_in_page + (range.end - range.start);
3231 let end_in_page = end_in_page.min(cur_page.num_rows);
3232 let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3233
3234 ranges_in_page.push(start_in_page..end_in_page);
3235 if last_in_range {
3236 self.range_idx += 1;
3237 if self.range_idx == self.ranges.len() {
3238 break;
3239 }
3240 range = self.ranges[self.range_idx].clone();
3241 } else {
3242 break;
3243 }
3244 }
3245
3246 trace!(
3247 "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3248 ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3249 ranges_in_page.len(),
3250 cur_page.num_rows,
3251 priority,
3252 self.scheduler.column_index,
3253 cur_page.page_index,
3254 );
3255
3256 self.global_row_offset += cur_page.num_rows;
3257 self.page_idx += 1;
3258
3259 let page_decoders = cur_page
3260 .scheduler
3261 .schedule_ranges(&ranges_in_page, context.io())?;
3262
3263 let cur_path = context.current_path();
3264 page_decoders
3265 .into_iter()
3266 .map(|page_load_task| {
3267 let cur_path = cur_path.clone();
3268 let page_decoder = page_load_task.decoder_fut;
3269 let unloaded_page = async move {
3270 let page_decoder = page_decoder.await?;
3271 Ok(LoadedPageShard {
3272 decoder: page_decoder,
3273 path: cur_path,
3274 })
3275 }
3276 .boxed();
3277 Ok(ScheduledScanLine {
3278 decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3279 rows_scheduled: page_load_task.num_rows,
3280 })
3281 })
3282 .collect::<Result<Vec<_>>>()
3283 }
3284}
3285
3286#[derive(Debug)]
3287struct PageInfoAndScheduler {
3288 page_index: usize,
3289 num_rows: u64,
3290 scheduler: Box<dyn StructuralPageScheduler>,
3291}
3292
3293#[derive(Debug)]
3298pub struct StructuralPrimitiveFieldScheduler {
3299 page_schedulers: Vec<PageInfoAndScheduler>,
3300 column_index: u32,
3301 view_tag: String,
3307}
3308
3309impl StructuralPrimitiveFieldScheduler {
3310 pub fn try_new(
3311 column_info: &ColumnInfo,
3312 decompressors: &dyn DecompressionStrategy,
3313 cache_repetition_index: bool,
3314 target_field: &Field,
3315 ) -> Result<Self> {
3316 let page_schedulers = column_info
3317 .page_infos
3318 .iter()
3319 .enumerate()
3320 .map(|(page_index, page_info)| {
3321 Self::page_info_to_scheduler(
3322 page_info,
3323 page_index,
3324 decompressors,
3325 cache_repetition_index,
3326 target_field,
3327 )
3328 })
3329 .collect::<Result<Vec<_>>>()?;
3330 Ok(Self {
3331 page_schedulers,
3332 column_index: column_info.index,
3333 view_tag: format!("{:?}", target_field.data_type()),
3334 })
3335 }
3336
3337 fn page_layout_to_scheduler(
3338 page_info: &PageInfo,
3339 page_layout: &PageLayout,
3340 decompressors: &dyn DecompressionStrategy,
3341 cache_repetition_index: bool,
3342 target_field: &Field,
3343 ) -> Result<Box<dyn StructuralPageScheduler>> {
3344 use pb21::page_layout::Layout;
3345 Ok(match page_layout.layout.as_ref().expect_ok()? {
3346 Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
3347 &page_info.buffer_offsets_and_sizes,
3348 page_info.priority,
3349 mini_block.num_items,
3350 mini_block,
3351 decompressors,
3352 )?),
3353 Layout::FullZipLayout(full_zip) => {
3354 let mut scheduler = FullZipScheduler::try_new(
3355 &page_info.buffer_offsets_and_sizes,
3356 page_info.priority,
3357 page_info.num_rows,
3358 full_zip,
3359 decompressors,
3360 )?;
3361 scheduler.enable_cache = cache_repetition_index;
3362 Box::new(scheduler)
3363 }
3364 Layout::ConstantLayout(constant_layout) => {
3365 let def_meaning = constant_layout
3366 .layers
3367 .iter()
3368 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3369 .collect::<Vec<_>>();
3370 let has_scalar_value = constant_layout.inline_value.is_some()
3371 || page_info.buffer_offsets_and_sizes.len() == 1
3372 || page_info.buffer_offsets_and_sizes.len() == 3;
3373 if has_scalar_value {
3374 Box::new(constant::ConstantPageScheduler::try_new(
3375 page_info.buffer_offsets_and_sizes.clone(),
3376 constant_layout.inline_value.clone(),
3377 target_field.data_type(),
3378 def_meaning.into(),
3379 )?) as Box<dyn StructuralPageScheduler>
3380 } else if def_meaning.len() == 1
3381 && def_meaning[0] == DefinitionInterpretation::NullableItem
3382 {
3383 Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
3384 } else {
3385 let rep_decompressor = constant_layout
3386 .rep_compression
3387 .as_ref()
3388 .map(|encoding| decompressors.create_block_decompressor(encoding))
3389 .transpose()?
3390 .map(Arc::from);
3391
3392 let def_decompressor = constant_layout
3393 .def_compression
3394 .as_ref()
3395 .map(|encoding| decompressors.create_block_decompressor(encoding))
3396 .transpose()?
3397 .map(Arc::from);
3398
3399 Box::new(ComplexAllNullScheduler::new(
3400 page_info.buffer_offsets_and_sizes.clone(),
3401 def_meaning.into(),
3402 rep_decompressor,
3403 def_decompressor,
3404 constant_layout.num_rep_values,
3405 constant_layout.num_def_values,
3406 )) as Box<dyn StructuralPageScheduler>
3407 }
3408 }
3409 Layout::BlobLayout(blob) => {
3410 let inner_scheduler = Self::page_layout_to_scheduler(
3411 page_info,
3412 blob.inner_layout.as_ref().expect_ok()?.as_ref(),
3413 decompressors,
3414 cache_repetition_index,
3415 target_field,
3416 )?;
3417 let def_meaning = blob
3418 .layers
3419 .iter()
3420 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3421 .collect::<Vec<_>>();
3422 if matches!(target_field.data_type(), DataType::Struct(_)) {
3423 Box::new(BlobDescriptionPageScheduler::new(
3425 inner_scheduler,
3426 def_meaning.into(),
3427 ))
3428 } else {
3429 Box::new(BlobPageScheduler::new(
3431 inner_scheduler,
3432 page_info.priority,
3433 page_info.num_rows,
3434 def_meaning.into(),
3435 ))
3436 }
3437 }
3438 })
3439 }
3440
3441 fn page_info_to_scheduler(
3442 page_info: &PageInfo,
3443 page_index: usize,
3444 decompressors: &dyn DecompressionStrategy,
3445 cache_repetition_index: bool,
3446 target_field: &Field,
3447 ) -> Result<PageInfoAndScheduler> {
3448 let page_layout = page_info.encoding.as_structural();
3449 let scheduler = Self::page_layout_to_scheduler(
3450 page_info,
3451 page_layout,
3452 decompressors,
3453 cache_repetition_index,
3454 target_field,
3455 )?;
3456 Ok(PageInfoAndScheduler {
3457 page_index,
3458 num_rows: page_info.num_rows,
3459 scheduler,
3460 })
3461 }
3462}
3463
3464pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
3465 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
3466}
3467
3468pub struct NoCachedPageData;
3469
3470impl DeepSizeOf for NoCachedPageData {
3471 fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
3472 0
3473 }
3474}
3475impl CachedPageData for NoCachedPageData {
3476 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3477 self
3478 }
3479}
3480
3481pub struct CachedFieldData {
3482 pages: Vec<Arc<dyn CachedPageData>>,
3483}
3484
3485impl DeepSizeOf for CachedFieldData {
3486 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
3487 self.pages.deep_size_of_children(ctx)
3488 }
3489}
3490
3491#[derive(Debug, Clone)]
3501pub struct FieldDataCacheKey {
3502 pub column_index: u32,
3503 pub view_tag: String,
3504}
3505
3506impl CacheKey for FieldDataCacheKey {
3507 type ValueType = CachedFieldData;
3508
3509 fn key(&self) -> std::borrow::Cow<'_, str> {
3510 format!("{}:{}", self.column_index, self.view_tag).into()
3511 }
3512
3513 fn type_name() -> &'static str {
3514 "FieldData"
3515 }
3516}
3517
3518impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
3519 fn initialize<'a>(
3520 &'a mut self,
3521 _filter: &'a FilterExpression,
3522 context: &'a SchedulerContext,
3523 ) -> BoxFuture<'a, Result<()>> {
3524 let cache_key = FieldDataCacheKey {
3525 column_index: self.column_index,
3526 view_tag: self.view_tag.clone(),
3527 };
3528 let cache = context.cache().clone();
3529
3530 async move {
3531 if let Some(cached_data) = cache.get_with_key(&cache_key).await {
3532 self.page_schedulers
3533 .iter_mut()
3534 .zip(cached_data.pages.iter())
3535 .for_each(|(page_scheduler, cached_data)| {
3536 page_scheduler.scheduler.load(cached_data);
3537 });
3538 return Ok(());
3539 }
3540
3541 let page_data = self
3542 .page_schedulers
3543 .iter_mut()
3544 .map(|s| s.scheduler.initialize(context.io()))
3545 .collect::<FuturesOrdered<_>>();
3546
3547 let page_data = page_data.try_collect::<Vec<_>>().await?;
3548 let cached_data = Arc::new(CachedFieldData { pages: page_data });
3549 cache.insert_with_key(&cache_key, cached_data).await;
3550 Ok(())
3551 }
3552 .boxed()
3553 }
3554
3555 fn schedule_ranges<'a>(
3556 &'a self,
3557 ranges: &[Range<u64>],
3558 _filter: &FilterExpression,
3559 ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
3560 let ranges = ranges.to_vec();
3561 Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
3562 self, ranges,
3563 )))
3564 }
3565}
3566
3567#[derive(Debug)]
3570pub struct StructuralCompositeDecodeArrayTask {
3571 tasks: Vec<Box<dyn DecodePageTask>>,
3572 should_validate: bool,
3573 data_type: DataType,
3574}
3575
3576impl StructuralCompositeDecodeArrayTask {
3577 fn restore_validity(
3578 array: Arc<dyn Array>,
3579 unraveler: &mut CompositeRepDefUnraveler,
3580 ) -> Arc<dyn Array> {
3581 let validity = unraveler.unravel_validity(array.len());
3582 let Some(validity) = validity else {
3583 return array;
3584 };
3585 if array.data_type() == &DataType::Null {
3586 return array;
3588 }
3589 assert_eq!(validity.len(), array.len());
3590 make_array(unsafe {
3593 array
3594 .to_data()
3595 .into_builder()
3596 .nulls(Some(validity))
3597 .build_unchecked()
3598 })
3599 }
3600}
3601
3602impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
3603 fn decode(self: Box<Self>) -> Result<DecodedArray> {
3604 let mut arrays = Vec::with_capacity(self.tasks.len());
3605 let mut unravelers = Vec::with_capacity(self.tasks.len());
3606 let mut data_size = 0u64;
3607 for task in self.tasks {
3608 let decoded = task.decode()?;
3609 data_size += decoded.data.data_size();
3610 unravelers.push(decoded.repdef);
3611
3612 let array = make_array(
3613 decoded
3614 .data
3615 .into_arrow(self.data_type.clone(), self.should_validate)?,
3616 );
3617
3618 arrays.push(array);
3619 }
3620 let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
3621 let array = arrow_select::concat::concat(&array_refs)?;
3622 let mut repdef = CompositeRepDefUnraveler::new(unravelers);
3623
3624 let array = Self::restore_validity(array, &mut repdef);
3625
3626 Ok(DecodedArray {
3627 array,
3628 repdef,
3629 data_size,
3630 })
3631 }
3632}
3633
3634#[derive(Debug)]
3635pub struct StructuralPrimitiveFieldDecoder {
3636 field: Arc<ArrowField>,
3637 page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
3638 should_validate: bool,
3639 rows_drained_in_current: u64,
3640}
3641
3642impl StructuralPrimitiveFieldDecoder {
3643 pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
3644 Self {
3645 field: field.clone(),
3646 page_decoders: VecDeque::new(),
3647 should_validate,
3648 rows_drained_in_current: 0,
3649 }
3650 }
3651}
3652
3653impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
3654 fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
3655 assert!(child.path.is_empty());
3656 self.page_decoders.push_back(child.decoder);
3657 Ok(())
3658 }
3659
3660 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
3661 let mut remaining = num_rows;
3662 let mut tasks = Vec::new();
3663 while remaining > 0 {
3664 let cur_page = self.page_decoders.front_mut().unwrap();
3665 let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
3666 let to_take = num_in_page.min(remaining);
3667
3668 let task = cur_page.drain(to_take)?;
3669 tasks.push(task);
3670
3671 if to_take == num_in_page {
3672 self.page_decoders.pop_front();
3673 self.rows_drained_in_current = 0;
3674 } else {
3675 self.rows_drained_in_current += to_take;
3676 }
3677
3678 remaining -= to_take;
3679 }
3680 Ok(Box::new(StructuralCompositeDecodeArrayTask {
3681 tasks,
3682 should_validate: self.should_validate,
3683 data_type: self.field.data_type().clone(),
3684 }))
3685 }
3686
3687 fn data_type(&self) -> &DataType {
3688 self.field.data_type()
3689 }
3690}
3691
3692struct SerializedFullZip {
3694 values: LanceBuffer,
3696 repetition_index: Option<LanceBuffer>,
3698}
3699
3700const MINIBLOCK_ALIGNMENT: usize = 8;
3720
3721pub struct PrimitiveStructuralEncoder {
3748 accumulation_queue: AccumulationQueue,
3750
3751 keep_original_array: bool,
3752 support_large_chunk: bool,
3753 accumulated_repdefs: Vec<RepDefBuilder>,
3754 compression_strategy: Arc<dyn CompressionStrategy>,
3756 column_index: u32,
3757 field: Field,
3758 encoding_metadata: Arc<HashMap<String, String>>,
3759 version: LanceFileVersion,
3760}
3761
3762struct CompressedLevelsChunk {
3763 data: LanceBuffer,
3764 num_levels: u16,
3765}
3766
3767struct CompressedLevels {
3768 data: Vec<CompressedLevelsChunk>,
3769 compression: CompressiveEncoding,
3770 rep_index: Option<LanceBuffer>,
3771}
3772
3773struct SerializedMiniBlockPage {
3774 num_buffers: u64,
3775 data: LanceBuffer,
3776 metadata: LanceBuffer,
3777}
3778
3779#[derive(Debug, Clone, Copy)]
3780struct DictEncodingBudget {
3781 max_dict_entries: u32,
3782 max_encoded_size: usize,
3783}
3784
3785impl PrimitiveStructuralEncoder {
3786 pub fn try_new(
3787 options: &EncodingOptions,
3788 compression_strategy: Arc<dyn CompressionStrategy>,
3789 column_index: u32,
3790 field: Field,
3791 encoding_metadata: Arc<HashMap<String, String>>,
3792 ) -> Result<Self> {
3793 Ok(Self {
3794 accumulation_queue: AccumulationQueue::new(
3795 options.cache_bytes_per_column,
3796 column_index,
3797 options.keep_original_array,
3798 ),
3799 support_large_chunk: options.support_large_chunk(),
3800 keep_original_array: options.keep_original_array,
3801 accumulated_repdefs: Vec::new(),
3802 column_index,
3803 compression_strategy,
3804 field,
3805 encoding_metadata,
3806 version: options.version,
3807 })
3808 }
3809
3810 fn is_narrow(data_block: &DataBlock) -> bool {
3818 const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
3819
3820 if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
3821 let max_len_array = max_len_array
3822 .as_any()
3823 .downcast_ref::<PrimitiveArray<UInt64Type>>()
3824 .unwrap();
3825 if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
3826 return true;
3827 }
3828 }
3829 false
3830 }
3831
3832 fn prefers_miniblock(
3833 data_block: &DataBlock,
3834 encoding_metadata: &HashMap<String, String>,
3835 ) -> bool {
3836 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3838 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
3839 }
3840 Self::is_narrow(data_block)
3842 }
3843
3844 fn repdef_too_sparse_for_miniblock(
3870 repdef: &crate::repdef::SerializedRepDefs,
3871 num_values: u64,
3872 ) -> bool {
3873 if num_values == 0 {
3874 return false;
3875 }
3876 let num_levels = repdef
3877 .repetition_levels
3878 .as_ref()
3879 .map(|r| r.len() as u64)
3880 .max(repdef.definition_levels.as_ref().map(|d| d.len() as u64))
3881 .unwrap_or(0);
3882 if num_levels == 0 {
3883 return false;
3884 }
3885
3886 let bits_per_rep = repdef
3888 .repetition_levels
3889 .as_ref()
3890 .and_then(|r| r.iter().max().copied())
3891 .map(|max_val| u16::BITS - max_val.leading_zeros())
3892 .unwrap_or(0) as u64;
3893 let bits_per_def = repdef
3894 .definition_levels
3895 .as_ref()
3896 .and_then(|d| d.iter().max().copied())
3897 .map(|max_val| u16::BITS - max_val.leading_zeros())
3898 .unwrap_or(0) as u64;
3899
3900 let bits_per_level = bits_per_rep + bits_per_def;
3901 if bits_per_level == 0 {
3902 return false;
3903 }
3904
3905 const REPDEF_BUDGET_BITS: u64 = 16 * 1024 * 8;
3907 let max_levels_per_chunk = REPDEF_BUDGET_BITS / bits_per_level;
3908 let levels_per_chunk =
3909 (num_levels as f64 / num_values as f64) * *miniblock::MAX_MINIBLOCK_VALUES as f64;
3910 if levels_per_chunk > max_levels_per_chunk as f64 {
3911 return true;
3912 }
3913
3914 Self::any_chunk_levels_overflow_u16(repdef, num_values)
3920 }
3921
3922 fn any_chunk_levels_overflow_u16(
3934 repdef: &crate::repdef::SerializedRepDefs,
3935 num_values: u64,
3936 ) -> bool {
3937 let Some(max_visible) = repdef.max_visible_level else {
3941 return false;
3942 };
3943 let Some(def_levels) = repdef.definition_levels.as_ref() else {
3944 return false;
3945 };
3946 let u16_max = u16::MAX as u64;
3950 if (def_levels.len() as u64) <= u16_max {
3951 return false;
3952 }
3953 let max_visible_per_chunk = *miniblock::MAX_MINIBLOCK_VALUES;
3954
3955 let mut visible_in_chunk: u64 = 0;
3956 let mut levels_in_current_chunk: u64 = 0;
3957 let mut visible_seen_total: u64 = 0;
3958 for &def_lvl in def_levels.iter() {
3959 levels_in_current_chunk += 1;
3960 if levels_in_current_chunk > u16_max {
3963 return true;
3964 }
3965 if def_lvl <= max_visible {
3966 visible_in_chunk += 1;
3967 visible_seen_total += 1;
3968 let is_last_value = visible_seen_total == num_values;
3969 if visible_in_chunk == max_visible_per_chunk && !is_last_value {
3970 visible_in_chunk = 0;
3974 levels_in_current_chunk = 0;
3975 }
3976 }
3977 }
3978 levels_in_current_chunk > u16_max
3981 }
3982
3983 fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
3984 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3988 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
3989 }
3990 true
3991 }
3992
3993 fn serialize_miniblocks(
4040 miniblocks: MiniBlockCompressed,
4041 rep: Option<Vec<CompressedLevelsChunk>>,
4042 def: Option<Vec<CompressedLevelsChunk>>,
4043 support_large_chunk: bool,
4044 ) -> Result<SerializedMiniBlockPage> {
4045 let bytes_rep = rep
4046 .as_ref()
4047 .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
4048 .unwrap_or(0);
4049 let bytes_def = def
4050 .as_ref()
4051 .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
4052 .unwrap_or(0);
4053 let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
4054 let mut num_buffers = miniblocks.data.len();
4055 if rep.is_some() {
4056 num_buffers += 1;
4057 }
4058 if def.is_some() {
4059 num_buffers += 1;
4060 }
4061 let max_extra = 9 * num_buffers;
4063 let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
4064 let chunk_size_bytes = if support_large_chunk { 4 } else { 2 };
4065 let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
4066
4067 let mut rep_iter = rep.map(|r| r.into_iter());
4068 let mut def_iter = def.map(|d| d.into_iter());
4069
4070 let mut buffer_offsets = vec![0; miniblocks.data.len()];
4071 for chunk in miniblocks.chunks {
4072 let start_pos = data_buffer.len();
4073 debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
4075
4076 let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
4077 let def = def_iter.as_mut().map(|d| d.next().unwrap());
4078
4079 let num_levels = rep
4081 .as_ref()
4082 .map(|r| r.num_levels)
4083 .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
4084 data_buffer.extend_from_slice(&num_levels.to_le_bytes());
4085
4086 if let Some(rep) = rep.as_ref() {
4088 let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
4089 Error::internal(format!(
4090 "Repetition buffer size ({} bytes) too large",
4091 rep.data.len()
4092 ))
4093 })?;
4094 data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
4095 }
4096 if let Some(def) = def.as_ref() {
4097 let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
4098 Error::internal(format!(
4099 "Definition buffer size ({} bytes) too large",
4100 def.data.len()
4101 ))
4102 })?;
4103 data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4104 }
4105
4106 if support_large_chunk {
4107 for &buffer_size in &chunk.buffer_sizes {
4108 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4109 }
4110 } else {
4111 for &buffer_size in &chunk.buffer_sizes {
4112 data_buffer.extend_from_slice(&(buffer_size as u16).to_le_bytes());
4113 }
4114 }
4115
4116 let add_padding = |data_buffer: &mut Vec<u8>| {
4118 let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4119 data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4120 };
4121 add_padding(&mut data_buffer);
4122
4123 if let Some(rep) = rep.as_ref() {
4125 data_buffer.extend_from_slice(&rep.data);
4126 add_padding(&mut data_buffer);
4127 }
4128 if let Some(def) = def.as_ref() {
4129 data_buffer.extend_from_slice(&def.data);
4130 add_padding(&mut data_buffer);
4131 }
4132 for (buffer_size, (buffer, buffer_offset)) in chunk
4133 .buffer_sizes
4134 .iter()
4135 .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4136 {
4137 let start = *buffer_offset;
4138 let end = start + *buffer_size as usize;
4139 *buffer_offset += *buffer_size as usize;
4140 data_buffer.extend_from_slice(&buffer[start..end]);
4141 add_padding(&mut data_buffer);
4142 }
4143
4144 let chunk_bytes = data_buffer.len() - start_pos;
4145 let max_chunk_size = if support_large_chunk {
4146 4 * 1024 * 1024 * 1024 } else {
4148 32 * 1024 };
4150 assert!(chunk_bytes <= max_chunk_size);
4151 assert!(chunk_bytes > 0);
4152 assert_eq!(chunk_bytes % 8, 0);
4153 assert!(chunk.log_num_values <= 12);
4155 let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4159 let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4160
4161 let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4162 if support_large_chunk {
4163 meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4164 } else {
4165 meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4166 }
4167 }
4168
4169 let data_buffer = LanceBuffer::from(data_buffer);
4170 let metadata_buffer = LanceBuffer::from(meta_buffer);
4171
4172 Ok(SerializedMiniBlockPage {
4173 num_buffers: miniblocks.data.len() as u64,
4174 data: data_buffer,
4175 metadata: metadata_buffer,
4176 })
4177 }
4178
4179 fn compress_levels(
4184 mut levels: RepDefSlicer<'_>,
4185 num_elements: u64,
4186 compression_strategy: &dyn CompressionStrategy,
4187 chunks: &[MiniBlockChunk],
4188 max_rep: u16,
4190 ) -> Result<CompressedLevels> {
4191 let mut rep_index = if max_rep > 0 {
4192 Vec::with_capacity(chunks.len())
4193 } else {
4194 vec![]
4195 };
4196 let num_levels = levels.num_levels() as u64;
4198 let levels_buf = levels.all_levels().clone();
4199
4200 let mut fixed_width_block = FixedWidthDataBlock {
4201 data: levels_buf,
4202 bits_per_value: 16,
4203 num_values: num_levels,
4204 block_info: BlockInfo::new(),
4205 };
4206 fixed_width_block.compute_stat();
4208
4209 let levels_block = DataBlock::FixedWidth(fixed_width_block);
4210 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4211 let (compressor, compressor_desc) =
4213 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4214 let mut level_chunks = Vec::with_capacity(chunks.len());
4216 let mut values_counter = 0;
4217 for (chunk_idx, chunk) in chunks.iter().enumerate() {
4218 let chunk_num_values = chunk.num_values(values_counter, num_elements);
4219 debug_assert!(chunk_num_values > 0);
4220 values_counter += chunk_num_values;
4221 let chunk_levels = if chunk_idx < chunks.len() - 1 {
4222 levels.slice_next(chunk_num_values as usize)
4223 } else {
4224 levels.slice_rest()
4225 };
4226 let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4227 if max_rep > 0 {
4228 let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4238 let rep_values = rep_values.as_ref();
4239
4240 let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4243 let num_leftovers = if chunk_idx < chunks.len() - 1 {
4244 rep_values
4245 .iter()
4246 .rev()
4247 .position(|v| *v == max_rep)
4248 .map(|pos| pos + 1)
4250 .unwrap_or(rep_values.len())
4251 } else {
4252 0
4254 };
4255
4256 if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
4257 let rep_len = rep_index.len();
4261 if rep_index[rep_len - 1] != 0 {
4262 rep_index[rep_len - 2] += 1;
4264 rep_index[rep_len - 1] = 0;
4265 }
4266 }
4267
4268 if chunk_idx == chunks.len() - 1 {
4269 num_rows += 1;
4271 }
4272 rep_index.push(num_rows as u64);
4273 rep_index.push(num_leftovers as u64);
4274 }
4275 let mut chunk_fixed_width = FixedWidthDataBlock {
4276 data: chunk_levels,
4277 bits_per_value: 16,
4278 num_values: num_chunk_levels,
4279 block_info: BlockInfo::new(),
4280 };
4281 chunk_fixed_width.compute_stat();
4282 let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
4283 let compressed_levels = compressor.compress(chunk_levels_block)?;
4284 let num_levels_u16 = u16::try_from(num_chunk_levels).map_err(|_| {
4293 Error::invalid_input_source(
4294 format!(
4295 "miniblock chunk has {} levels but per-chunk header is u16 \
4296 (max {}); this shape needs fullzip encoding",
4297 num_chunk_levels,
4298 u16::MAX,
4299 )
4300 .into(),
4301 )
4302 })?;
4303 level_chunks.push(CompressedLevelsChunk {
4304 data: compressed_levels,
4305 num_levels: num_levels_u16,
4306 });
4307 }
4308 debug_assert_eq!(levels.num_levels_remaining(), 0);
4309 let rep_index = if rep_index.is_empty() {
4310 None
4311 } else {
4312 Some(LanceBuffer::reinterpret_vec(rep_index))
4313 };
4314 Ok(CompressedLevels {
4315 data: level_chunks,
4316 compression: compressor_desc,
4317 rep_index,
4318 })
4319 }
4320
4321 fn encode_simple_all_null(
4322 column_idx: u32,
4323 num_rows: u64,
4324 row_number: u64,
4325 ) -> Result<EncodedPage> {
4326 let description =
4327 ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
4328 Ok(EncodedPage {
4329 column_idx,
4330 data: vec![],
4331 description: PageEncoding::Structural(description),
4332 num_rows,
4333 row_number,
4334 })
4335 }
4336
4337 fn encode_complex_all_null_vals(
4338 data: &Arc<[u16]>,
4339 compression_strategy: &dyn CompressionStrategy,
4340 ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
4341 let buffer = LanceBuffer::reinterpret_slice(data.clone());
4342 let mut fixed_width_block = FixedWidthDataBlock {
4343 data: buffer,
4344 bits_per_value: 16,
4345 num_values: data.len() as u64,
4346 block_info: BlockInfo::new(),
4347 };
4348 fixed_width_block.compute_stat();
4349
4350 let levels_block = DataBlock::FixedWidth(fixed_width_block);
4351 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4352 let (compressor, encoding) =
4353 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4354 let compressed_buffer = compressor.compress(levels_block)?;
4355 Ok((compressed_buffer, encoding))
4356 }
4357
4358 fn encode_complex_all_null(
4362 column_idx: u32,
4363 repdef: crate::repdef::SerializedRepDefs,
4364 row_number: u64,
4365 num_rows: u64,
4366 version: LanceFileVersion,
4367 compression_strategy: &dyn CompressionStrategy,
4368 ) -> Result<EncodedPage> {
4369 if version.resolve() < LanceFileVersion::V2_2 {
4370 let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
4371 LanceBuffer::reinterpret_slice(rep.clone())
4372 } else {
4373 LanceBuffer::empty()
4374 };
4375
4376 let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
4377 LanceBuffer::reinterpret_slice(def.clone())
4378 } else {
4379 LanceBuffer::empty()
4380 };
4381
4382 let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
4383 return Ok(EncodedPage {
4384 column_idx,
4385 data: vec![rep_bytes, def_bytes],
4386 description: PageEncoding::Structural(description),
4387 num_rows,
4388 row_number,
4389 });
4390 }
4391
4392 let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
4393 repdef.repetition_levels.as_ref()
4394 {
4395 let num_values = rep.len() as u64;
4396 let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
4397 (buffer, Some(encoding), num_values)
4398 } else {
4399 (LanceBuffer::empty(), None, 0)
4400 };
4401
4402 let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
4403 repdef.definition_levels.as_ref()
4404 {
4405 let num_values = def.len() as u64;
4406 let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
4407 (buffer, Some(encoding), num_values)
4408 } else {
4409 (LanceBuffer::empty(), None, 0)
4410 };
4411
4412 let description = ProtobufUtils21::compressed_all_null_constant_layout(
4413 &repdef.def_meaning,
4414 rep_encoding,
4415 def_encoding,
4416 num_rep_values,
4417 num_def_values,
4418 );
4419 Ok(EncodedPage {
4420 column_idx,
4421 data: vec![rep_bytes, def_bytes],
4422 description: PageEncoding::Structural(description),
4423 num_rows,
4424 row_number,
4425 })
4426 }
4427
4428 fn leaf_validity(
4429 repdef: &crate::repdef::SerializedRepDefs,
4430 num_values: usize,
4431 ) -> Result<Option<BooleanBuffer>> {
4432 let rep = repdef
4433 .repetition_levels
4434 .as_ref()
4435 .map(|rep| rep.as_ref().to_vec());
4436 let def = repdef
4437 .definition_levels
4438 .as_ref()
4439 .map(|def| def.as_ref().to_vec());
4440 let mut unraveler = RepDefUnraveler::new(
4441 rep,
4442 def,
4443 repdef.def_meaning.clone().into(),
4444 num_values as u64,
4445 );
4446 if unraveler.is_all_valid() {
4447 return Ok(None);
4448 }
4449 let mut validity = BooleanBufferBuilder::new(num_values);
4450 unraveler.unravel_validity(&mut validity);
4451 Ok(Some(validity.finish()))
4452 }
4453
4454 fn is_constant_values(
4455 arrays: &[ArrayRef],
4456 scalar: &ArrayRef,
4457 validity: Option<&BooleanBuffer>,
4458 ) -> Result<bool> {
4459 debug_assert_eq!(scalar.len(), 1);
4460 debug_assert_eq!(scalar.null_count(), 0);
4461
4462 match scalar.data_type() {
4463 DataType::Boolean => {
4464 let mut global_idx = 0usize;
4465 let scalar_val = scalar.as_boolean().value(0);
4466 for arr in arrays {
4467 let bool_arr = arr.as_boolean();
4468 for i in 0..arr.len() {
4469 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4470 global_idx += 1;
4471 if !is_valid {
4472 continue;
4473 }
4474 if bool_arr.value(i) != scalar_val {
4475 return Ok(false);
4476 }
4477 }
4478 }
4479 Ok(true)
4480 }
4481 DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
4482 DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
4483 DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
4484 DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
4485 data_type => {
4486 let mut global_idx = 0usize;
4487 let Some(byte_width) = data_type.byte_width_opt() else {
4488 return Ok(false);
4489 };
4490 let scalar_data = scalar.to_data();
4491 if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
4492 return Ok(false);
4493 }
4494 let scalar_bytes = scalar_data.buffers()[0].as_slice();
4495 if scalar_bytes.len() != byte_width {
4496 return Ok(false);
4497 }
4498
4499 for arr in arrays {
4500 let data = arr.to_data();
4501 if data.buffers().is_empty() {
4502 return Ok(false);
4503 }
4504 let buf = data.buffers()[0].as_slice();
4505 let base = data.offset();
4506 for i in 0..arr.len() {
4507 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4508 global_idx += 1;
4509 if !is_valid {
4510 continue;
4511 }
4512 let start = (base + i) * byte_width;
4513 if buf[start..start + byte_width] != scalar_bytes[..] {
4514 return Ok(false);
4515 }
4516 }
4517 }
4518 Ok(true)
4519 }
4520 }
4521 }
4522
4523 fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
4524 arrays: &[ArrayRef],
4525 scalar: &ArrayRef,
4526 validity: Option<&BooleanBuffer>,
4527 ) -> Result<bool> {
4528 debug_assert_eq!(scalar.len(), 1);
4529 let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
4530 let mut global_idx = 0usize;
4531 for arr in arrays {
4532 let str_arr = arr.as_string::<O>();
4533 for i in 0..arr.len() {
4534 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4535 global_idx += 1;
4536 if !is_valid {
4537 continue;
4538 }
4539 if str_arr.value(i).as_bytes() != scalar_val {
4540 return Ok(false);
4541 }
4542 }
4543 }
4544 Ok(true)
4545 }
4546
4547 fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
4548 arrays: &[ArrayRef],
4549 scalar: &ArrayRef,
4550 validity: Option<&BooleanBuffer>,
4551 ) -> Result<bool> {
4552 debug_assert_eq!(scalar.len(), 1);
4553 let scalar_val = scalar.as_binary::<O>().value(0);
4554 let mut global_idx = 0usize;
4555 for arr in arrays {
4556 let bin_arr = arr.as_binary::<O>();
4557 for i in 0..arr.len() {
4558 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4559 global_idx += 1;
4560 if !is_valid {
4561 continue;
4562 }
4563 if bin_arr.value(i) != scalar_val {
4564 return Ok(false);
4565 }
4566 }
4567 }
4568 Ok(true)
4569 }
4570
4571 fn find_constant_scalar(
4572 arrays: &[ArrayRef],
4573 validity: Option<&BooleanBuffer>,
4574 ) -> Result<Option<ArrayRef>> {
4575 if arrays.is_empty() {
4576 return Ok(None);
4577 }
4578
4579 let global_scalar_idx = if let Some(validity) = validity {
4580 let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
4581 return Ok(None);
4582 };
4583 idx
4584 } else {
4585 0
4586 };
4587
4588 let mut idx_remaining = global_scalar_idx;
4589 let mut scalar_arr_idx = 0usize;
4590 while scalar_arr_idx < arrays.len() {
4591 let len = arrays[scalar_arr_idx].len();
4592 if idx_remaining < len {
4593 break;
4594 }
4595 idx_remaining -= len;
4596 scalar_arr_idx += 1;
4597 }
4598
4599 if scalar_arr_idx >= arrays.len() {
4600 return Ok(None);
4601 }
4602
4603 let scalar =
4604 lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
4605 if scalar.null_count() != 0 {
4606 return Ok(None);
4607 }
4608 if !Self::is_constant_values(arrays, &scalar, validity)? {
4609 return Ok(None);
4610 }
4611 Ok(Some(scalar))
4612 }
4613
4614 fn resolve_dict_values_compression_metadata(
4615 field_metadata: &HashMap<String, String>,
4616 env_compression: Option<String>,
4617 env_compression_level: Option<String>,
4618 ) -> HashMap<String, String> {
4619 let mut metadata = HashMap::new();
4620
4621 let compression = field_metadata
4622 .get(DICT_VALUES_COMPRESSION_META_KEY)
4623 .cloned()
4624 .or(env_compression)
4625 .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
4626 metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
4627
4628 if let Some(compression_level) = field_metadata
4629 .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
4630 .cloned()
4631 .or(env_compression_level)
4632 {
4633 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
4634 }
4635
4636 metadata
4637 }
4638
4639 fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
4640 let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
4645 dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
4646 &field.metadata,
4647 env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
4648 env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
4649 );
4650 Ok(dict_values_field)
4651 }
4652
4653 #[allow(clippy::too_many_arguments)]
4654 fn encode_miniblock(
4655 column_idx: u32,
4656 field: &Field,
4657 compression_strategy: &dyn CompressionStrategy,
4658 data: DataBlock,
4659 repdef: crate::repdef::SerializedRepDefs,
4660 row_number: u64,
4661 dictionary_data: Option<DataBlock>,
4662 num_rows: u64,
4663 support_large_chunk: bool,
4664 ) -> Result<EncodedPage> {
4665 if let DataBlock::AllNull(_null_block) = data {
4666 unreachable!()
4669 }
4670
4671 let num_items = data.num_values();
4672
4673 let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
4674 let (compressed_data, value_encoding) = compressor.compress(data)?;
4675
4676 let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
4677
4678 let mut compressed_rep = repdef
4679 .rep_slicer()
4680 .map(|rep_slicer| {
4681 Self::compress_levels(
4682 rep_slicer,
4683 num_items,
4684 compression_strategy,
4685 &compressed_data.chunks,
4686 max_rep,
4687 )
4688 })
4689 .transpose()?;
4690
4691 let (rep_index, rep_index_depth) =
4692 match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
4693 Some(rep_index) => (Some(rep_index.clone()), 1),
4694 None => (None, 0),
4695 };
4696
4697 let mut compressed_def = repdef
4698 .def_slicer()
4699 .map(|def_slicer| {
4700 Self::compress_levels(
4701 def_slicer,
4702 num_items,
4703 compression_strategy,
4704 &compressed_data.chunks,
4705 0,
4706 )
4707 })
4708 .transpose()?;
4709
4710 let rep_data = compressed_rep
4716 .as_mut()
4717 .map(|cr| std::mem::take(&mut cr.data));
4718 let def_data = compressed_def
4719 .as_mut()
4720 .map(|cd| std::mem::take(&mut cd.data));
4721
4722 let serialized =
4723 Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?;
4724
4725 let mut data = Vec::with_capacity(4);
4727 data.push(serialized.metadata);
4728 data.push(serialized.data);
4729
4730 if let Some(dictionary_data) = dictionary_data {
4731 let num_dictionary_items = dictionary_data.num_values();
4732 let dict_values_field = Self::build_dict_values_compressor_field(field)?;
4733
4734 let (compressor, dictionary_encoding) = compression_strategy
4735 .create_block_compressor(&dict_values_field, &dictionary_data)?;
4736 let dictionary_buffer = compressor.compress(dictionary_data)?;
4737
4738 data.push(dictionary_buffer);
4739 if let Some(rep_index) = rep_index {
4740 data.push(rep_index);
4741 }
4742
4743 let description = ProtobufUtils21::miniblock_layout(
4744 compressed_rep.map(|cr| cr.compression),
4745 compressed_def.map(|cd| cd.compression),
4746 value_encoding,
4747 rep_index_depth,
4748 serialized.num_buffers,
4749 Some((dictionary_encoding, num_dictionary_items)),
4750 &repdef.def_meaning,
4751 num_items,
4752 support_large_chunk,
4753 );
4754 Ok(EncodedPage {
4755 num_rows,
4756 column_idx,
4757 data,
4758 description: PageEncoding::Structural(description),
4759 row_number,
4760 })
4761 } else {
4762 let description = ProtobufUtils21::miniblock_layout(
4763 compressed_rep.map(|cr| cr.compression),
4764 compressed_def.map(|cd| cd.compression),
4765 value_encoding,
4766 rep_index_depth,
4767 serialized.num_buffers,
4768 None,
4769 &repdef.def_meaning,
4770 num_items,
4771 support_large_chunk,
4772 );
4773
4774 if let Some(rep_index) = rep_index {
4775 let view = rep_index.borrow_to_typed_slice::<u64>();
4776 let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
4777 debug_assert_eq!(total, num_rows);
4778
4779 data.push(rep_index);
4780 }
4781
4782 Ok(EncodedPage {
4783 num_rows,
4784 column_idx,
4785 data,
4786 description: PageEncoding::Structural(description),
4787 row_number,
4788 })
4789 }
4790 }
4791
4792 fn serialize_full_zip_fixed(
4794 fixed: FixedWidthDataBlock,
4795 mut repdef: ControlWordIterator,
4796 num_values: u64,
4797 ) -> SerializedFullZip {
4798 let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
4799 let mut zipped_data = Vec::with_capacity(len);
4800
4801 let max_rep_index_val = if repdef.has_repetition() {
4802 len as u64
4803 } else {
4804 0
4806 };
4807 let mut rep_index_builder =
4808 BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
4809
4810 assert_eq!(
4813 fixed.bits_per_value % 8,
4814 0,
4815 "Non-byte aligned full-zip compression not yet supported"
4816 );
4817
4818 let bytes_per_value = fixed.bits_per_value as usize / 8;
4819 let mut offset = 0;
4820
4821 if bytes_per_value == 0 {
4822 while let Some(control) = repdef.append_next(&mut zipped_data) {
4824 if control.is_new_row {
4825 debug_assert!(offset <= len);
4827 unsafe { rep_index_builder.append(offset as u64) };
4829 }
4830 offset = zipped_data.len();
4831 }
4832 } else {
4833 let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
4835 while let Some(control) = repdef.append_next(&mut zipped_data) {
4836 if control.is_new_row {
4837 debug_assert!(offset <= len);
4839 unsafe { rep_index_builder.append(offset as u64) };
4841 }
4842 if control.is_visible {
4843 let value = data_iter.next().unwrap();
4844 zipped_data.extend_from_slice(value);
4845 }
4846 offset = zipped_data.len();
4847 }
4848 }
4849
4850 debug_assert_eq!(zipped_data.len(), len);
4851 unsafe {
4854 rep_index_builder.append(zipped_data.len() as u64);
4855 }
4856
4857 let zipped_data = LanceBuffer::from(zipped_data);
4858 let rep_index = rep_index_builder.into_data();
4859 let rep_index = if rep_index.is_empty() {
4860 None
4861 } else {
4862 Some(LanceBuffer::from(rep_index))
4863 };
4864 SerializedFullZip {
4865 values: zipped_data,
4866 repetition_index: rep_index,
4867 }
4868 }
4869
4870 fn serialize_full_zip_variable(
4874 variable: VariableWidthBlock,
4875 mut repdef: ControlWordIterator,
4876 num_items: u64,
4877 ) -> SerializedFullZip {
4878 let bytes_per_offset = variable.bits_per_offset as usize / 8;
4879 assert_eq!(
4880 variable.bits_per_offset % 8,
4881 0,
4882 "Only byte-aligned offsets supported"
4883 );
4884 let len = variable.data.len()
4885 + repdef.bytes_per_word() * num_items as usize
4886 + bytes_per_offset * variable.num_values as usize;
4887 let mut buf = Vec::with_capacity(len);
4888
4889 let max_rep_index_val = len as u64;
4890 let mut rep_index_builder =
4891 BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
4892
4893 match bytes_per_offset {
4895 4 => {
4896 let offs = variable.offsets.borrow_to_typed_slice::<u32>();
4897 let mut rep_offset = 0;
4898 let mut windows_iter = offs.as_ref().windows(2);
4899 while let Some(control) = repdef.append_next(&mut buf) {
4900 if control.is_new_row {
4901 debug_assert!(rep_offset <= len);
4903 unsafe { rep_index_builder.append(rep_offset as u64) };
4905 }
4906 if control.is_visible {
4907 let window = windows_iter.next().unwrap();
4908 if control.is_valid_item {
4909 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4910 buf.extend_from_slice(
4911 &variable.data[window[0] as usize..window[1] as usize],
4912 );
4913 }
4914 }
4915 rep_offset = buf.len();
4916 }
4917 }
4918 8 => {
4919 let offs = variable.offsets.borrow_to_typed_slice::<u64>();
4920 let mut rep_offset = 0;
4921 let mut windows_iter = offs.as_ref().windows(2);
4922 while let Some(control) = repdef.append_next(&mut buf) {
4923 if control.is_new_row {
4924 debug_assert!(rep_offset <= len);
4926 unsafe { rep_index_builder.append(rep_offset as u64) };
4928 }
4929 if control.is_visible {
4930 let window = windows_iter.next().unwrap();
4931 if control.is_valid_item {
4932 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4933 buf.extend_from_slice(
4934 &variable.data[window[0] as usize..window[1] as usize],
4935 );
4936 }
4937 }
4938 rep_offset = buf.len();
4939 }
4940 }
4941 _ => panic!("Unsupported offset size"),
4942 }
4943
4944 debug_assert!(buf.len() <= len);
4947 unsafe {
4950 rep_index_builder.append(buf.len() as u64);
4951 }
4952
4953 let zipped_data = LanceBuffer::from(buf);
4954 let rep_index = rep_index_builder.into_data();
4955 debug_assert!(!rep_index.is_empty());
4956 let rep_index = Some(LanceBuffer::from(rep_index));
4957 SerializedFullZip {
4958 values: zipped_data,
4959 repetition_index: rep_index,
4960 }
4961 }
4962
4963 fn serialize_full_zip(
4966 compressed_data: PerValueDataBlock,
4967 repdef: ControlWordIterator,
4968 num_items: u64,
4969 ) -> SerializedFullZip {
4970 match compressed_data {
4971 PerValueDataBlock::Fixed(fixed) => {
4972 Self::serialize_full_zip_fixed(fixed, repdef, num_items)
4973 }
4974 PerValueDataBlock::Variable(var) => {
4975 Self::serialize_full_zip_variable(var, repdef, num_items)
4976 }
4977 }
4978 }
4979
4980 fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
4981 debug_assert_eq!(fixed.bits_per_value, 1);
4982 let num_values = fixed.num_values as usize;
4983 let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
4984 let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
4985 FixedWidthDataBlock {
4986 data: LanceBuffer::from(expanded),
4987 bits_per_value: 8,
4988 num_values: fixed.num_values,
4989 block_info: BlockInfo::new(),
4990 }
4991 }
4992
4993 fn encode_full_zip(
4994 column_idx: u32,
4995 field: &Field,
4996 compression_strategy: &dyn CompressionStrategy,
4997 data: DataBlock,
4998 repdef: crate::repdef::SerializedRepDefs,
4999 row_number: u64,
5000 num_lists: u64,
5001 ) -> Result<EncodedPage> {
5002 let max_rep = repdef
5003 .repetition_levels
5004 .as_ref()
5005 .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
5006 let max_def = repdef
5007 .definition_levels
5008 .as_ref()
5009 .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
5010
5011 let (num_items, num_visible_items) =
5015 if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
5016 (rep_levels.len() as u64, data.num_values())
5019 } else {
5020 (data.num_values(), data.num_values())
5022 };
5023
5024 let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
5025
5026 let repdef_iter = build_control_word_iterator(
5027 repdef.repetition_levels.as_deref(),
5028 max_rep,
5029 repdef.definition_levels.as_deref(),
5030 max_def,
5031 max_visible_def,
5032 num_items as usize,
5033 );
5034 let bits_rep = repdef_iter.bits_rep();
5035 let bits_def = repdef_iter.bits_def();
5036
5037 let data = match data {
5039 DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
5040 DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
5041 }
5042 other => other,
5043 };
5044
5045 let compressor = compression_strategy.create_per_value(field, &data)?;
5046 let (compressed_data, value_encoding) = compressor.compress(data)?;
5047
5048 let description = match &compressed_data {
5049 PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
5050 bits_rep,
5051 bits_def,
5052 fixed.bits_per_value as u32,
5053 value_encoding,
5054 &repdef.def_meaning,
5055 num_items as u32,
5056 num_visible_items as u32,
5057 ),
5058 PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
5059 bits_rep,
5060 bits_def,
5061 variable.bits_per_offset as u32,
5062 value_encoding,
5063 &repdef.def_meaning,
5064 num_items as u32,
5065 num_visible_items as u32,
5066 ),
5067 };
5068
5069 let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items);
5070
5071 let data = if let Some(repindex) = zipped.repetition_index {
5072 vec![zipped.values, repindex]
5073 } else {
5074 vec![zipped.values]
5075 };
5076
5077 Ok(EncodedPage {
5078 num_rows: num_lists,
5079 column_idx,
5080 data,
5081 description: PageEncoding::Structural(description),
5082 row_number,
5083 })
5084 }
5085
5086 fn should_dictionary_encode(
5087 data_block: &DataBlock,
5088 field: &Field,
5089 version: LanceFileVersion,
5090 ) -> Option<DictEncodingBudget> {
5091 const DEFAULT_SAMPLE_SIZE: usize = 4096;
5092 const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5093
5094 match data_block {
5097 DataBlock::FixedWidth(fixed) => {
5098 if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
5099 return None;
5100 }
5101 if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5102 return None;
5103 }
5104 if fixed.bits_per_value % 8 != 0 {
5105 return None;
5106 }
5107 }
5108 DataBlock::VariableWidth(var) => {
5109 if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5110 return None;
5111 }
5112 }
5113 _ => return None,
5114 }
5115
5116 let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5118 .ok()
5119 .and_then(|val| val.parse().ok())
5120 .unwrap_or(100);
5121 if data_block.num_values() < too_small {
5122 return None;
5123 }
5124
5125 let num_values = data_block.num_values();
5126
5127 let divisor: u64 = field
5130 .metadata
5131 .get(DICT_DIVISOR_META_KEY)
5132 .and_then(|val| val.parse().ok())
5133 .or_else(|| {
5134 env::var("LANCE_ENCODING_DICT_DIVISOR")
5135 .ok()
5136 .and_then(|val| val.parse().ok())
5137 })
5138 .unwrap_or(DEFAULT_DICT_DIVISOR);
5139
5140 let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5141 .ok()
5142 .and_then(|val| val.parse().ok())
5143 .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5144
5145 let threshold_cardinality = num_values
5146 .checked_div(divisor.max(1))
5147 .unwrap_or(0)
5148 .min(max_cardinality);
5149 if threshold_cardinality == 0 {
5150 return None;
5151 }
5152
5153 let threshold_ratio = field
5155 .metadata
5156 .get(DICT_SIZE_RATIO_META_KEY)
5157 .and_then(|val| val.parse::<f64>().ok())
5158 .or_else(|| {
5159 env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5160 .ok()
5161 .and_then(|val| val.parse().ok())
5162 })
5163 .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5164
5165 if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5166 panic!(
5167 "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5168 threshold_ratio
5169 );
5170 }
5171
5172 let data_size = data_block.data_size();
5173 if data_size == 0 {
5174 return None;
5175 }
5176
5177 let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5178 let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5179
5180 if Self::sample_is_near_unique(
5182 data_block,
5183 DEFAULT_SAMPLE_SIZE,
5184 DEFAULT_SAMPLE_UNIQUE_RATIO,
5185 )? {
5186 return None;
5187 }
5188
5189 let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5190 Some(DictEncodingBudget {
5191 max_dict_entries,
5192 max_encoded_size,
5193 })
5194 }
5195
5196 fn sample_is_near_unique(
5202 data_block: &DataBlock,
5203 max_samples: usize,
5204 unique_ratio_threshold: f64,
5205 ) -> Option<bool> {
5206 use std::collections::HashSet;
5207
5208 if unique_ratio_threshold <= 0.0 || unique_ratio_threshold > 1.0 {
5209 return None;
5210 }
5211
5212 let num_values = usize::try_from(data_block.num_values()).ok()?;
5213 if num_values == 0 {
5214 return Some(false);
5215 }
5216
5217 let sample_count = num_values.min(max_samples).max(1);
5218 let step = (num_values / sample_count).max(1);
5220
5221 match data_block {
5222 DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
5223 64 => {
5224 let values = fixed.data.borrow_to_typed_slice::<u64>();
5225 let values = values.as_ref();
5226 let mut unique: HashSet<u64> = HashSet::with_capacity(sample_count.min(1024));
5227 for idx in (0..num_values).step_by(step).take(sample_count) {
5228 unique.insert(values.get(idx).copied()?);
5229 }
5230 let ratio = unique.len() as f64 / sample_count as f64;
5231 Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5233 }
5234 128 => {
5235 let values = fixed.data.borrow_to_typed_slice::<u128>();
5236 let values = values.as_ref();
5237 let mut unique: HashSet<u128> = HashSet::with_capacity(sample_count.min(1024));
5238 for idx in (0..num_values).step_by(step).take(sample_count) {
5239 unique.insert(values.get(idx).copied()?);
5240 }
5241 let ratio = unique.len() as f64 / sample_count as f64;
5242 Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5243 }
5244 _ => Some(false),
5245 },
5246 DataBlock::VariableWidth(var) => {
5247 use xxhash_rust::xxh3::xxh3_64;
5248
5249 let mut unique: HashSet<u64> = HashSet::with_capacity(sample_count.min(1024));
5251 match var.bits_per_offset {
5252 32 => {
5253 let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
5254 let offsets: &[u32] = offsets_ref.as_ref();
5255 for i in (0..num_values).step_by(step).take(sample_count) {
5256 let start = usize::try_from(*offsets.get(i)?).ok()?;
5257 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5258 if start > end || end > var.data.len() {
5259 return None;
5260 }
5261 unique.insert(xxh3_64(&var.data[start..end]));
5262 }
5263 }
5264 64 => {
5265 let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
5266 let offsets: &[u64] = offsets_ref.as_ref();
5267 for i in (0..num_values).step_by(step).take(sample_count) {
5268 let start = usize::try_from(*offsets.get(i)?).ok()?;
5269 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5270 if start > end || end > var.data.len() {
5271 return None;
5272 }
5273 unique.insert(xxh3_64(&var.data[start..end]));
5274 }
5275 }
5276 _ => return Some(false),
5277 }
5278 let ratio = unique.len() as f64 / sample_count as f64;
5279 Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5280 }
5281 _ => Some(false),
5282 }
5283 }
5284
5285 fn do_flush(
5287 &mut self,
5288 arrays: Vec<ArrayRef>,
5289 repdefs: Vec<RepDefBuilder>,
5290 row_number: u64,
5291 num_rows: u64,
5292 ) -> Result<Vec<EncodeTask>> {
5293 let column_idx = self.column_index;
5294 let compression_strategy = self.compression_strategy.clone();
5295 let field = self.field.clone();
5296 let encoding_metadata = self.encoding_metadata.clone();
5297 let support_large_chunk = self.support_large_chunk;
5298 let version = self.version;
5299 let task = spawn_cpu(move || {
5300 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5301 let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
5302 let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
5303 let repdef = RepDefBuilder::serialize(repdefs);
5304
5305 if num_values == 0 {
5306 log::debug!("Encoding column {} with {} items ({} rows) using complex-null layout", column_idx, num_values, num_rows);
5310 return Self::encode_complex_all_null(
5311 column_idx,
5312 repdef,
5313 row_number,
5314 num_rows,
5315 version,
5316 compression_strategy.as_ref(),
5317 );
5318 }
5319
5320 let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
5321 let all_null = leaf_validity
5322 .as_ref()
5323 .map(|validity| validity.count_set_bits() == 0)
5324 .unwrap_or(false);
5325
5326 if all_null {
5327 return if is_simple_validity {
5328 log::debug!(
5329 "Encoding column {} with {} items ({} rows) using simple-null layout",
5330 column_idx,
5331 num_values,
5332 num_rows
5333 );
5334 Self::encode_simple_all_null(column_idx, num_values, row_number)
5335 } else {
5336 log::debug!(
5337 "Encoding column {} with {} items ({} rows) using complex-null layout",
5338 column_idx,
5339 num_values,
5340 num_rows
5341 );
5342 Self::encode_complex_all_null(
5343 column_idx,
5344 repdef,
5345 row_number,
5346 num_rows,
5347 version,
5348 compression_strategy.as_ref(),
5349 )
5350 };
5351 }
5352
5353 if let DataType::Struct(fields) = &field.data_type()
5354 && fields.is_empty()
5355 {
5356 if has_repdef_info {
5357 return Err(Error::invalid_input_source(format!("Empty structs with rep/def information are not yet supported. The field {} is an empty struct that either has nulls or is in a list.", field.name).into()));
5358 }
5359 return Self::encode_simple_all_null(column_idx, num_values, row_number);
5362 }
5363
5364 let data_block = DataBlock::from_arrays(&arrays, num_values);
5365
5366 if version.resolve() >= LanceFileVersion::V2_2
5367 && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())?
5368 {
5369 log::debug!(
5370 "Encoding column {} with {} items ({} rows) using constant layout",
5371 column_idx,
5372 num_values,
5373 num_rows
5374 );
5375 return constant::encode_constant_page(
5376 column_idx,
5377 scalar,
5378 repdef,
5379 row_number,
5380 num_rows,
5381 );
5382 }
5383
5384 let requires_full_zip_packed_struct =
5385 if let DataBlock::Struct(ref struct_data_block) = data_block {
5386 struct_data_block.has_variable_width_child()
5387 } else {
5388 false
5389 };
5390
5391 if requires_full_zip_packed_struct {
5392 log::debug!(
5393 "Encoding column {} with {} items using full-zip packed struct layout",
5394 column_idx,
5395 num_values
5396 );
5397 return Self::encode_full_zip(
5398 column_idx,
5399 &field,
5400 compression_strategy.as_ref(),
5401 data_block,
5402 repdef,
5403 row_number,
5404 num_rows,
5405 );
5406 }
5407
5408 let too_sparse = Self::repdef_too_sparse_for_miniblock(&repdef, num_values);
5412
5413 if !too_sparse {
5414 if let DataBlock::Dictionary(dict) = data_block {
5415 log::debug!("Encoding column {} with {} items using dictionary encoding (already dictionary encoded)", column_idx, num_values);
5416 let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
5417 indices_data_block.compute_stat();
5422 return Self::encode_miniblock(
5423 column_idx,
5424 &field,
5425 compression_strategy.as_ref(),
5426 indices_data_block,
5427 repdef,
5428 row_number,
5429 Some(dictionary_data_block),
5430 num_rows,
5431 support_large_chunk,
5432 );
5433 }
5434 } else {
5435 log::debug!(
5436 "Encoding column {} with {} items using full-zip layout \
5437 (rep/def too sparse for mini-block)",
5438 column_idx,
5439 num_values
5440 );
5441 }
5442
5443 {
5444 let dict_result = if too_sparse {
5447 None
5448 } else {
5449 Self::should_dictionary_encode(&data_block, &field, version)
5450 .and_then(|budget| {
5451 log::debug!(
5452 "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
5453 column_idx,
5454 num_values
5455 );
5456 dict::dictionary_encode(
5457 &data_block,
5458 budget.max_dict_entries,
5459 budget.max_encoded_size,
5460 )
5461 })
5462 };
5463
5464 if let Some((indices_data_block, dictionary_data_block)) = dict_result {
5465 Self::encode_miniblock(
5466 column_idx,
5467 &field,
5468 compression_strategy.as_ref(),
5469 indices_data_block,
5470 repdef,
5471 row_number,
5472 Some(dictionary_data_block),
5473 num_rows,
5474 support_large_chunk,
5475 )
5476 } else if !too_sparse && Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
5477 log::debug!(
5478 "Encoding column {} with {} items using mini-block layout",
5479 column_idx,
5480 num_values
5481 );
5482 Self::encode_miniblock(
5483 column_idx,
5484 &field,
5485 compression_strategy.as_ref(),
5486 data_block,
5487 repdef,
5488 row_number,
5489 None,
5490 num_rows,
5491 support_large_chunk,
5492 )
5493 } else if too_sparse || Self::prefers_fullzip(encoding_metadata.as_ref()) {
5494 let user_forced_miniblock = encoding_metadata
5498 .get(STRUCTURAL_ENCODING_META_KEY)
5499 .map(|v| v.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK)
5500 .unwrap_or(false);
5501 if too_sparse && user_forced_miniblock {
5502 log::warn!(
5503 "Field '{}' requested {}={} but the rep/def shape would \
5504 overflow miniblock per-chunk limits; encoding as fullzip \
5505 instead to avoid corruption.",
5506 field.name,
5507 STRUCTURAL_ENCODING_META_KEY,
5508 STRUCTURAL_ENCODING_MINIBLOCK,
5509 );
5510 }
5511 log::debug!(
5512 "Encoding column {} with {} items using full-zip layout",
5513 column_idx,
5514 num_values
5515 );
5516 Self::encode_full_zip(
5517 column_idx,
5518 &field,
5519 compression_strategy.as_ref(),
5520 data_block,
5521 repdef,
5522 row_number,
5523 num_rows,
5524 )
5525 } else {
5526 Err(Error::invalid_input_source(format!("Cannot determine structural encoding for field {}. This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into()))
5527 }
5528 }
5529 })
5530 .boxed();
5531 Ok(vec![task])
5532 }
5533
5534 fn extract_validity_buf(
5535 array: Arc<dyn Array>,
5536 repdef: &mut RepDefBuilder,
5537 keep_original_array: bool,
5538 ) -> Result<Arc<dyn Array>> {
5539 if let Some(validity) = array.nulls() {
5540 if keep_original_array {
5541 repdef.add_validity_bitmap(validity.clone());
5542 } else {
5543 repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
5544 }
5545 let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
5546 Ok(make_array(data_no_nulls))
5547 } else {
5548 repdef.add_no_null(array.len());
5549 Ok(array)
5550 }
5551 }
5552
5553 fn extract_validity(
5554 mut array: Arc<dyn Array>,
5555 repdef: &mut RepDefBuilder,
5556 keep_original_array: bool,
5557 ) -> Result<Arc<dyn Array>> {
5558 match array.data_type() {
5559 DataType::Null => {
5560 repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
5561 Ok(array)
5562 }
5563 DataType::Dictionary(_, _) => {
5564 array = dict::normalize_dict_nulls(array)?;
5565 Self::extract_validity_buf(array, repdef, keep_original_array)
5566 }
5567 _ => Self::extract_validity_buf(array, repdef, keep_original_array),
5576 }
5577 }
5578}
5579
5580impl FieldEncoder for PrimitiveStructuralEncoder {
5581 fn maybe_encode(
5583 &mut self,
5584 array: ArrayRef,
5585 _external_buffers: &mut OutOfLineBuffers,
5586 mut repdef: RepDefBuilder,
5587 row_number: u64,
5588 num_rows: u64,
5589 ) -> Result<Vec<EncodeTask>> {
5590 let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
5591 self.accumulated_repdefs.push(repdef);
5592
5593 if let Some((arrays, row_number, num_rows)) =
5594 self.accumulation_queue.insert(array, row_number, num_rows)
5595 {
5596 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5597 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5598 } else {
5599 Ok(vec![])
5600 }
5601 }
5602
5603 fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
5605 if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
5606 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5607 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5608 } else {
5609 Ok(vec![])
5610 }
5611 }
5612
5613 fn num_columns(&self) -> u32 {
5614 1
5615 }
5616
5617 fn finish(
5618 &mut self,
5619 _external_buffers: &mut OutOfLineBuffers,
5620 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
5621 std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
5622 }
5623}
5624
5625#[cfg(test)]
5626#[allow(clippy::single_range_in_vec_init)]
5627mod tests {
5628 use super::{
5629 ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
5630 FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource,
5631 FullZipRepIndexDetails, FullZipScheduler, MiniBlockRepIndex, PerValueDecompressor,
5632 PreambleAction, StructuralPageScheduler, VariableFullZipDecoder,
5633 };
5634 use crate::buffer::LanceBuffer;
5635 use crate::compression::DefaultDecompressionStrategy;
5636 use crate::constants::{
5637 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
5638 DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
5639 STRUCTURAL_ENCODING_MINIBLOCK,
5640 };
5641 use crate::data::BlockInfo;
5642 use crate::decoder::PageEncoding;
5643 use crate::encodings::logical::primitive::{
5644 ChunkDrainInstructions, PrimitiveStructuralEncoder,
5645 };
5646 use crate::format::ProtobufUtils21;
5647 use crate::format::pb21;
5648 use crate::format::pb21::compressive_encoding::Compression;
5649 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
5650 use crate::version::LanceFileVersion;
5651 use arrow_array::{ArrayRef, Int8Array, StringArray};
5652 use arrow_schema::DataType;
5653 use std::collections::HashMap;
5654 use std::{collections::VecDeque, sync::Arc};
5655
5656 #[test]
5657 fn test_is_narrow() {
5658 let int8_array = Int8Array::from(vec![1, 2, 3]);
5659 let array_ref: ArrayRef = Arc::new(int8_array);
5660 let block = DataBlock::from_array(array_ref);
5661
5662 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5663
5664 let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
5665 let block = DataBlock::from_array(string_array);
5666 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5667
5668 let string_array = StringArray::from(vec![
5669 Some("hello world".repeat(100)),
5670 Some("world".to_string()),
5671 ]);
5672 let block = DataBlock::from_array(string_array);
5673 assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
5674 }
5675
5676 #[test]
5677 fn test_map_range() {
5678 let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
5681 let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
5682 let max_visible_def = 0;
5683 let total_items = 8;
5684 let max_rep = 1;
5685
5686 let check = |range, expected_item_range, expected_level_range| {
5687 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5688 range,
5689 rep.as_ref(),
5690 def.as_ref(),
5691 max_rep,
5692 max_visible_def,
5693 total_items,
5694 PreambleAction::Absent,
5695 );
5696 assert_eq!(item_range, expected_item_range);
5697 assert_eq!(level_range, expected_level_range);
5698 };
5699
5700 check(0..1, 0..3, 0..3);
5701 check(1..2, 3..5, 3..5);
5702 check(2..3, 5..5, 5..6);
5703 check(3..4, 5..8, 6..9);
5704 check(0..2, 0..5, 0..5);
5705 check(1..3, 3..5, 3..6);
5706 check(2..4, 5..8, 5..9);
5707 check(0..3, 0..5, 0..6);
5708 check(1..4, 3..8, 3..9);
5709 check(0..4, 0..8, 0..9);
5710
5711 let rep = Some(vec![1, 1, 0, 1]);
5714 let def = Some(vec![1, 0, 0, 0]);
5715 let max_visible_def = 0;
5716 let total_items = 3;
5717
5718 let check = |range, expected_item_range, expected_level_range| {
5719 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5720 range,
5721 rep.as_ref(),
5722 def.as_ref(),
5723 max_rep,
5724 max_visible_def,
5725 total_items,
5726 PreambleAction::Absent,
5727 );
5728 assert_eq!(item_range, expected_item_range);
5729 assert_eq!(level_range, expected_level_range);
5730 };
5731
5732 check(0..1, 0..0, 0..1);
5733 check(1..2, 0..2, 1..3);
5734 check(2..3, 2..3, 3..4);
5735 check(0..2, 0..2, 0..3);
5736 check(1..3, 0..3, 1..4);
5737 check(0..3, 0..3, 0..4);
5738
5739 let rep = Some(vec![1, 1, 0, 1]);
5742 let def = Some(vec![0, 0, 0, 1]);
5743 let max_visible_def = 0;
5744 let total_items = 3;
5745
5746 let check = |range, expected_item_range, expected_level_range| {
5747 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5748 range,
5749 rep.as_ref(),
5750 def.as_ref(),
5751 max_rep,
5752 max_visible_def,
5753 total_items,
5754 PreambleAction::Absent,
5755 );
5756 assert_eq!(item_range, expected_item_range);
5757 assert_eq!(level_range, expected_level_range);
5758 };
5759
5760 check(0..1, 0..1, 0..1);
5761 check(1..2, 1..3, 1..3);
5762 check(2..3, 3..3, 3..4);
5763 check(0..2, 0..3, 0..3);
5764 check(1..3, 1..3, 1..4);
5765 check(0..3, 0..3, 0..4);
5766
5767 let rep = Some(vec![1, 0, 1, 0, 1, 0]);
5770 let def: Option<&[u16]> = None;
5771 let max_visible_def = 0;
5772 let total_items = 6;
5773
5774 let check = |range, expected_item_range, expected_level_range| {
5775 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5776 range,
5777 rep.as_ref(),
5778 def.as_ref(),
5779 max_rep,
5780 max_visible_def,
5781 total_items,
5782 PreambleAction::Absent,
5783 );
5784 assert_eq!(item_range, expected_item_range);
5785 assert_eq!(level_range, expected_level_range);
5786 };
5787
5788 check(0..1, 0..2, 0..2);
5789 check(1..2, 2..4, 2..4);
5790 check(2..3, 4..6, 4..6);
5791 check(0..2, 0..4, 0..4);
5792 check(1..3, 2..6, 2..6);
5793 check(0..3, 0..6, 0..6);
5794
5795 let rep: Option<&[u16]> = None;
5798 let def = Some(vec![0, 0, 1, 0]);
5799 let max_visible_def = 1;
5800 let total_items = 4;
5801
5802 let check = |range, expected_item_range, expected_level_range| {
5803 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5804 range,
5805 rep.as_ref(),
5806 def.as_ref(),
5807 max_rep,
5808 max_visible_def,
5809 total_items,
5810 PreambleAction::Absent,
5811 );
5812 assert_eq!(item_range, expected_item_range);
5813 assert_eq!(level_range, expected_level_range);
5814 };
5815
5816 check(0..1, 0..1, 0..1);
5817 check(1..2, 1..2, 1..2);
5818 check(2..3, 2..3, 2..3);
5819 check(0..2, 0..2, 0..2);
5820 check(1..3, 1..3, 1..3);
5821 check(0..3, 0..3, 0..3);
5822
5823 let rep = Some(vec![0, 1, 0, 1]);
5828 let def = Some(vec![0, 0, 0, 1]);
5829 let max_visible_def = 0;
5830 let total_items = 3;
5831
5832 let check = |range, expected_item_range, expected_level_range| {
5833 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5834 range,
5835 rep.as_ref(),
5836 def.as_ref(),
5837 max_rep,
5838 max_visible_def,
5839 total_items,
5840 PreambleAction::Take,
5841 );
5842 assert_eq!(item_range, expected_item_range);
5843 assert_eq!(level_range, expected_level_range);
5844 };
5845
5846 check(0..1, 0..3, 0..3);
5848 check(0..2, 0..3, 0..4);
5849
5850 let check = |range, expected_item_range, expected_level_range| {
5851 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5852 range,
5853 rep.as_ref(),
5854 def.as_ref(),
5855 max_rep,
5856 max_visible_def,
5857 total_items,
5858 PreambleAction::Skip,
5859 );
5860 assert_eq!(item_range, expected_item_range);
5861 assert_eq!(level_range, expected_level_range);
5862 };
5863
5864 check(0..1, 1..3, 1..3);
5865 check(1..2, 3..3, 3..4);
5866 check(0..2, 1..3, 1..4);
5867
5868 let rep = Some(vec![0, 1, 1, 0]);
5873 let def = Some(vec![0, 1, 0, 0]);
5874 let max_visible_def = 0;
5875 let total_items = 4;
5876
5877 let check = |range, expected_item_range, expected_level_range| {
5878 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5879 range,
5880 rep.as_ref(),
5881 def.as_ref(),
5882 max_rep,
5883 max_visible_def,
5884 total_items,
5885 PreambleAction::Take,
5886 );
5887 assert_eq!(item_range, expected_item_range);
5888 assert_eq!(level_range, expected_level_range);
5889 };
5890
5891 check(0..1, 0..1, 0..2);
5893 check(0..2, 0..3, 0..4);
5894
5895 let check = |range, expected_item_range, expected_level_range| {
5896 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5897 range,
5898 rep.as_ref(),
5899 def.as_ref(),
5900 max_rep,
5901 max_visible_def,
5902 total_items,
5903 PreambleAction::Skip,
5904 );
5905 assert_eq!(item_range, expected_item_range);
5906 assert_eq!(level_range, expected_level_range);
5907 };
5908
5909 check(0..1, 1..1, 1..2);
5911 check(1..2, 1..3, 2..4);
5912 check(0..2, 1..3, 1..4);
5913
5914 let rep = Some(vec![0, 1, 0, 1]);
5917 let def: Option<Vec<u16>> = None;
5918 let max_visible_def = 0;
5919 let total_items = 4;
5920
5921 let check = |range, expected_item_range, expected_level_range| {
5922 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5923 range,
5924 rep.as_ref(),
5925 def.as_ref(),
5926 max_rep,
5927 max_visible_def,
5928 total_items,
5929 PreambleAction::Take,
5930 );
5931 assert_eq!(item_range, expected_item_range);
5932 assert_eq!(level_range, expected_level_range);
5933 };
5934
5935 check(0..1, 0..3, 0..3);
5937 check(0..2, 0..4, 0..4);
5938
5939 let check = |range, expected_item_range, expected_level_range| {
5940 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5941 range,
5942 rep.as_ref(),
5943 def.as_ref(),
5944 max_rep,
5945 max_visible_def,
5946 total_items,
5947 PreambleAction::Skip,
5948 );
5949 assert_eq!(item_range, expected_item_range);
5950 assert_eq!(level_range, expected_level_range);
5951 };
5952
5953 check(0..1, 1..3, 1..3);
5954 check(1..2, 3..4, 3..4);
5955 check(0..2, 1..4, 1..4);
5956
5957 let rep = Some(vec![2, 1, 2, 0, 1, 2]);
5961 let def = Some(vec![0, 1, 2, 0, 0, 0]);
5962 let max_rep = 2;
5963 let max_visible_def = 0;
5964 let total_items = 4;
5965
5966 let check = |range, expected_item_range, expected_level_range| {
5967 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5968 range,
5969 rep.as_ref(),
5970 def.as_ref(),
5971 max_rep,
5972 max_visible_def,
5973 total_items,
5974 PreambleAction::Absent,
5975 );
5976 assert_eq!(item_range, expected_item_range);
5977 assert_eq!(level_range, expected_level_range);
5978 };
5979
5980 check(0..3, 0..4, 0..6);
5981 check(0..1, 0..1, 0..2);
5982 check(1..2, 1..3, 2..5);
5983 check(2..3, 3..4, 5..6);
5984
5985 let rep = Some(vec![0, 0, 1, 0, 1, 1]);
5987 let def = Some(vec![0, 1, 0, 0, 0, 0]);
5988 let max_rep = 1;
5989 let max_visible_def = 0;
5990 let total_items = 5;
5991
5992 let check = |range, expected_item_range, expected_level_range| {
5993 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5994 range,
5995 rep.as_ref(),
5996 def.as_ref(),
5997 max_rep,
5998 max_visible_def,
5999 total_items,
6000 PreambleAction::Take,
6001 );
6002 assert_eq!(item_range, expected_item_range);
6003 assert_eq!(level_range, expected_level_range);
6004 };
6005
6006 check(0..0, 0..1, 0..2);
6007 check(0..1, 0..3, 0..4);
6008 check(0..2, 0..4, 0..5);
6009
6010 let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
6013 let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
6014 let max_rep = 1;
6015 let max_visible_def = 0;
6016 let total_items = 5;
6017
6018 let check = |range, expected_item_range, expected_level_range| {
6019 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6020 range,
6021 rep.as_ref(),
6022 def.as_ref(),
6023 max_rep,
6024 max_visible_def,
6025 total_items,
6026 PreambleAction::Skip,
6027 );
6028 assert_eq!(item_range, expected_item_range);
6029 assert_eq!(level_range, expected_level_range);
6030 };
6031
6032 check(2..3, 2..4, 5..7);
6033 }
6034
6035 #[test]
6036 fn test_slice_batch_data_and_rebase_offsets_u32() {
6037 let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
6038 let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
6039
6040 let (sliced_data, normalized_offsets) =
6041 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6042 .unwrap();
6043
6044 assert_eq!(sliced_data.as_ref(), b"6789ab");
6045 let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
6046 assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
6047 }
6048
6049 #[test]
6050 fn test_slice_batch_data_and_rebase_offsets_u64() {
6051 let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
6052 let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
6053
6054 let (sliced_data, normalized_offsets) =
6055 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
6056 .unwrap();
6057
6058 assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
6059 let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
6060 assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
6061 }
6062
6063 #[test]
6064 fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
6065 let data = LanceBuffer::copy_slice(b"abcd");
6066 let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
6067
6068 let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6069 .expect_err("offset end before start should error");
6070 assert!(err.to_string().contains("less than base"));
6071 }
6072
6073 #[test]
6074 fn test_schedule_instructions() {
6075 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6077 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6078 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6079
6080 let check = |user_ranges, expected_instructions| {
6081 let instructions =
6082 ChunkInstructions::schedule_instructions(&repetition_index, user_ranges);
6083 assert_eq!(instructions, expected_instructions);
6084 };
6085
6086 let expected_take_all = vec![
6088 ChunkInstructions {
6089 chunk_idx: 0,
6090 preamble: PreambleAction::Absent,
6091 rows_to_skip: 0,
6092 rows_to_take: 6,
6093 take_trailer: true,
6094 },
6095 ChunkInstructions {
6096 chunk_idx: 1,
6097 preamble: PreambleAction::Take,
6098 rows_to_skip: 0,
6099 rows_to_take: 2,
6100 take_trailer: false,
6101 },
6102 ChunkInstructions {
6103 chunk_idx: 2,
6104 preamble: PreambleAction::Absent,
6105 rows_to_skip: 0,
6106 rows_to_take: 5,
6107 take_trailer: true,
6108 },
6109 ChunkInstructions {
6110 chunk_idx: 3,
6111 preamble: PreambleAction::Take,
6112 rows_to_skip: 0,
6113 rows_to_take: 1,
6114 take_trailer: false,
6115 },
6116 ];
6117
6118 check(&[0..14], expected_take_all.clone());
6120
6121 check(
6123 &[
6124 0..1,
6125 1..2,
6126 2..3,
6127 3..4,
6128 4..5,
6129 5..6,
6130 6..7,
6131 7..8,
6132 8..9,
6133 9..10,
6134 10..11,
6135 11..12,
6136 12..13,
6137 13..14,
6138 ],
6139 expected_take_all,
6140 );
6141
6142 check(
6146 &[0..1, 3..4],
6147 vec![
6148 ChunkInstructions {
6149 chunk_idx: 0,
6150 preamble: PreambleAction::Absent,
6151 rows_to_skip: 0,
6152 rows_to_take: 1,
6153 take_trailer: false,
6154 },
6155 ChunkInstructions {
6156 chunk_idx: 0,
6157 preamble: PreambleAction::Absent,
6158 rows_to_skip: 3,
6159 rows_to_take: 1,
6160 take_trailer: false,
6161 },
6162 ],
6163 );
6164
6165 check(
6167 &[5..6],
6168 vec![
6169 ChunkInstructions {
6170 chunk_idx: 0,
6171 preamble: PreambleAction::Absent,
6172 rows_to_skip: 5,
6173 rows_to_take: 1,
6174 take_trailer: true,
6175 },
6176 ChunkInstructions {
6177 chunk_idx: 1,
6178 preamble: PreambleAction::Take,
6179 rows_to_skip: 0,
6180 rows_to_take: 0,
6181 take_trailer: false,
6182 },
6183 ],
6184 );
6185
6186 check(
6188 &[7..10],
6189 vec![
6190 ChunkInstructions {
6191 chunk_idx: 1,
6192 preamble: PreambleAction::Skip,
6193 rows_to_skip: 1,
6194 rows_to_take: 1,
6195 take_trailer: false,
6196 },
6197 ChunkInstructions {
6198 chunk_idx: 2,
6199 preamble: PreambleAction::Absent,
6200 rows_to_skip: 0,
6201 rows_to_take: 2,
6202 take_trailer: false,
6203 },
6204 ],
6205 );
6206 }
6207
6208 #[test]
6209 fn test_drain_instructions() {
6210 fn drain_from_instructions(
6211 instructions: &mut VecDeque<ChunkInstructions>,
6212 mut rows_desired: u64,
6213 need_preamble: &mut bool,
6214 skip_in_chunk: &mut u64,
6215 ) -> Vec<ChunkDrainInstructions> {
6216 let mut drain_instructions = Vec::with_capacity(instructions.len());
6218 while rows_desired > 0 || *need_preamble {
6219 let (next_instructions, consumed_chunk) = instructions
6220 .front()
6221 .unwrap()
6222 .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
6223 if consumed_chunk {
6224 instructions.pop_front();
6225 }
6226 drain_instructions.push(next_instructions);
6227 }
6228 drain_instructions
6229 }
6230
6231 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6233 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6234 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6235 let user_ranges = vec![1..7, 10..14];
6236
6237 let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6239
6240 let mut to_drain = VecDeque::from(scheduled.clone());
6241
6242 let mut need_preamble = false;
6245 let mut skip_in_chunk = 0;
6246
6247 let next_batch =
6248 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6249
6250 assert!(!need_preamble);
6251 assert_eq!(skip_in_chunk, 4);
6252 assert_eq!(
6253 next_batch,
6254 vec![ChunkDrainInstructions {
6255 chunk_instructions: scheduled[0].clone(),
6256 rows_to_take: 4,
6257 rows_to_skip: 0,
6258 preamble_action: PreambleAction::Absent,
6259 }]
6260 );
6261
6262 let next_batch =
6263 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6264
6265 assert!(!need_preamble);
6266 assert_eq!(skip_in_chunk, 2);
6267
6268 assert_eq!(
6269 next_batch,
6270 vec![
6271 ChunkDrainInstructions {
6272 chunk_instructions: scheduled[0].clone(),
6273 rows_to_take: 1,
6274 rows_to_skip: 4,
6275 preamble_action: PreambleAction::Absent,
6276 },
6277 ChunkDrainInstructions {
6278 chunk_instructions: scheduled[1].clone(),
6279 rows_to_take: 1,
6280 rows_to_skip: 0,
6281 preamble_action: PreambleAction::Take,
6282 },
6283 ChunkDrainInstructions {
6284 chunk_instructions: scheduled[2].clone(),
6285 rows_to_take: 2,
6286 rows_to_skip: 0,
6287 preamble_action: PreambleAction::Absent,
6288 }
6289 ]
6290 );
6291
6292 let next_batch =
6293 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6294
6295 assert!(!need_preamble);
6296 assert_eq!(skip_in_chunk, 0);
6297
6298 assert_eq!(
6299 next_batch,
6300 vec![
6301 ChunkDrainInstructions {
6302 chunk_instructions: scheduled[2].clone(),
6303 rows_to_take: 1,
6304 rows_to_skip: 2,
6305 preamble_action: PreambleAction::Absent,
6306 },
6307 ChunkDrainInstructions {
6308 chunk_instructions: scheduled[3].clone(),
6309 rows_to_take: 1,
6310 rows_to_skip: 0,
6311 preamble_action: PreambleAction::Take,
6312 },
6313 ]
6314 );
6315
6316 let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
6318 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6319 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6320 let user_ranges = vec![0..28];
6321
6322 let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6324
6325 let mut to_drain = VecDeque::from(scheduled.clone());
6326
6327 let mut need_preamble = false;
6330 let mut skip_in_chunk = 0;
6331
6332 let next_batch =
6333 drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
6334
6335 assert_eq!(
6336 next_batch,
6337 vec![
6338 ChunkDrainInstructions {
6339 chunk_instructions: scheduled[0].clone(),
6340 rows_to_take: 6,
6341 rows_to_skip: 0,
6342 preamble_action: PreambleAction::Absent,
6343 },
6344 ChunkDrainInstructions {
6345 chunk_instructions: scheduled[1].clone(),
6346 rows_to_take: 1,
6347 rows_to_skip: 0,
6348 preamble_action: PreambleAction::Take,
6349 },
6350 ]
6351 );
6352
6353 assert!(!need_preamble);
6354 assert_eq!(skip_in_chunk, 1);
6355
6356 let next_batch =
6359 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6360
6361 assert_eq!(
6362 next_batch,
6363 vec![
6364 ChunkDrainInstructions {
6365 chunk_instructions: scheduled[1].clone(),
6366 rows_to_take: 2,
6367 rows_to_skip: 1,
6368 preamble_action: PreambleAction::Skip,
6369 },
6370 ChunkDrainInstructions {
6371 chunk_instructions: scheduled[2].clone(),
6372 rows_to_take: 0,
6373 rows_to_skip: 0,
6374 preamble_action: PreambleAction::Take,
6375 },
6376 ]
6377 );
6378
6379 assert!(!need_preamble);
6380 assert_eq!(skip_in_chunk, 0);
6381 }
6382
6383 #[tokio::test]
6384 async fn test_fullzip_initialize_is_lazy() {
6385 use futures::{FutureExt, future::BoxFuture};
6386 use std::ops::Range;
6387 use std::sync::Mutex;
6388
6389 #[derive(Debug, Clone)]
6390 struct RecordingScheduler {
6391 data: bytes::Bytes,
6392 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6393 }
6394
6395 impl RecordingScheduler {
6396 fn new(data: bytes::Bytes) -> Self {
6397 Self {
6398 data,
6399 requests: Arc::new(Mutex::new(Vec::new())),
6400 }
6401 }
6402
6403 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6404 self.requests.lock().unwrap().clone()
6405 }
6406 }
6407
6408 impl crate::EncodingsIo for RecordingScheduler {
6409 fn submit_request(
6410 &self,
6411 ranges: Vec<Range<u64>>,
6412 _priority: u64,
6413 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6414 self.requests.lock().unwrap().push(ranges.clone());
6415 let data = ranges
6416 .into_iter()
6417 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6418 .collect::<Vec<_>>();
6419 std::future::ready(Ok(data)).boxed()
6420 }
6421 }
6422
6423 #[derive(Debug)]
6424 struct TestFixedDecompressor;
6425
6426 impl FixedPerValueDecompressor for TestFixedDecompressor {
6427 fn decompress(
6428 &self,
6429 _data: FixedWidthDataBlock,
6430 _num_rows: u64,
6431 ) -> crate::Result<DataBlock> {
6432 unimplemented!("Test decompressor")
6433 }
6434
6435 fn bits_per_value(&self) -> u64 {
6436 32
6437 }
6438 }
6439
6440 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
6441 0;
6442 16 * 1024
6443 ])));
6444 let mut scheduler = FullZipScheduler {
6445 data_buf_position: 0,
6446 data_buf_size: 4096,
6447 rep_index: Some(FullZipRepIndexDetails {
6448 buf_position: 1000,
6449 bytes_per_value: 4,
6450 }),
6451 priority: 0,
6452 rows_in_page: 100,
6453 bits_per_offset: 32,
6454 details: Arc::new(FullZipDecodeDetails {
6455 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6456 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6457 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6458 max_rep: 0,
6459 max_visible_def: 0,
6460 }),
6461 cached_state: None,
6462 enable_cache: false,
6463 };
6464
6465 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6466 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6467
6468 assert!(
6469 cached_data
6470 .as_arc_any()
6471 .downcast_ref::<super::NoCachedPageData>()
6472 .is_some(),
6473 "FullZip initialize should not eagerly load repetition index data"
6474 );
6475 assert!(scheduler.cached_state.is_none());
6476 assert!(
6477 io.requests().is_empty(),
6478 "FullZip initialize should not issue any I/O"
6479 );
6480 }
6481
6482 #[tokio::test]
6483 async fn test_fullzip_read_source_slices_prefetched_page() {
6484 let page_start = 200_u64;
6485 let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
6486 let source = FullZipReadSource::PrefetchedPage {
6487 base_offset: page_start,
6488 data: page_data,
6489 };
6490 let ranges = vec![
6491 page_start..(page_start + 3),
6492 (page_start + 4)..(page_start + 8),
6493 ];
6494 let mut data = source.fetch(&ranges, 0).await.unwrap();
6495 assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
6496 assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
6497 }
6498
6499 #[tokio::test]
6500 async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
6501 use futures::{FutureExt, future::BoxFuture};
6502 use std::ops::Range;
6503 use std::sync::Mutex;
6504
6505 #[derive(Debug, Clone)]
6506 struct RecordingScheduler {
6507 data: bytes::Bytes,
6508 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6509 }
6510
6511 impl RecordingScheduler {
6512 fn new(data: bytes::Bytes) -> Self {
6513 Self {
6514 data,
6515 requests: Arc::new(Mutex::new(Vec::new())),
6516 }
6517 }
6518
6519 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6520 self.requests.lock().unwrap().clone()
6521 }
6522 }
6523
6524 impl crate::EncodingsIo for RecordingScheduler {
6525 fn submit_request(
6526 &self,
6527 ranges: Vec<Range<u64>>,
6528 _priority: u64,
6529 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6530 self.requests.lock().unwrap().push(ranges.clone());
6531 let data = ranges
6532 .into_iter()
6533 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6534 .collect::<Vec<_>>();
6535 std::future::ready(Ok(data)).boxed()
6536 }
6537 }
6538
6539 #[derive(Debug)]
6540 struct TestFixedDecompressor;
6541
6542 impl FixedPerValueDecompressor for TestFixedDecompressor {
6543 fn decompress(
6544 &self,
6545 _data: FixedWidthDataBlock,
6546 _num_rows: u64,
6547 ) -> crate::Result<DataBlock> {
6548 unimplemented!("Test decompressor")
6549 }
6550
6551 fn bits_per_value(&self) -> u64 {
6552 32
6553 }
6554 }
6555
6556 let rows_in_page = 100_u64;
6557 let bytes_per_value = 4_u64;
6558 let rep_start = 1000_u64;
6559 let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
6560 let mut data = vec![0_u8; 16 * 1024];
6561 data[rep_start as usize..rep_start as usize + rep_size].fill(7);
6562 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
6563
6564 let mut scheduler = FullZipScheduler {
6565 data_buf_position: 0,
6566 data_buf_size: 4096,
6567 rep_index: Some(FullZipRepIndexDetails {
6568 buf_position: rep_start,
6569 bytes_per_value,
6570 }),
6571 priority: 0,
6572 rows_in_page,
6573 bits_per_offset: 32,
6574 details: Arc::new(FullZipDecodeDetails {
6575 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6576 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6577 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6578 max_rep: 0,
6579 max_visible_def: 0,
6580 }),
6581 cached_state: None,
6582 enable_cache: true,
6583 };
6584
6585 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6586 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6587 assert!(
6588 cached_data
6589 .as_arc_any()
6590 .downcast_ref::<FullZipCacheableState>()
6591 .is_some()
6592 );
6593 assert!(scheduler.cached_state.is_some());
6594 assert_eq!(
6595 io.requests(),
6596 vec![vec![
6597 rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
6598 ]]
6599 );
6600 }
6601
6602 #[tokio::test]
6603 async fn test_fullzip_full_page_bypasses_rep_index_io() {
6604 use futures::{FutureExt, future::BoxFuture};
6605 use std::ops::Range;
6606 use std::sync::Mutex;
6607
6608 #[derive(Debug, Clone)]
6609 struct RecordingScheduler {
6610 data: bytes::Bytes,
6611 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6612 }
6613
6614 impl RecordingScheduler {
6615 fn new(data: bytes::Bytes) -> Self {
6616 Self {
6617 data,
6618 requests: Arc::new(Mutex::new(Vec::new())),
6619 }
6620 }
6621
6622 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6623 self.requests.lock().unwrap().clone()
6624 }
6625 }
6626
6627 impl crate::EncodingsIo for RecordingScheduler {
6628 fn submit_request(
6629 &self,
6630 ranges: Vec<Range<u64>>,
6631 _priority: u64,
6632 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6633 self.requests.lock().unwrap().push(ranges.clone());
6634 let data = ranges
6635 .into_iter()
6636 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6637 .collect::<Vec<_>>();
6638 std::future::ready(Ok(data)).boxed()
6639 }
6640 }
6641
6642 #[derive(Debug)]
6643 struct TestFixedDecompressor;
6644
6645 impl FixedPerValueDecompressor for TestFixedDecompressor {
6646 fn decompress(
6647 &self,
6648 _data: FixedWidthDataBlock,
6649 _num_rows: u64,
6650 ) -> crate::Result<DataBlock> {
6651 unimplemented!("Test decompressor")
6652 }
6653
6654 fn bits_per_value(&self) -> u64 {
6655 32
6656 }
6657 }
6658
6659 let rows_in_page = 100_u64;
6660 let data_start = 256_u64;
6661 let data_size = 500_u64;
6662 let rep_start = 4096_u64;
6663 let bytes_per_value = 4_u64;
6664
6665 let mut bytes = vec![0_u8; 16 * 1024];
6666 for i in 0..=rows_in_page {
6667 let offset = (i * 5) as u32;
6668 let pos = rep_start as usize + (i * bytes_per_value) as usize;
6669 bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
6670 }
6671 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
6672
6673 let scheduler = FullZipScheduler {
6674 data_buf_position: data_start,
6675 data_buf_size: data_size,
6676 rep_index: Some(FullZipRepIndexDetails {
6677 buf_position: rep_start,
6678 bytes_per_value,
6679 }),
6680 priority: 0,
6681 rows_in_page,
6682 bits_per_offset: 32,
6683 details: Arc::new(FullZipDecodeDetails {
6684 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6685 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6686 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6687 max_rep: 0,
6688 max_visible_def: 0,
6689 }),
6690 cached_state: None,
6691 enable_cache: false,
6692 };
6693
6694 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6695 let tasks = scheduler
6696 .schedule_ranges_rep(
6697 &[0..rows_in_page],
6698 &io_dyn,
6699 FullZipRepIndexDetails {
6700 buf_position: rep_start,
6701 bytes_per_value,
6702 },
6703 )
6704 .unwrap();
6705
6706 let requests = io.requests();
6707 assert_eq!(requests.len(), 1);
6708 assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
6709
6710 let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
6711 let requests_after_await = io.requests();
6712 assert_eq!(
6713 requests_after_await.len(),
6714 1,
6715 "full page path should not issue rep-index I/O"
6716 );
6717 }
6718
6719 #[tokio::test]
6721 async fn test_fuzz_issue_4492_empty_rep_values() {
6722 use lance_datagen::{RowCount, Seed, array, gen_batch};
6723
6724 let seed = 1823859942947654717u64;
6725 let num_rows = 2741usize;
6726
6727 let batch_gen = gen_batch().with_seed(Seed::from(seed));
6729 let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
6730 let list_generator = array::rand_list_any(base_generator, false);
6731
6732 let batch = batch_gen
6733 .anon_col(list_generator)
6734 .into_batch_rows(RowCount::from(num_rows as u64))
6735 .unwrap();
6736
6737 let list_array = batch.column(0).clone();
6738
6739 let mut metadata = HashMap::new();
6741 metadata.insert(
6742 STRUCTURAL_ENCODING_META_KEY.to_string(),
6743 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6744 );
6745
6746 let test_cases = TestCases::default()
6747 .with_min_file_version(LanceFileVersion::V2_1)
6748 .with_batch_size(100)
6749 .with_range(0..num_rows.min(500) as u64)
6750 .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
6751
6752 check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
6753 }
6754
6755 async fn test_minichunk_size_helper(
6756 string_data: Vec<Option<String>>,
6757 minichunk_size: u64,
6758 file_version: LanceFileVersion,
6759 ) {
6760 use crate::constants::MINICHUNK_SIZE_META_KEY;
6761 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6762 use arrow_array::{ArrayRef, StringArray};
6763 use std::sync::Arc;
6764
6765 let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
6766
6767 let mut metadata = HashMap::new();
6768 metadata.insert(
6769 MINICHUNK_SIZE_META_KEY.to_string(),
6770 minichunk_size.to_string(),
6771 );
6772 metadata.insert(
6773 STRUCTURAL_ENCODING_META_KEY.to_string(),
6774 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6775 );
6776
6777 let test_cases = TestCases::default()
6778 .with_min_file_version(file_version)
6779 .with_batch_size(1000);
6780
6781 check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
6782 }
6783
6784 #[tokio::test]
6785 async fn test_minichunk_size_roundtrip() {
6786 let mut string_data = Vec::new();
6788 for i in 0..100 {
6789 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6790 }
6791 test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await;
6793 }
6794
6795 #[tokio::test]
6796 async fn test_minichunk_size_128kb_v2_2() {
6797 let mut string_data = Vec::new();
6799 for i in 0..10000 {
6801 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6802 }
6803 test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6804 }
6805
6806 #[tokio::test]
6807 async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
6808 let mut string_data = Vec::new();
6809 for i in 0..10000 {
6811 string_data.push(Some(format!("t_{}", i)));
6812 }
6813 test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6814 }
6815
6816 #[tokio::test]
6817 async fn test_large_dictionary_general_compression() {
6818 use arrow_array::{ArrayRef, StringArray};
6819 use std::collections::HashMap;
6820 use std::sync::Arc;
6821
6822 let unique_values: Vec<String> = (0..100)
6825 .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
6826 .collect();
6827
6828 let repeated_strings: Vec<_> = unique_values
6830 .iter()
6831 .cycle()
6832 .take(100_000)
6833 .map(|s| Some(s.as_str()))
6834 .collect();
6835
6836 let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
6837
6838 let test_cases = TestCases::default()
6840 .with_min_file_version(LanceFileVersion::V2_2)
6841 .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
6842 assert_eq!(cols.len(), 1);
6843 let col = &cols[0];
6844
6845 if let Some(PageEncoding::Structural(page_layout)) =
6847 &col.final_pages.first().map(|p| &p.description)
6848 && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
6849 &page_layout.layout
6850 && let Some(dictionary_encoding) = &mini_block.dictionary
6851 {
6852 match dictionary_encoding.compression.as_ref() {
6853 Some(Compression::General(general)) => {
6854 let compression = general.compression.as_ref().unwrap();
6856 assert!(
6857 compression.scheme()
6858 == pb21::CompressionScheme::CompressionAlgorithmLz4
6859 || compression.scheme()
6860 == pb21::CompressionScheme::CompressionAlgorithmZstd,
6861 "Expected LZ4 or Zstd compression for large dictionary"
6862 );
6863 }
6864 _ => panic!("Expected General compression for large dictionary"),
6865 }
6866 }
6867 }));
6868
6869 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
6870 }
6871
6872 fn dictionary_encoding_from_page(
6873 page: &crate::encoder::EncodedPage,
6874 ) -> &crate::format::pb21::CompressiveEncoding {
6875 let PageEncoding::Structural(layout) = &page.description else {
6876 panic!("Expected structural page encoding");
6877 };
6878 let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
6879 else {
6880 panic!("Expected mini-block layout");
6881 };
6882 layout
6883 .dictionary
6884 .as_ref()
6885 .unwrap_or_else(|| panic!("Expected dictionary encoding"))
6886 }
6887
6888 async fn encode_variable_dict_page(
6889 metadata: HashMap<String, String>,
6890 ) -> crate::encoder::EncodedPage {
6891 use arrow_array::types::Int32Type;
6892 use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
6893
6894 let values = Arc::new(StringArray::from(
6895 (0..128)
6896 .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
6897 .collect::<Vec<_>>(),
6898 )) as ArrayRef;
6899 let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
6900 let dict_array =
6901 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
6902
6903 let field = arrow_schema::Field::new(
6904 "dict_col",
6905 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6906 false,
6907 )
6908 .with_metadata(metadata);
6909
6910 encode_first_page(field, dict_array, LanceFileVersion::V2_2).await
6911 }
6912
6913 async fn encode_auto_fixed_dict_page(
6914 metadata: HashMap<String, String>,
6915 ) -> crate::encoder::EncodedPage {
6916 use arrow_array::{ArrayRef, Decimal128Array};
6917
6918 let values = (0..20_000)
6920 .map(|i| match i % 3 {
6921 0 => 10_i128,
6922 1 => 20_i128,
6923 _ => 30_i128,
6924 })
6925 .collect::<Vec<_>>();
6926 let decimal = Decimal128Array::from_iter_values(values)
6927 .with_precision_and_scale(38, 0)
6928 .unwrap();
6929 let decimal = Arc::new(decimal) as ArrayRef;
6930
6931 let mut field_metadata = metadata;
6932 field_metadata.insert(
6934 "lance-encoding:dict-size-ratio".to_string(),
6935 "0.99".to_string(),
6936 );
6937 let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
6938 .with_metadata(field_metadata);
6939
6940 encode_first_page(field, decimal, LanceFileVersion::V2_2).await
6941 }
6942
6943 #[tokio::test]
6944 async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
6945 let page = encode_variable_dict_page(HashMap::new()).await;
6946 let dictionary_encoding = dictionary_encoding_from_page(&page);
6947 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6948 panic!("Expected General compression for dictionary values");
6949 };
6950 let compression = general.compression.as_ref().unwrap();
6951 assert_eq!(
6952 compression.scheme(),
6953 pb21::CompressionScheme::CompressionAlgorithmLz4
6954 );
6955 }
6956
6957 #[tokio::test]
6958 async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
6959 let page = encode_auto_fixed_dict_page(HashMap::new()).await;
6960 let dictionary_encoding = dictionary_encoding_from_page(&page);
6961 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6962 panic!("Expected General compression for dictionary values");
6963 };
6964 let compression = general.compression.as_ref().unwrap();
6965 assert_eq!(
6966 compression.scheme(),
6967 pb21::CompressionScheme::CompressionAlgorithmLz4
6968 );
6969 }
6970
6971 #[tokio::test]
6972 async fn test_dict_values_general_compression_zstd() {
6973 let mut metadata = HashMap::new();
6974 metadata.insert(
6975 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
6976 "zstd".to_string(),
6977 );
6978 let page = encode_variable_dict_page(metadata).await;
6979 let dictionary_encoding = dictionary_encoding_from_page(&page);
6980 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6981 panic!("Expected General compression for dictionary values");
6982 };
6983 let compression = general.compression.as_ref().unwrap();
6984 assert_eq!(
6985 compression.scheme(),
6986 pb21::CompressionScheme::CompressionAlgorithmZstd
6987 );
6988 }
6989
6990 #[tokio::test]
6991 async fn test_dict_values_general_compression_none() {
6992 let mut metadata = HashMap::new();
6993 metadata.insert(
6994 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
6995 "none".to_string(),
6996 );
6997 let page = encode_variable_dict_page(metadata).await;
6998 let dictionary_encoding = dictionary_encoding_from_page(&page);
6999 assert!(
7000 !matches!(
7001 dictionary_encoding.compression.as_ref(),
7002 Some(Compression::General(_))
7003 ),
7004 "Expected dictionary values to avoid General compression"
7005 );
7006 }
7007
7008 #[test]
7009 fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
7010 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7011 &HashMap::new(),
7012 None,
7013 None,
7014 );
7015 assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
7016 assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
7017 }
7018
7019 #[test]
7020 fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
7021 let field_metadata = HashMap::from([
7022 (
7023 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7024 "none".to_string(),
7025 ),
7026 (
7027 DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
7028 "7".to_string(),
7029 ),
7030 ]);
7031 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7032 &field_metadata,
7033 Some("zstd".to_string()),
7034 Some("3".to_string()),
7035 );
7036 assert_eq!(
7037 metadata.get(COMPRESSION_META_KEY),
7038 Some(&"none".to_string()),
7039 );
7040 assert_eq!(
7041 metadata.get(COMPRESSION_LEVEL_META_KEY),
7042 Some(&"7".to_string()),
7043 );
7044 }
7045
7046 #[test]
7047 fn test_resolve_dict_values_compression_metadata_env_fallback() {
7048 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7049 &HashMap::new(),
7050 Some("zstd".to_string()),
7051 Some("9".to_string()),
7052 );
7053 assert_eq!(
7054 metadata.get(COMPRESSION_META_KEY),
7055 Some(&"zstd".to_string()),
7056 );
7057 assert_eq!(
7058 metadata.get(COMPRESSION_LEVEL_META_KEY),
7059 Some(&"9".to_string()),
7060 );
7061 }
7062
7063 #[tokio::test]
7064 async fn test_dictionary_encode_int64() {
7065 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7066 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7067 use crate::version::LanceFileVersion;
7068 use arrow_array::{ArrayRef, Int64Array};
7069 use std::collections::HashMap;
7070 use std::sync::Arc;
7071
7072 let values = (0..1000)
7074 .map(|i| match i % 3 {
7075 0 => 10i64,
7076 1 => 20i64,
7077 _ => 30i64,
7078 })
7079 .collect::<Vec<_>>();
7080 let array = Arc::new(Int64Array::from(values)) as ArrayRef;
7081
7082 let mut metadata = HashMap::new();
7083 metadata.insert(
7084 STRUCTURAL_ENCODING_META_KEY.to_string(),
7085 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7086 );
7087 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7088
7089 let test_cases = TestCases::default()
7090 .with_min_file_version(LanceFileVersion::V2_2)
7091 .with_batch_size(1000)
7092 .with_range(0..1000)
7093 .with_indices(vec![0, 1, 10, 999])
7094 .with_expected_encoding("dictionary");
7095
7096 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7097 }
7098
7099 #[tokio::test]
7100 async fn test_dictionary_encode_float64() {
7101 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7102 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7103 use crate::version::LanceFileVersion;
7104 use arrow_array::{ArrayRef, Float64Array};
7105 use std::collections::HashMap;
7106 use std::sync::Arc;
7107
7108 let values = (0..1000)
7110 .map(|i| match i % 3 {
7111 0 => 0.1f64,
7112 1 => 0.2f64,
7113 _ => 0.3f64,
7114 })
7115 .collect::<Vec<_>>();
7116 let array = Arc::new(Float64Array::from(values)) as ArrayRef;
7117
7118 let mut metadata = HashMap::new();
7119 metadata.insert(
7120 STRUCTURAL_ENCODING_META_KEY.to_string(),
7121 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7122 );
7123 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7124
7125 let test_cases = TestCases::default()
7126 .with_min_file_version(LanceFileVersion::V2_2)
7127 .with_batch_size(1000)
7128 .with_range(0..1000)
7129 .with_indices(vec![0, 1, 10, 999])
7130 .with_expected_encoding("dictionary");
7131
7132 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7133 }
7134
7135 #[test]
7136 fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
7137 let rows = 10_000;
7138 let unique_values = 2_000;
7139
7140 let dictionary_encoding =
7141 ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
7142 let layout = pb21::MiniBlockLayout {
7143 rep_compression: None,
7144 def_compression: None,
7145 value_compression: Some(ProtobufUtils21::flat(64, None)),
7146 dictionary: Some(dictionary_encoding),
7147 num_dictionary_items: unique_values,
7148 layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
7149 num_buffers: 1,
7150 repetition_index_depth: 0,
7151 num_items: rows,
7152 has_large_chunk: false,
7153 };
7154
7155 let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
7156 let scheduler = super::MiniBlockScheduler::try_new(
7157 &buffer_offsets_and_sizes,
7158 0,
7159 rows,
7160 &layout,
7161 &DefaultDecompressionStrategy::default(),
7162 )
7163 .unwrap();
7164
7165 let dictionary = scheduler.dictionary.unwrap();
7166 assert_eq!(dictionary.num_dictionary_items, unique_values);
7167 assert_eq!(
7168 dictionary.dictionary_data_alignment,
7169 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
7170 );
7171 }
7172
7173 fn create_test_fixed_data_block(
7175 num_values: u64,
7176 cardinality: u64,
7177 bits_per_value: u64,
7178 ) -> DataBlock {
7179 assert!(cardinality > 0);
7180 assert!(cardinality <= num_values);
7181 let block_info = BlockInfo::default();
7182
7183 assert_eq!(bits_per_value % 8, 0);
7184 let data = match bits_per_value {
7185 32 => {
7186 let values = (0..num_values)
7187 .map(|i| (i % cardinality) as u32)
7188 .collect::<Vec<_>>();
7189 crate::buffer::LanceBuffer::reinterpret_vec(values)
7190 }
7191 64 => {
7192 let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
7193 crate::buffer::LanceBuffer::reinterpret_vec(values)
7194 }
7195 128 => {
7196 let values = (0..num_values)
7197 .map(|i| (i % cardinality) as u128)
7198 .collect::<Vec<_>>();
7199 crate::buffer::LanceBuffer::reinterpret_vec(values)
7200 }
7201 _ => unreachable!(),
7202 };
7203 DataBlock::FixedWidth(FixedWidthDataBlock {
7204 bits_per_value,
7205 data,
7206 num_values,
7207 block_info,
7208 })
7209 }
7210
7211 fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7213 use arrow_array::StringArray;
7214
7215 assert!(cardinality <= num_values && cardinality > 0);
7216
7217 let mut values = Vec::with_capacity(num_values as usize);
7218 for i in 0..num_values {
7219 values.push(format!("value_{:016}", i % cardinality));
7220 }
7221
7222 let array = StringArray::from(values);
7223 DataBlock::from_array(Arc::new(array) as ArrayRef)
7224 }
7225
7226 #[test]
7227 fn test_should_dictionary_encode() {
7228 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7229 use lance_core::datatypes::Field as LanceField;
7230
7231 let block = create_test_variable_width_block(1000, 10);
7233
7234 let mut metadata = HashMap::new();
7235 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7236 let arrow_field =
7237 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7238 let field = LanceField::try_from(&arrow_field).unwrap();
7239
7240 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7241 &block,
7242 &field,
7243 LanceFileVersion::V2_1,
7244 );
7245
7246 assert!(
7247 result.is_some(),
7248 "Should use dictionary encode based on size"
7249 );
7250 }
7251
7252 #[test]
7253 fn test_should_not_dictionary_encode_unsupported_bits() {
7254 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7255 use lance_core::datatypes::Field as LanceField;
7256
7257 let block = create_test_fixed_data_block(1000, 1000, 32);
7258
7259 let mut metadata = HashMap::new();
7260 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7261 let arrow_field =
7262 arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
7263 let field = LanceField::try_from(&arrow_field).unwrap();
7264
7265 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7266 &block,
7267 &field,
7268 LanceFileVersion::V2_1,
7269 );
7270
7271 assert!(
7272 result.is_none(),
7273 "Should not use dictionary encode for unsupported bit width"
7274 );
7275 }
7276
7277 #[test]
7278 fn test_should_not_dictionary_encode_near_unique_sample() {
7279 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7280 use lance_core::datatypes::Field as LanceField;
7281
7282 let num_values = 5000;
7283 let block = create_test_variable_width_block(num_values, num_values);
7284
7285 let mut metadata = HashMap::new();
7286 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
7287 let arrow_field =
7288 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7289 let field = LanceField::try_from(&arrow_field).unwrap();
7290
7291 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7292 &block,
7293 &field,
7294 LanceFileVersion::V2_1,
7295 );
7296
7297 assert!(
7298 result.is_none(),
7299 "Should not probe dictionary encoding for near-unique data"
7300 );
7301 }
7302
7303 async fn encode_first_page(
7304 field: arrow_schema::Field,
7305 array: ArrayRef,
7306 version: LanceFileVersion,
7307 ) -> crate::encoder::EncodedPage {
7308 use crate::encoder::{
7309 ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
7310 default_encoding_strategy,
7311 };
7312 use crate::repdef::RepDefBuilder;
7313
7314 let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
7315 let encoding_strategy = default_encoding_strategy(version);
7316 let mut column_index_seq = ColumnIndexSequence::default();
7317 let encoding_options = EncodingOptions {
7318 cache_bytes_per_column: 1,
7319 max_page_bytes: 32 * 1024 * 1024,
7320 keep_original_array: true,
7321 buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
7322 version,
7323 };
7324
7325 let mut encoder = encoding_strategy
7326 .create_field_encoder(
7327 encoding_strategy.as_ref(),
7328 &lance_field,
7329 &mut column_index_seq,
7330 &encoding_options,
7331 )
7332 .unwrap();
7333
7334 let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
7335 let repdef = RepDefBuilder::default();
7336 let num_rows = array.len() as u64;
7337 let mut pages = Vec::new();
7338 for task in encoder
7339 .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
7340 .unwrap()
7341 {
7342 pages.push(task.await.unwrap());
7343 }
7344 for task in encoder.flush(&mut external_buffers).unwrap() {
7345 pages.push(task.await.unwrap());
7346 }
7347 pages.into_iter().next().unwrap()
7348 }
7349
7350 #[tokio::test]
7351 async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
7352 use crate::format::pb21::page_layout::Layout;
7353
7354 let val = vec![0xABu8; 33];
7355 let arr: ArrayRef = Arc::new(
7356 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7357 std::iter::repeat_n(Some(val.as_slice()), 256),
7358 33,
7359 )
7360 .unwrap(),
7361 );
7362 let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
7363 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7364
7365 let PageEncoding::Structural(layout) = &page.description else {
7366 panic!("Expected structural encoding");
7367 };
7368 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7369 panic!("Expected constant layout in slot 2");
7370 };
7371 assert!(layout.inline_value.is_none());
7372 assert_eq!(page.data.len(), 1);
7373
7374 let test_cases = TestCases::default()
7375 .with_min_file_version(LanceFileVersion::V2_2)
7376 .with_max_file_version(LanceFileVersion::V2_2)
7377 .with_page_sizes(vec![4096]);
7378 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7379 }
7380
7381 #[tokio::test]
7382 async fn test_constant_layout_out_of_line_utf8_v2_2() {
7383 use crate::format::pb21::page_layout::Layout;
7384
7385 let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
7386 std::iter::repeat_n("hello", 512),
7387 ));
7388 let field = arrow_schema::Field::new("c", DataType::Utf8, true);
7389 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7390
7391 let PageEncoding::Structural(layout) = &page.description else {
7392 panic!("Expected structural encoding");
7393 };
7394 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7395 panic!("Expected constant layout in slot 2");
7396 };
7397 assert!(layout.inline_value.is_none());
7398 assert_eq!(page.data.len(), 1);
7399
7400 let test_cases = TestCases::default()
7401 .with_min_file_version(LanceFileVersion::V2_2)
7402 .with_max_file_version(LanceFileVersion::V2_2)
7403 .with_page_sizes(vec![4096]);
7404 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7405 }
7406
7407 #[tokio::test]
7408 async fn test_constant_layout_nullable_item_v2_2() {
7409 use crate::format::pb21::page_layout::Layout;
7410
7411 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
7412 Some(7),
7413 None,
7414 Some(7),
7415 None,
7416 Some(7),
7417 ]));
7418 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7419 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7420
7421 let PageEncoding::Structural(layout) = &page.description else {
7422 panic!("Expected structural encoding");
7423 };
7424 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7425 panic!("Expected constant layout in slot 2");
7426 };
7427 assert!(layout.inline_value.is_some());
7428 assert_eq!(page.data.len(), 2);
7429
7430 let test_cases = TestCases::default()
7431 .with_min_file_version(LanceFileVersion::V2_2)
7432 .with_max_file_version(LanceFileVersion::V2_2)
7433 .with_page_sizes(vec![4096]);
7434 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7435 }
7436
7437 #[tokio::test]
7438 async fn test_constant_layout_list_repdef_v2_2() {
7439 use crate::format::pb21::page_layout::Layout;
7440 use arrow_array::builder::{Int32Builder, ListBuilder};
7441
7442 let mut builder = ListBuilder::new(Int32Builder::new());
7443 builder.values().append_value(7);
7444 builder.values().append_null();
7445 builder.values().append_value(7);
7446 builder.append(true);
7447
7448 builder.append(true);
7449
7450 builder.values().append_value(7);
7451 builder.append(true);
7452
7453 builder.append_null();
7454
7455 let arr: ArrayRef = Arc::new(builder.finish());
7456 let field = arrow_schema::Field::new(
7457 "c",
7458 DataType::List(Arc::new(arrow_schema::Field::new(
7459 "item",
7460 DataType::Int32,
7461 true,
7462 ))),
7463 true,
7464 );
7465 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7466
7467 let PageEncoding::Structural(layout) = &page.description else {
7468 panic!("Expected structural encoding");
7469 };
7470 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7471 panic!("Expected constant layout in slot 2");
7472 };
7473 assert!(layout.inline_value.is_some());
7474 assert_eq!(page.data.len(), 2);
7475
7476 let test_cases = TestCases::default()
7477 .with_min_file_version(LanceFileVersion::V2_2)
7478 .with_max_file_version(LanceFileVersion::V2_2)
7479 .with_page_sizes(vec![4096]);
7480 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7481 }
7482
7483 #[tokio::test]
7484 async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
7485 use crate::format::pb21::page_layout::Layout;
7486 use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
7487
7488 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
7489 for _ in 0..64 {
7490 builder.values().append_value(1);
7491 builder.values().append_null();
7492 builder.values().append_value(3);
7493 builder.append(true);
7494 }
7495 let arr: ArrayRef = Arc::new(builder.finish());
7496 let field = arrow_schema::Field::new(
7497 "c",
7498 DataType::FixedSizeList(
7499 Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
7500 3,
7501 ),
7502 true,
7503 );
7504 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7505
7506 if let PageEncoding::Structural(layout) = &page.description {
7507 assert!(
7508 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7509 "FixedSizeList should not use constant layout yet"
7510 );
7511 }
7512
7513 let test_cases = TestCases::default()
7514 .with_min_file_version(LanceFileVersion::V2_2)
7515 .with_max_file_version(LanceFileVersion::V2_2)
7516 .with_page_sizes(vec![4096]);
7517 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7518 }
7519
7520 #[tokio::test]
7521 async fn test_constant_layout_not_written_before_v2_2() {
7522 use crate::format::pb21::page_layout::Layout;
7523
7524 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
7525 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7526 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await;
7527
7528 let PageEncoding::Structural(layout) = &page.description else {
7529 return;
7530 };
7531 assert!(
7532 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7533 "Should not emit constant layout before v2.2"
7534 );
7535
7536 let test_cases = TestCases::default()
7537 .with_min_file_version(LanceFileVersion::V2_1)
7538 .with_max_file_version(LanceFileVersion::V2_1)
7539 .with_page_sizes(vec![4096]);
7540 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7541 }
7542
7543 #[tokio::test]
7544 async fn test_all_null_constant_layout_still_works_v2_2() {
7545 use crate::format::pb21::page_layout::Layout;
7546
7547 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
7548 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7549 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7550
7551 let PageEncoding::Structural(layout) = &page.description else {
7552 panic!("Expected structural encoding");
7553 };
7554 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7555 panic!("Expected layout in slot 2");
7556 };
7557 assert!(layout.inline_value.is_none());
7558 assert_eq!(page.data.len(), 0);
7559
7560 let test_cases = TestCases::default()
7561 .with_min_file_version(LanceFileVersion::V2_2)
7562 .with_max_file_version(LanceFileVersion::V2_2)
7563 .with_page_sizes(vec![4096]);
7564 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7565 }
7566
7567 #[test]
7568 fn test_encode_decode_complex_all_null_vals_roundtrip() {
7569 use crate::compression::{
7570 DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy,
7571 };
7572
7573 let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
7574
7575 let compression_strategy = DefaultCompressionStrategy::default();
7576 let decompression_strategy = DefaultDecompressionStrategy::default();
7577
7578 let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
7579 &values,
7580 &compression_strategy,
7581 )
7582 .unwrap();
7583
7584 let decompressor = decompression_strategy
7585 .create_block_decompressor(&encoding)
7586 .unwrap();
7587 let decompressed = decompressor
7588 .decompress(compressed_buf, values.len() as u64)
7589 .unwrap();
7590 let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
7591 assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
7592 assert_eq!(decompressed_fixed_width.bits_per_value, 16);
7593 let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
7594 assert_eq!(rep_result.as_ref(), values.as_ref());
7595 }
7596
7597 #[tokio::test]
7598 async fn test_complex_all_null_compression_gated_by_version() {
7599 use crate::format::pb21::page_layout::Layout;
7600 use arrow_array::ListArray;
7601
7602 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7603 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7604 );
7605 let arr: ArrayRef = Arc::new(list_array);
7606 let field = arrow_schema::Field::new(
7607 "c",
7608 DataType::List(Arc::new(arrow_schema::Field::new(
7609 "item",
7610 DataType::Int32,
7611 true,
7612 ))),
7613 true,
7614 );
7615
7616 let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await;
7617 let PageEncoding::Structural(layout_v21) = &page_v21.description else {
7618 panic!("Expected structural encoding");
7619 };
7620 let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
7621 panic!("Expected constant layout");
7622 };
7623 assert!(layout_v21.rep_compression.is_none());
7624 assert!(layout_v21.def_compression.is_none());
7625 assert_eq!(layout_v21.num_rep_values, 0);
7626 assert_eq!(layout_v21.num_def_values, 0);
7627
7628 let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await;
7629 let PageEncoding::Structural(layout_v22) = &page_v22.description else {
7630 panic!("Expected structural encoding");
7631 };
7632 let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
7633 panic!("Expected constant layout");
7634 };
7635 assert!(layout_v22.def_compression.is_some());
7636 assert!(layout_v22.num_def_values > 0);
7637 }
7638
7639 #[tokio::test]
7640 async fn test_complex_all_null_round_trip() {
7641 use arrow_array::ListArray;
7642
7643 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7644 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7645 );
7646
7647 let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2);
7648 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
7649 .await;
7650 }
7651
7652 #[tokio::test]
7654 async fn test_sparse_boolean_list_roundtrip() {
7655 use arrow_array::builder::{BooleanBuilder, ListBuilder};
7656
7657 let mut list_builder = ListBuilder::new(BooleanBuilder::new());
7658 for i in 0..1000i32 {
7659 if i % 64 == 0 {
7660 list_builder.values().append_value(i % 128 == 0);
7662 list_builder.append(true);
7663 } else {
7664 list_builder.append(false);
7665 }
7666 }
7667 let list_array = Arc::new(list_builder.finish());
7668
7669 let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
7670 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
7671 }
7672
7673 fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
7674 use crate::compression::VariablePerValueDecompressor;
7675 use crate::encodings::physical::binary::VariableDecoder;
7676 use crate::repdef::{ControlWordParser, DefinitionInterpretation};
7677 use std::sync::Arc;
7678 Arc::new(super::FullZipDecodeDetails {
7679 value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
7680 VariableDecoder::default(),
7681 )
7682 as Arc<dyn VariablePerValueDecompressor>),
7683 def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
7684 ctrl_word_parser: ControlWordParser::new(0, 0),
7685 max_rep: 0,
7686 max_visible_def: 0,
7687 })
7688 }
7689
7690 fn decode_variable_full_zip(
7691 buf: Vec<u8>,
7692 bits_per_offset: u8,
7693 ) -> lance_core::Result<super::VariableFullZipDecoder> {
7694 use std::collections::VecDeque;
7695 let mut data = VecDeque::new();
7696 data.push_back(crate::buffer::LanceBuffer::from(buf));
7697 super::VariableFullZipDecoder::new(
7698 truncated_tail_details(),
7699 data,
7700 1,
7701 bits_per_offset,
7702 bits_per_offset,
7703 )
7704 }
7705
7706 #[test]
7715 fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
7716 use lance_core::Error;
7717
7718 for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
7719 let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
7720 .expect_err("a truncated length prefix must not decode");
7721 assert!(
7722 matches!(err, Error::CorruptFile { .. }),
7723 "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
7724 bits,
7725 buf_len,
7726 err
7727 );
7728 let msg = err.to_string();
7729 assert!(
7730 msg.contains("truncated length prefix"),
7731 "error should say what is wrong, got: {msg}"
7732 );
7733 }
7734 }
7735}