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