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 RepDefSlicer, SerializedRepDefs, StructuralPagePlan, 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
2705fn corrupt_file_named(name: &str, message: impl Into<String>) -> Error {
2706 Error::corrupt_file(name.into(), message)
2707}
2708
2709impl VariableFullZipDecoder {
2710 fn new(
2711 details: Arc<FullZipDecodeDetails>,
2712 data: VecDeque<LanceBuffer>,
2713 num_rows: u64,
2714 in_bits_per_length: u8,
2715 out_bits_per_offset: u8,
2716 ) -> Result<Self> {
2717 let decompressor = match details.value_decompressor {
2718 PerValueDecompressor::Variable(ref d) => d.clone(),
2719 _ => unreachable!(),
2720 };
2721
2722 assert_eq!(in_bits_per_length % 8, 0);
2723 assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
2724
2725 let mut decoder = Self {
2726 details,
2727 decompressor,
2728 data: LanceBuffer::empty(),
2729 offsets: LanceBuffer::empty(),
2730 rep: LanceBuffer::empty().borrow_to_typed_slice(),
2731 def: LanceBuffer::empty().borrow_to_typed_slice(),
2732 bits_per_offset: out_bits_per_offset,
2733 repdef_starts: Vec::with_capacity(num_rows as usize + 1),
2734 data_starts: Vec::with_capacity(num_rows as usize + 1),
2735 offset_starts: Vec::with_capacity(num_rows as usize + 1),
2736 visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
2737 current_idx: 0,
2738 num_rows,
2739 };
2740
2741 decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
2762
2763 Ok(decoder)
2764 }
2765
2766 fn slice_batch_data_and_rebase_offsets_typed<T>(
2767 data: &LanceBuffer,
2768 offsets: &LanceBuffer,
2769 ) -> Result<(LanceBuffer, LanceBuffer)>
2770 where
2771 T: arrow_buffer::ArrowNativeType
2772 + Copy
2773 + PartialOrd
2774 + std::ops::Sub<Output = T>
2775 + std::fmt::Display
2776 + TryInto<usize>,
2777 {
2778 let offsets_slice = offsets.borrow_to_typed_slice::<T>();
2779 let offsets_slice = offsets_slice.as_ref();
2780 if offsets_slice.is_empty() {
2781 return Err(Error::internal(
2782 "Variable offsets cannot be empty".to_string(),
2783 ));
2784 }
2785
2786 let base = offsets_slice[0];
2787 let end = *offsets_slice.last().unwrap();
2788 if end < base {
2789 return Err(Error::internal(format!(
2790 "Invalid variable offsets: end ({end}) is less than base ({base})"
2791 )));
2792 }
2793
2794 let data_start = base.try_into().map_err(|_| {
2795 Error::internal(format!("Variable offset ({base}) does not fit into usize"))
2796 })?;
2797 let data_end = end.try_into().map_err(|_| {
2798 Error::internal(format!("Variable offset ({end}) does not fit into usize"))
2799 })?;
2800 if data_end > data.len() {
2801 return Err(Error::internal(format!(
2802 "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
2803 data.len()
2804 )));
2805 }
2806
2807 let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
2808 for &offset in offsets_slice {
2809 if offset < base {
2810 return Err(Error::internal(format!(
2811 "Invalid variable offsets: offset ({offset}) is less than base ({base})"
2812 )));
2813 }
2814 rebased_offsets.push(offset - base);
2815 }
2816
2817 let sliced_data = data.slice_with_length(data_start, data_end - data_start);
2818 let sliced_data = LanceBuffer::copy_slice(&sliced_data);
2820 let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
2821 Ok((sliced_data, rebased_offsets))
2822 }
2823
2824 fn slice_batch_data_and_rebase_offsets(
2825 data: &LanceBuffer,
2826 offsets: &LanceBuffer,
2827 bits_per_offset: u8,
2828 ) -> Result<(LanceBuffer, LanceBuffer)> {
2829 match bits_per_offset {
2830 32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
2831 64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
2832 _ => Err(Error::internal(format!(
2833 "Unsupported bits_per_offset={bits_per_offset}"
2834 ))),
2835 }
2836 }
2837
2838 fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
2845 let width = bits_per_offset as usize / 8;
2846 if data.len() < width {
2847 return Err(corrupt_file_named(
2848 "variable_full_zip",
2849 format!(
2850 "truncated length prefix: {} byte(s) remain in the page buffer but a \
2851 {}-bit length prefix requires {}",
2852 data.len(),
2853 bits_per_offset,
2854 width
2855 ),
2856 ));
2857 }
2858 Ok(match bits_per_offset {
2859 8 => data[0] as u64,
2860 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
2861 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
2862 64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
2863 _ => unreachable!(),
2864 })
2865 }
2866
2867 fn unzip(
2868 &mut self,
2869 data: VecDeque<LanceBuffer>,
2870 in_bits_per_length: u8,
2871 out_bits_per_offset: u8,
2872 num_rows: u64,
2873 ) -> Result<()> {
2874 let mut rep = Vec::with_capacity(num_rows as usize);
2876 let mut def = Vec::with_capacity(num_rows as usize);
2877 let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
2878
2879 let bytes_per_offset = out_bits_per_offset as usize / 8;
2882 let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
2883 let mut offsets_data = Vec::with_capacity(bytes_offsets);
2884
2885 let bytes_per_length = in_bits_per_length as usize / 8;
2886 let bytes_lengths = bytes_per_length * num_rows as usize;
2887
2888 let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
2889 let mut unzipped_data =
2892 Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
2893
2894 let mut current_offset = 0_u64;
2895 let mut visible_item_count = 0_u64;
2896 for databuf in data.into_iter() {
2897 let mut databuf = databuf.as_ref();
2898 while !databuf.is_empty() {
2899 let data_start = unzipped_data.len();
2900 let offset_start = offsets_data.len();
2901 let repdef_start = rep.len().max(def.len());
2904 let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
2906 databuf,
2907 self.details.max_rep,
2908 self.details.max_visible_def,
2909 );
2910 self.details
2911 .ctrl_word_parser
2912 .parse(databuf, &mut rep, &mut def);
2913 databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
2914
2915 if ctrl_desc.is_new_row {
2916 self.repdef_starts.push(repdef_start);
2917 self.data_starts.push(data_start);
2918 self.offset_starts.push(offset_start);
2919 self.visible_item_counts.push(visible_item_count);
2920 }
2921 if ctrl_desc.is_visible {
2922 visible_item_count += 1;
2923 if ctrl_desc.is_valid_item {
2924 let length = Self::parse_length(databuf, in_bits_per_length)?;
2925 match out_bits_per_offset {
2926 32 => offsets_data
2927 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2928 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2929 _ => unreachable!(),
2930 };
2931 databuf = &databuf[bytes_per_offset..];
2932 unzipped_data.extend_from_slice(&databuf[..length as usize]);
2933 databuf = &databuf[length as usize..];
2934 current_offset += length;
2935 } else {
2936 match out_bits_per_offset {
2938 32 => offsets_data
2939 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2940 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2941 _ => unreachable!(),
2942 }
2943 }
2944 }
2945 }
2946 }
2947 self.repdef_starts.push(rep.len().max(def.len()));
2948 self.data_starts.push(unzipped_data.len());
2949 self.offset_starts.push(offsets_data.len());
2950 self.visible_item_counts.push(visible_item_count);
2951 match out_bits_per_offset {
2952 32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
2953 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
2954 _ => unreachable!(),
2955 };
2956 self.rep = ScalarBuffer::from(rep);
2957 self.def = ScalarBuffer::from(def);
2958 self.data = LanceBuffer::from(unzipped_data);
2959 self.offsets = LanceBuffer::from(offsets_data);
2960 Ok(())
2961 }
2962}
2963
2964impl StructuralPageDecoder for VariableFullZipDecoder {
2965 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2966 let start = self.current_idx;
2967 let end = start + num_rows as usize;
2968
2969 let offset_start = self.offset_starts[start];
2970 let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
2971 let offsets = self
2972 .offsets
2973 .slice_with_length(offset_start, offset_end - offset_start);
2974 let (data, offsets) =
2976 Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
2977
2978 let repdef_start = self.repdef_starts[start];
2979 let repdef_end = self.repdef_starts[end];
2980 let rep = if self.rep.is_empty() {
2981 self.rep.clone()
2982 } else {
2983 self.rep.slice(repdef_start, repdef_end - repdef_start)
2984 };
2985 let def = if self.def.is_empty() {
2986 self.def.clone()
2987 } else {
2988 self.def.slice(repdef_start, repdef_end - repdef_start)
2989 };
2990
2991 let visible_item_counts_start = self.visible_item_counts[start];
2992 let visible_item_counts_end = self.visible_item_counts[end];
2993 let num_visible_items = visible_item_counts_end - visible_item_counts_start;
2994
2995 self.current_idx += num_rows as usize;
2996
2997 Ok(Box::new(VariableFullZipDecodeTask {
2998 details: self.details.clone(),
2999 decompressor: self.decompressor.clone(),
3000 data,
3001 offsets,
3002 bits_per_offset: self.bits_per_offset,
3003 num_visible_items,
3004 rep,
3005 def,
3006 }))
3007 }
3008
3009 fn num_rows(&self) -> u64 {
3010 self.num_rows
3011 }
3012}
3013
3014#[derive(Debug)]
3015struct VariableFullZipDecodeTask {
3016 details: Arc<FullZipDecodeDetails>,
3017 decompressor: Arc<dyn VariablePerValueDecompressor>,
3018 data: LanceBuffer,
3019 offsets: LanceBuffer,
3020 bits_per_offset: u8,
3021 num_visible_items: u64,
3022 rep: ScalarBuffer<u16>,
3023 def: ScalarBuffer<u16>,
3024}
3025
3026impl DecodePageTask for VariableFullZipDecodeTask {
3027 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3028 let block = VariableWidthBlock {
3029 data: self.data,
3030 offsets: self.offsets,
3031 bits_per_offset: self.bits_per_offset,
3032 num_values: self.num_visible_items,
3033 block_info: BlockInfo::new(),
3034 };
3035 let decomopressed = self.decompressor.decompress(block)?;
3036 let rep = if self.rep.is_empty() {
3037 None
3038 } else {
3039 Some(self.rep.to_vec())
3040 };
3041 let def = if self.def.is_empty() {
3042 None
3043 } else {
3044 Some(self.def.to_vec())
3045 };
3046 let unraveler = RepDefUnraveler::new(
3047 rep,
3048 def,
3049 self.details.def_meaning.clone(),
3050 self.num_visible_items,
3051 );
3052 Ok(DecodedPage {
3053 data: decomopressed,
3054 repdef: unraveler,
3055 })
3056 }
3057}
3058
3059#[derive(Debug)]
3060struct FullZipDecodeTaskItem {
3061 data: PerValueDataBlock,
3062 rows_in_buf: u64,
3063}
3064
3065#[derive(Debug)]
3068struct FixedFullZipDecodeTask {
3069 details: Arc<FullZipDecodeDetails>,
3070 data: Vec<FullZipDecodeTaskItem>,
3071 num_rows: usize,
3072 bytes_per_value: usize,
3073}
3074
3075impl DecodePageTask for FixedFullZipDecodeTask {
3076 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3077 let estimated_size_bytes = self
3079 .data
3080 .iter()
3081 .map(|task_item| task_item.data.data_size() as usize)
3082 .sum::<usize>()
3083 * 2;
3084 let mut data_builder =
3085 DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3086
3087 if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3088 for task_item in self.data.into_iter() {
3092 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3093 unreachable!()
3094 };
3095 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3096 else {
3097 unreachable!()
3098 };
3099 debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3100 let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3101 data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3102 }
3103
3104 let unraveler = RepDefUnraveler::new(
3105 None,
3106 None,
3107 self.details.def_meaning.clone(),
3108 self.num_rows as u64,
3109 );
3110
3111 Ok(DecodedPage {
3112 data: data_builder.finish(),
3113 repdef: unraveler,
3114 })
3115 } else {
3116 let mut rep = Vec::with_capacity(self.num_rows);
3118 let mut def = Vec::with_capacity(self.num_rows);
3119
3120 for task_item in self.data.into_iter() {
3121 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3122 unreachable!()
3123 };
3124 let mut buf_slice = fixed_data.data.as_ref();
3125 let num_values = fixed_data.num_values as usize;
3126 let mut values = Vec::with_capacity(
3129 fixed_data.data.len()
3130 - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3131 );
3132 let mut visible_items = 0;
3133 for _ in 0..num_values {
3134 self.details
3136 .ctrl_word_parser
3137 .parse(buf_slice, &mut rep, &mut def);
3138 buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3139
3140 let is_visible = def
3141 .last()
3142 .map(|d| *d <= self.details.max_visible_def)
3143 .unwrap_or(true);
3144 if is_visible {
3145 values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3147 buf_slice = &buf_slice[self.bytes_per_value..];
3148 visible_items += 1;
3149 }
3150 }
3151
3152 let values_buf = LanceBuffer::from(values);
3154 let fixed_data = FixedWidthDataBlock {
3155 bits_per_value: self.bytes_per_value as u64 * 8,
3156 block_info: BlockInfo::new(),
3157 data: values_buf,
3158 num_values: visible_items,
3159 };
3160 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3161 else {
3162 unreachable!()
3163 };
3164 let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3165 data_builder.append(&decompressed, 0..visible_items);
3166 }
3167
3168 let repetition = if rep.is_empty() { None } else { Some(rep) };
3169 let definition = if def.is_empty() { None } else { Some(def) };
3170
3171 let unraveler = RepDefUnraveler::new(
3172 repetition,
3173 definition,
3174 self.details.def_meaning.clone(),
3175 self.num_rows as u64,
3176 );
3177 let data = data_builder.finish();
3178
3179 Ok(DecodedPage {
3180 data,
3181 repdef: unraveler,
3182 })
3183 }
3184 }
3185}
3186
3187#[derive(Debug)]
3188struct StructuralPrimitiveFieldSchedulingJob<'a> {
3189 scheduler: &'a StructuralPrimitiveFieldScheduler,
3190 ranges: Vec<Range<u64>>,
3191 page_idx: usize,
3192 range_idx: usize,
3193 global_row_offset: u64,
3194}
3195
3196impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3197 pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3198 Self {
3199 scheduler,
3200 ranges,
3201 page_idx: 0,
3202 range_idx: 0,
3203 global_row_offset: 0,
3204 }
3205 }
3206}
3207
3208impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3209 fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3210 if self.range_idx >= self.ranges.len() {
3211 return Ok(Vec::new());
3212 }
3213 let mut range = self.ranges[self.range_idx].clone();
3215 let priority = range.start;
3216
3217 let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3218 trace!(
3219 "Current range is {:?} and current page has {} rows",
3220 range, cur_page.num_rows
3221 );
3222 while cur_page.num_rows + self.global_row_offset <= range.start {
3224 self.global_row_offset += cur_page.num_rows;
3225 self.page_idx += 1;
3226 trace!("Skipping entire page of {} rows", cur_page.num_rows);
3227 cur_page = &self.scheduler.page_schedulers[self.page_idx];
3228 }
3229
3230 let mut ranges_in_page = Vec::new();
3234 while cur_page.num_rows + self.global_row_offset > range.start {
3235 range.start = range.start.max(self.global_row_offset);
3236 let start_in_page = range.start - self.global_row_offset;
3237 let end_in_page = start_in_page + (range.end - range.start);
3238 let end_in_page = end_in_page.min(cur_page.num_rows);
3239 let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3240
3241 ranges_in_page.push(start_in_page..end_in_page);
3242 if last_in_range {
3243 self.range_idx += 1;
3244 if self.range_idx == self.ranges.len() {
3245 break;
3246 }
3247 range = self.ranges[self.range_idx].clone();
3248 } else {
3249 break;
3250 }
3251 }
3252
3253 trace!(
3254 "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3255 ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3256 ranges_in_page.len(),
3257 cur_page.num_rows,
3258 priority,
3259 self.scheduler.column_index,
3260 cur_page.page_index,
3261 );
3262
3263 self.global_row_offset += cur_page.num_rows;
3264 self.page_idx += 1;
3265
3266 let page_decoders = cur_page
3267 .scheduler
3268 .schedule_ranges(&ranges_in_page, context.io())?;
3269
3270 let cur_path = context.current_path();
3271 page_decoders
3272 .into_iter()
3273 .map(|page_load_task| {
3274 let cur_path = cur_path.clone();
3275 let page_decoder = page_load_task.decoder_fut;
3276 let unloaded_page = async move {
3277 let page_decoder = page_decoder.await?;
3278 Ok(LoadedPageShard {
3279 decoder: page_decoder,
3280 path: cur_path,
3281 })
3282 }
3283 .boxed();
3284 Ok(ScheduledScanLine {
3285 decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3286 rows_scheduled: page_load_task.num_rows,
3287 })
3288 })
3289 .collect::<Result<Vec<_>>>()
3290 }
3291}
3292
3293#[derive(Debug)]
3294struct PageInfoAndScheduler {
3295 page_index: usize,
3296 num_rows: u64,
3297 scheduler: Box<dyn StructuralPageScheduler>,
3298}
3299
3300#[derive(Debug)]
3305pub struct StructuralPrimitiveFieldScheduler {
3306 page_schedulers: Vec<PageInfoAndScheduler>,
3307 column_index: u32,
3308 view_tag: String,
3314}
3315
3316impl StructuralPrimitiveFieldScheduler {
3317 pub fn try_new(
3318 column_info: &ColumnInfo,
3319 decompressors: &dyn DecompressionStrategy,
3320 cache_repetition_index: bool,
3321 target_field: &Field,
3322 ) -> Result<Self> {
3323 let page_schedulers = column_info
3324 .page_infos
3325 .iter()
3326 .enumerate()
3327 .map(|(page_index, page_info)| {
3328 Self::page_info_to_scheduler(
3329 page_info,
3330 page_index,
3331 decompressors,
3332 cache_repetition_index,
3333 target_field,
3334 )
3335 })
3336 .collect::<Result<Vec<_>>>()?;
3337 Ok(Self {
3338 page_schedulers,
3339 column_index: column_info.index,
3340 view_tag: format!("{:?}", target_field.data_type()),
3341 })
3342 }
3343
3344 fn page_layout_to_scheduler(
3345 page_info: &PageInfo,
3346 page_layout: &PageLayout,
3347 decompressors: &dyn DecompressionStrategy,
3348 cache_repetition_index: bool,
3349 target_field: &Field,
3350 ) -> Result<Box<dyn StructuralPageScheduler>> {
3351 use pb21::page_layout::Layout;
3352 Ok(match page_layout.layout.as_ref().expect_ok()? {
3353 Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
3354 &page_info.buffer_offsets_and_sizes,
3355 page_info.priority,
3356 mini_block.num_items,
3357 mini_block,
3358 decompressors,
3359 )?),
3360 Layout::FullZipLayout(full_zip) => {
3361 let mut scheduler = FullZipScheduler::try_new(
3362 &page_info.buffer_offsets_and_sizes,
3363 page_info.priority,
3364 page_info.num_rows,
3365 full_zip,
3366 decompressors,
3367 )?;
3368 scheduler.enable_cache = cache_repetition_index;
3369 Box::new(scheduler)
3370 }
3371 Layout::ConstantLayout(constant_layout) => {
3372 let def_meaning = constant_layout
3373 .layers
3374 .iter()
3375 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3376 .collect::<Vec<_>>();
3377 let has_scalar_value = constant_layout.inline_value.is_some()
3378 || page_info.buffer_offsets_and_sizes.len() == 1
3379 || page_info.buffer_offsets_and_sizes.len() == 3;
3380 if has_scalar_value {
3381 Box::new(constant::ConstantPageScheduler::try_new(
3382 page_info.buffer_offsets_and_sizes.clone(),
3383 constant_layout.inline_value.clone(),
3384 target_field.data_type(),
3385 def_meaning.into(),
3386 )?) as Box<dyn StructuralPageScheduler>
3387 } else if def_meaning.len() == 1
3388 && def_meaning[0] == DefinitionInterpretation::NullableItem
3389 {
3390 Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
3391 } else {
3392 let rep_decompressor = constant_layout
3393 .rep_compression
3394 .as_ref()
3395 .map(|encoding| decompressors.create_block_decompressor(encoding))
3396 .transpose()?
3397 .map(Arc::from);
3398
3399 let def_decompressor = constant_layout
3400 .def_compression
3401 .as_ref()
3402 .map(|encoding| decompressors.create_block_decompressor(encoding))
3403 .transpose()?
3404 .map(Arc::from);
3405
3406 Box::new(ComplexAllNullScheduler::new(
3407 page_info.buffer_offsets_and_sizes.clone(),
3408 def_meaning.into(),
3409 rep_decompressor,
3410 def_decompressor,
3411 constant_layout.num_rep_values,
3412 constant_layout.num_def_values,
3413 )) as Box<dyn StructuralPageScheduler>
3414 }
3415 }
3416 Layout::BlobLayout(blob) => {
3417 let inner_scheduler = Self::page_layout_to_scheduler(
3418 page_info,
3419 blob.inner_layout.as_ref().expect_ok()?.as_ref(),
3420 decompressors,
3421 cache_repetition_index,
3422 target_field,
3423 )?;
3424 let def_meaning = blob
3425 .layers
3426 .iter()
3427 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3428 .collect::<Vec<_>>();
3429 if matches!(target_field.data_type(), DataType::Struct(_)) {
3430 Box::new(BlobDescriptionPageScheduler::new(
3432 inner_scheduler,
3433 def_meaning.into(),
3434 ))
3435 } else {
3436 Box::new(BlobPageScheduler::new(
3438 inner_scheduler,
3439 page_info.priority,
3440 page_info.num_rows,
3441 def_meaning.into(),
3442 ))
3443 }
3444 }
3445 })
3446 }
3447
3448 fn page_info_to_scheduler(
3449 page_info: &PageInfo,
3450 page_index: usize,
3451 decompressors: &dyn DecompressionStrategy,
3452 cache_repetition_index: bool,
3453 target_field: &Field,
3454 ) -> Result<PageInfoAndScheduler> {
3455 let page_layout = page_info.encoding.as_structural();
3456 let scheduler = Self::page_layout_to_scheduler(
3457 page_info,
3458 page_layout,
3459 decompressors,
3460 cache_repetition_index,
3461 target_field,
3462 )?;
3463 Ok(PageInfoAndScheduler {
3464 page_index,
3465 num_rows: page_info.num_rows,
3466 scheduler,
3467 })
3468 }
3469}
3470
3471pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
3472 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
3473}
3474
3475pub struct NoCachedPageData;
3476
3477impl DeepSizeOf for NoCachedPageData {
3478 fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
3479 0
3480 }
3481}
3482impl CachedPageData for NoCachedPageData {
3483 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3484 self
3485 }
3486}
3487
3488pub struct CachedFieldData {
3489 pages: Vec<Arc<dyn CachedPageData>>,
3490}
3491
3492impl DeepSizeOf for CachedFieldData {
3493 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
3494 self.pages.deep_size_of_children(ctx)
3495 }
3496}
3497
3498#[derive(Debug, Clone)]
3508pub struct FieldDataCacheKey {
3509 pub column_index: u32,
3510 pub view_tag: String,
3511}
3512
3513impl CacheKey for FieldDataCacheKey {
3514 type ValueType = CachedFieldData;
3515
3516 fn key(&self) -> std::borrow::Cow<'_, str> {
3517 format!("{}:{}", self.column_index, self.view_tag).into()
3518 }
3519
3520 fn type_name() -> &'static str {
3521 "FieldData"
3522 }
3523}
3524
3525impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
3526 fn initialize<'a>(
3527 &'a mut self,
3528 _filter: &'a FilterExpression,
3529 context: &'a SchedulerContext,
3530 ) -> BoxFuture<'a, Result<()>> {
3531 let cache_key = FieldDataCacheKey {
3532 column_index: self.column_index,
3533 view_tag: self.view_tag.clone(),
3534 };
3535 let cache = context.cache().clone();
3536
3537 async move {
3538 if let Some(cached_data) = cache.get_with_key(&cache_key).await {
3539 self.page_schedulers
3540 .iter_mut()
3541 .zip(cached_data.pages.iter())
3542 .for_each(|(page_scheduler, cached_data)| {
3543 page_scheduler.scheduler.load(cached_data);
3544 });
3545 return Ok(());
3546 }
3547
3548 let page_data = self
3549 .page_schedulers
3550 .iter_mut()
3551 .map(|s| s.scheduler.initialize(context.io()))
3552 .collect::<FuturesOrdered<_>>();
3553
3554 let page_data = page_data.try_collect::<Vec<_>>().await?;
3555 let cached_data = Arc::new(CachedFieldData { pages: page_data });
3556 cache.insert_with_key(&cache_key, cached_data).await;
3557 Ok(())
3558 }
3559 .boxed()
3560 }
3561
3562 fn schedule_ranges<'a>(
3563 &'a self,
3564 ranges: &[Range<u64>],
3565 _filter: &FilterExpression,
3566 ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
3567 let ranges = ranges.to_vec();
3568 Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
3569 self, ranges,
3570 )))
3571 }
3572}
3573
3574#[derive(Debug)]
3577pub struct StructuralCompositeDecodeArrayTask {
3578 tasks: Vec<Box<dyn DecodePageTask>>,
3579 should_validate: bool,
3580 data_type: DataType,
3581}
3582
3583impl StructuralCompositeDecodeArrayTask {
3584 fn restore_validity(
3585 array: Arc<dyn Array>,
3586 unraveler: &mut CompositeRepDefUnraveler,
3587 ) -> Arc<dyn Array> {
3588 let validity = unraveler.unravel_validity(array.len());
3589 let Some(validity) = validity else {
3590 return array;
3591 };
3592 if array.data_type() == &DataType::Null {
3593 return array;
3595 }
3596 assert_eq!(validity.len(), array.len());
3597 make_array(unsafe {
3600 array
3601 .to_data()
3602 .into_builder()
3603 .nulls(Some(validity))
3604 .build_unchecked()
3605 })
3606 }
3607}
3608
3609impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
3610 fn decode(self: Box<Self>) -> Result<DecodedArray> {
3611 let mut arrays = Vec::with_capacity(self.tasks.len());
3612 let mut unravelers = Vec::with_capacity(self.tasks.len());
3613 let mut data_size = 0u64;
3614 for task in self.tasks {
3615 let decoded = task.decode()?;
3616 data_size += decoded.data.data_size();
3617 unravelers.push(decoded.repdef);
3618
3619 let array = make_array(
3620 decoded
3621 .data
3622 .into_arrow(self.data_type.clone(), self.should_validate)?,
3623 );
3624
3625 arrays.push(array);
3626 }
3627 let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
3628 let array = arrow_select::concat::concat(&array_refs)?;
3629 let mut repdef = CompositeRepDefUnraveler::new(unravelers);
3630
3631 let array = Self::restore_validity(array, &mut repdef);
3632
3633 Ok(DecodedArray {
3634 array,
3635 repdef,
3636 data_size,
3637 })
3638 }
3639}
3640
3641#[derive(Debug)]
3642pub struct StructuralPrimitiveFieldDecoder {
3643 field: Arc<ArrowField>,
3644 page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
3645 should_validate: bool,
3646 rows_drained_in_current: u64,
3647}
3648
3649impl StructuralPrimitiveFieldDecoder {
3650 pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
3651 Self {
3652 field: field.clone(),
3653 page_decoders: VecDeque::new(),
3654 should_validate,
3655 rows_drained_in_current: 0,
3656 }
3657 }
3658}
3659
3660impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
3661 fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
3662 assert!(child.path.is_empty());
3663 self.page_decoders.push_back(child.decoder);
3664 Ok(())
3665 }
3666
3667 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
3668 let mut remaining = num_rows;
3669 let mut tasks = Vec::new();
3670 while remaining > 0 {
3671 let cur_page = self.page_decoders.front_mut().unwrap();
3672 let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
3673 let to_take = num_in_page.min(remaining);
3674
3675 let task = cur_page.drain(to_take)?;
3676 tasks.push(task);
3677
3678 if to_take == num_in_page {
3679 self.page_decoders.pop_front();
3680 self.rows_drained_in_current = 0;
3681 } else {
3682 self.rows_drained_in_current += to_take;
3683 }
3684
3685 remaining -= to_take;
3686 }
3687 Ok(Box::new(StructuralCompositeDecodeArrayTask {
3688 tasks,
3689 should_validate: self.should_validate,
3690 data_type: self.field.data_type().clone(),
3691 }))
3692 }
3693
3694 fn data_type(&self) -> &DataType {
3695 self.field.data_type()
3696 }
3697}
3698
3699struct SerializedFullZip {
3701 values: LanceBuffer,
3703 repetition_index: Option<LanceBuffer>,
3705}
3706
3707const MINIBLOCK_ALIGNMENT: usize = 8;
3722
3723pub struct PrimitiveStructuralEncoder {
3750 accumulation_queue: AccumulationQueue,
3752
3753 keep_original_array: bool,
3754 support_large_chunk: bool,
3755 accumulated_repdefs: Vec<RepDefBuilder>,
3756 compression_strategy: Arc<dyn CompressionStrategy>,
3758 column_index: u32,
3759 field: Field,
3760 encoding_metadata: Arc<HashMap<String, String>>,
3761 version: LanceFileVersion,
3762}
3763
3764struct CompressedLevelsChunk {
3765 data: LanceBuffer,
3766 num_levels: u16,
3767}
3768
3769struct CompressedLevels {
3770 data: Vec<CompressedLevelsChunk>,
3771 compression: CompressiveEncoding,
3772 rep_index: Option<LanceBuffer>,
3773}
3774
3775struct SerializedMiniBlockPage {
3776 num_buffers: u64,
3777 data: LanceBuffer,
3778 metadata: LanceBuffer,
3779}
3780
3781#[derive(Debug, Clone, Copy)]
3782struct DictEncodingBudget {
3783 max_dict_entries: u32,
3784 max_encoded_size: usize,
3785}
3786
3787struct PrimitivePageData {
3789 arrays: Vec<ArrayRef>,
3791 repdef: SerializedRepDefs,
3793 row_number: u64,
3795 num_rows: u64,
3797 unsplittable_miniblock_levels: Option<u64>,
3799}
3800
3801#[derive(Clone)]
3806struct PrimitiveEncodeContext {
3807 column_idx: u32,
3809 field: Field,
3811 compression_strategy: Arc<dyn CompressionStrategy>,
3813 encoding_metadata: Arc<HashMap<String, String>>,
3815 support_large_chunk: bool,
3817 version: LanceFileVersion,
3819 is_simple_validity: bool,
3821 has_repdef_info: bool,
3823}
3824
3825impl PrimitiveStructuralEncoder {
3826 pub fn try_new(
3827 options: &EncodingOptions,
3828 compression_strategy: Arc<dyn CompressionStrategy>,
3829 column_index: u32,
3830 field: Field,
3831 encoding_metadata: Arc<HashMap<String, String>>,
3832 ) -> Result<Self> {
3833 Ok(Self {
3834 accumulation_queue: AccumulationQueue::new(
3835 options.cache_bytes_per_column,
3836 column_index,
3837 options.keep_original_array,
3838 ),
3839 support_large_chunk: options.support_large_chunk(),
3840 keep_original_array: options.keep_original_array,
3841 accumulated_repdefs: Vec::new(),
3842 column_index,
3843 compression_strategy,
3844 field,
3845 encoding_metadata,
3846 version: options.version,
3847 })
3848 }
3849
3850 fn is_narrow(data_block: &DataBlock) -> bool {
3858 const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
3859
3860 if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
3861 let max_len_array = max_len_array
3862 .as_any()
3863 .downcast_ref::<PrimitiveArray<UInt64Type>>()
3864 .unwrap();
3865 if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
3866 return true;
3867 }
3868 }
3869 false
3870 }
3871
3872 fn prefers_miniblock(
3873 data_block: &DataBlock,
3874 encoding_metadata: &HashMap<String, String>,
3875 ) -> bool {
3876 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3878 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
3879 }
3880 Self::is_narrow(data_block)
3882 }
3883
3884 fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
3885 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3889 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
3890 }
3891 true
3892 }
3893
3894 fn serialize_miniblocks(
3941 miniblocks: MiniBlockCompressed,
3942 rep: Option<Vec<CompressedLevelsChunk>>,
3943 def: Option<Vec<CompressedLevelsChunk>>,
3944 support_large_chunk: bool,
3945 ) -> Result<SerializedMiniBlockPage> {
3946 let bytes_rep = rep
3947 .as_ref()
3948 .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
3949 .unwrap_or(0);
3950 let bytes_def = def
3951 .as_ref()
3952 .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
3953 .unwrap_or(0);
3954 let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
3955 let mut num_buffers = miniblocks.data.len();
3956 if rep.is_some() {
3957 num_buffers += 1;
3958 }
3959 if def.is_some() {
3960 num_buffers += 1;
3961 }
3962 let max_extra = 9 * num_buffers;
3964 let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
3965 let chunk_size_bytes = if support_large_chunk { 4 } else { 2 };
3966 let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
3967
3968 let mut rep_iter = rep.map(|r| r.into_iter());
3969 let mut def_iter = def.map(|d| d.into_iter());
3970
3971 let mut buffer_offsets = vec![0; miniblocks.data.len()];
3972 for chunk in miniblocks.chunks {
3973 let start_pos = data_buffer.len();
3974 debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
3976
3977 let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
3978 let def = def_iter.as_mut().map(|d| d.next().unwrap());
3979
3980 let num_levels = rep
3982 .as_ref()
3983 .map(|r| r.num_levels)
3984 .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
3985 data_buffer.extend_from_slice(&num_levels.to_le_bytes());
3986
3987 if let Some(rep) = rep.as_ref() {
3989 let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
3990 Error::internal(format!(
3991 "Repetition buffer size ({} bytes) too large",
3992 rep.data.len()
3993 ))
3994 })?;
3995 data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
3996 }
3997 if let Some(def) = def.as_ref() {
3998 let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
3999 Error::internal(format!(
4000 "Definition buffer size ({} bytes) too large",
4001 def.data.len()
4002 ))
4003 })?;
4004 data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4005 }
4006
4007 if support_large_chunk {
4008 for &buffer_size in &chunk.buffer_sizes {
4009 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4010 }
4011 } else {
4012 for &buffer_size in &chunk.buffer_sizes {
4013 let buffer_size = u16::try_from(buffer_size).map_err(|_| {
4014 Error::internal(format!(
4015 "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
4016 buffer_size
4017 ))
4018 })?;
4019 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4020 }
4021 }
4022
4023 let add_padding = |data_buffer: &mut Vec<u8>| {
4025 let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4026 data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4027 };
4028 add_padding(&mut data_buffer);
4029
4030 if let Some(rep) = rep.as_ref() {
4032 data_buffer.extend_from_slice(&rep.data);
4033 add_padding(&mut data_buffer);
4034 }
4035 if let Some(def) = def.as_ref() {
4036 data_buffer.extend_from_slice(&def.data);
4037 add_padding(&mut data_buffer);
4038 }
4039 for (buffer_size, (buffer, buffer_offset)) in chunk
4040 .buffer_sizes
4041 .iter()
4042 .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4043 {
4044 let start = *buffer_offset;
4045 let end = start + *buffer_size as usize;
4046 *buffer_offset += *buffer_size as usize;
4047 data_buffer.extend_from_slice(&buffer[start..end]);
4048 add_padding(&mut data_buffer);
4049 }
4050
4051 let chunk_bytes = data_buffer.len() - start_pos;
4052 let max_chunk_size = if support_large_chunk {
4053 1_u64 << 31 } else {
4055 32 * 1024 };
4057 if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
4058 return Err(Error::internal(format!(
4059 "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
4060 chunk_bytes, max_chunk_size
4061 )));
4062 }
4063 if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
4064 return Err(Error::internal(format!(
4065 "Mini-block chunk size {} bytes is not aligned to {} bytes",
4066 chunk_bytes, MINIBLOCK_ALIGNMENT
4067 )));
4068 }
4069 if chunk.log_num_values > 15 {
4070 return Err(Error::internal(format!(
4071 "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
4072 chunk.log_num_values
4073 )));
4074 }
4075 let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4079 let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4080
4081 let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4082 if support_large_chunk {
4083 meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4084 } else {
4085 meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4086 }
4087 }
4088
4089 let data_buffer = LanceBuffer::from(data_buffer);
4090 let metadata_buffer = LanceBuffer::from(meta_buffer);
4091
4092 Ok(SerializedMiniBlockPage {
4093 num_buffers: miniblocks.data.len() as u64,
4094 data: data_buffer,
4095 metadata: metadata_buffer,
4096 })
4097 }
4098
4099 fn compress_levels(
4104 mut levels: RepDefSlicer<'_>,
4105 num_elements: u64,
4106 compression_strategy: &dyn CompressionStrategy,
4107 chunks: &[MiniBlockChunk],
4108 max_rep: u16,
4110 ) -> Result<CompressedLevels> {
4111 let mut rep_index = if max_rep > 0 {
4112 Vec::with_capacity(chunks.len())
4113 } else {
4114 vec![]
4115 };
4116 let num_levels = levels.num_levels() as u64;
4118 let levels_buf = levels.all_levels().clone();
4119
4120 let mut fixed_width_block = FixedWidthDataBlock {
4121 data: levels_buf,
4122 bits_per_value: 16,
4123 num_values: num_levels,
4124 block_info: BlockInfo::new(),
4125 };
4126 fixed_width_block.compute_stat();
4128
4129 let levels_block = DataBlock::FixedWidth(fixed_width_block);
4130 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4131 let (compressor, compressor_desc) =
4133 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4134 let mut level_chunks = Vec::with_capacity(chunks.len());
4136 let mut values_counter = 0;
4137 for (chunk_idx, chunk) in chunks.iter().enumerate() {
4138 let chunk_num_values = chunk.num_values(values_counter, num_elements);
4139 debug_assert!(chunk_num_values > 0);
4140 values_counter += chunk_num_values;
4141 let chunk_levels = if chunk_idx < chunks.len() - 1 {
4142 levels.slice_next(chunk_num_values as usize)
4143 } else {
4144 levels.slice_rest()
4145 };
4146 let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4147 if max_rep > 0 {
4148 let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4158 let rep_values = rep_values.as_ref();
4159
4160 let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4163 let num_leftovers = if chunk_idx < chunks.len() - 1 {
4164 rep_values
4165 .iter()
4166 .rev()
4167 .position(|v| *v == max_rep)
4168 .map(|pos| pos + 1)
4170 .unwrap_or(rep_values.len())
4171 } else {
4172 0
4174 };
4175
4176 if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
4177 let rep_len = rep_index.len();
4181 if rep_index[rep_len - 1] != 0 {
4182 rep_index[rep_len - 2] += 1;
4184 rep_index[rep_len - 1] = 0;
4185 }
4186 }
4187
4188 if chunk_idx == chunks.len() - 1 {
4189 num_rows += 1;
4191 }
4192 rep_index.push(num_rows as u64);
4193 rep_index.push(num_leftovers as u64);
4194 }
4195 let mut chunk_fixed_width = FixedWidthDataBlock {
4196 data: chunk_levels,
4197 bits_per_value: 16,
4198 num_values: num_chunk_levels,
4199 block_info: BlockInfo::new(),
4200 };
4201 chunk_fixed_width.compute_stat();
4202 let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
4203 let compressed_levels = compressor.compress(chunk_levels_block)?;
4204 let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
4205 Error::invalid_input_source(
4206 format!(
4207 "Mini-block cannot encode {} rep/def levels in one chunk. \
4208 This usually means a top-level row contains too much nested structure \
4209 for the current layout.",
4210 num_chunk_levels
4211 )
4212 .into(),
4213 )
4214 })?;
4215 level_chunks.push(CompressedLevelsChunk {
4216 data: compressed_levels,
4217 num_levels,
4218 });
4219 }
4220 debug_assert_eq!(levels.num_levels_remaining(), 0);
4221 let rep_index = if rep_index.is_empty() {
4222 None
4223 } else {
4224 Some(LanceBuffer::reinterpret_vec(rep_index))
4225 };
4226 Ok(CompressedLevels {
4227 data: level_chunks,
4228 compression: compressor_desc,
4229 rep_index,
4230 })
4231 }
4232
4233 fn encode_simple_all_null(
4234 column_idx: u32,
4235 num_rows: u64,
4236 row_number: u64,
4237 ) -> Result<EncodedPage> {
4238 let description =
4239 ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
4240 Ok(EncodedPage {
4241 column_idx,
4242 data: vec![],
4243 description: PageEncoding::Structural(description),
4244 num_rows,
4245 row_number,
4246 })
4247 }
4248
4249 fn encode_complex_all_null_vals(
4250 data: &Arc<[u16]>,
4251 compression_strategy: &dyn CompressionStrategy,
4252 ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
4253 let buffer = LanceBuffer::reinterpret_slice(data.clone());
4254 let mut fixed_width_block = FixedWidthDataBlock {
4255 data: buffer,
4256 bits_per_value: 16,
4257 num_values: data.len() as u64,
4258 block_info: BlockInfo::new(),
4259 };
4260 fixed_width_block.compute_stat();
4261
4262 let levels_block = DataBlock::FixedWidth(fixed_width_block);
4263 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4264 let (compressor, encoding) =
4265 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4266 let compressed_buffer = compressor.compress(levels_block)?;
4267 Ok((compressed_buffer, encoding))
4268 }
4269
4270 fn encode_complex_all_null(
4274 column_idx: u32,
4275 repdef: crate::repdef::SerializedRepDefs,
4276 row_number: u64,
4277 num_rows: u64,
4278 version: LanceFileVersion,
4279 compression_strategy: &dyn CompressionStrategy,
4280 ) -> Result<EncodedPage> {
4281 if version.resolve() < LanceFileVersion::V2_2 {
4282 let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
4283 LanceBuffer::reinterpret_slice(rep.clone())
4284 } else {
4285 LanceBuffer::empty()
4286 };
4287
4288 let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
4289 LanceBuffer::reinterpret_slice(def.clone())
4290 } else {
4291 LanceBuffer::empty()
4292 };
4293
4294 let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
4295 return Ok(EncodedPage {
4296 column_idx,
4297 data: vec![rep_bytes, def_bytes],
4298 description: PageEncoding::Structural(description),
4299 num_rows,
4300 row_number,
4301 });
4302 }
4303
4304 let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
4305 repdef.repetition_levels.as_ref()
4306 {
4307 let num_values = rep.len() as u64;
4308 let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
4309 (buffer, Some(encoding), num_values)
4310 } else {
4311 (LanceBuffer::empty(), None, 0)
4312 };
4313
4314 let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
4315 repdef.definition_levels.as_ref()
4316 {
4317 let num_values = def.len() as u64;
4318 let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
4319 (buffer, Some(encoding), num_values)
4320 } else {
4321 (LanceBuffer::empty(), None, 0)
4322 };
4323
4324 let description = ProtobufUtils21::compressed_all_null_constant_layout(
4325 &repdef.def_meaning,
4326 rep_encoding,
4327 def_encoding,
4328 num_rep_values,
4329 num_def_values,
4330 );
4331 Ok(EncodedPage {
4332 column_idx,
4333 data: vec![rep_bytes, def_bytes],
4334 description: PageEncoding::Structural(description),
4335 num_rows,
4336 row_number,
4337 })
4338 }
4339
4340 fn leaf_validity(
4341 repdef: &crate::repdef::SerializedRepDefs,
4342 num_values: usize,
4343 ) -> Result<Option<BooleanBuffer>> {
4344 let rep = repdef
4345 .repetition_levels
4346 .as_ref()
4347 .map(|rep| rep.as_ref().to_vec());
4348 let def = repdef
4349 .definition_levels
4350 .as_ref()
4351 .map(|def| def.as_ref().to_vec());
4352 let mut unraveler = RepDefUnraveler::new(
4353 rep,
4354 def,
4355 repdef.def_meaning.clone().into(),
4356 num_values as u64,
4357 );
4358 if unraveler.is_all_valid() {
4359 return Ok(None);
4360 }
4361 let mut validity = BooleanBufferBuilder::new(num_values);
4362 unraveler.unravel_validity(&mut validity);
4363 Ok(Some(validity.finish()))
4364 }
4365
4366 fn is_constant_values(
4367 arrays: &[ArrayRef],
4368 scalar: &ArrayRef,
4369 validity: Option<&BooleanBuffer>,
4370 ) -> Result<bool> {
4371 debug_assert_eq!(scalar.len(), 1);
4372 debug_assert_eq!(scalar.null_count(), 0);
4373
4374 match scalar.data_type() {
4375 DataType::Boolean => {
4376 let mut global_idx = 0usize;
4377 let scalar_val = scalar.as_boolean().value(0);
4378 for arr in arrays {
4379 let bool_arr = arr.as_boolean();
4380 for i in 0..arr.len() {
4381 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4382 global_idx += 1;
4383 if !is_valid {
4384 continue;
4385 }
4386 if bool_arr.value(i) != scalar_val {
4387 return Ok(false);
4388 }
4389 }
4390 }
4391 Ok(true)
4392 }
4393 DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
4394 DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
4395 DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
4396 DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
4397 data_type => {
4398 let mut global_idx = 0usize;
4399 let Some(byte_width) = data_type.byte_width_opt() else {
4400 return Ok(false);
4401 };
4402 let scalar_data = scalar.to_data();
4403 if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
4404 return Ok(false);
4405 }
4406 let scalar_bytes = scalar_data.buffers()[0].as_slice();
4407 if scalar_bytes.len() != byte_width {
4408 return Ok(false);
4409 }
4410
4411 for arr in arrays {
4412 let data = arr.to_data();
4413 if data.buffers().is_empty() {
4414 return Ok(false);
4415 }
4416 let buf = data.buffers()[0].as_slice();
4417 let base = data.offset();
4418 for i in 0..arr.len() {
4419 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4420 global_idx += 1;
4421 if !is_valid {
4422 continue;
4423 }
4424 let start = (base + i) * byte_width;
4425 if buf[start..start + byte_width] != scalar_bytes[..] {
4426 return Ok(false);
4427 }
4428 }
4429 }
4430 Ok(true)
4431 }
4432 }
4433 }
4434
4435 fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
4436 arrays: &[ArrayRef],
4437 scalar: &ArrayRef,
4438 validity: Option<&BooleanBuffer>,
4439 ) -> Result<bool> {
4440 debug_assert_eq!(scalar.len(), 1);
4441 let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
4442 let mut global_idx = 0usize;
4443 for arr in arrays {
4444 let str_arr = arr.as_string::<O>();
4445 for i in 0..arr.len() {
4446 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4447 global_idx += 1;
4448 if !is_valid {
4449 continue;
4450 }
4451 if str_arr.value(i).as_bytes() != scalar_val {
4452 return Ok(false);
4453 }
4454 }
4455 }
4456 Ok(true)
4457 }
4458
4459 fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
4460 arrays: &[ArrayRef],
4461 scalar: &ArrayRef,
4462 validity: Option<&BooleanBuffer>,
4463 ) -> Result<bool> {
4464 debug_assert_eq!(scalar.len(), 1);
4465 let scalar_val = scalar.as_binary::<O>().value(0);
4466 let mut global_idx = 0usize;
4467 for arr in arrays {
4468 let bin_arr = arr.as_binary::<O>();
4469 for i in 0..arr.len() {
4470 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4471 global_idx += 1;
4472 if !is_valid {
4473 continue;
4474 }
4475 if bin_arr.value(i) != scalar_val {
4476 return Ok(false);
4477 }
4478 }
4479 }
4480 Ok(true)
4481 }
4482
4483 fn find_constant_scalar(
4484 arrays: &[ArrayRef],
4485 validity: Option<&BooleanBuffer>,
4486 ) -> Result<Option<ArrayRef>> {
4487 if arrays.is_empty() {
4488 return Ok(None);
4489 }
4490
4491 let global_scalar_idx = if let Some(validity) = validity {
4492 let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
4493 return Ok(None);
4494 };
4495 idx
4496 } else {
4497 0
4498 };
4499
4500 let mut idx_remaining = global_scalar_idx;
4501 let mut scalar_arr_idx = 0usize;
4502 while scalar_arr_idx < arrays.len() {
4503 let len = arrays[scalar_arr_idx].len();
4504 if idx_remaining < len {
4505 break;
4506 }
4507 idx_remaining -= len;
4508 scalar_arr_idx += 1;
4509 }
4510
4511 if scalar_arr_idx >= arrays.len() {
4512 return Ok(None);
4513 }
4514
4515 let scalar =
4516 lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
4517 if scalar.null_count() != 0 {
4518 return Ok(None);
4519 }
4520 if !Self::is_constant_values(arrays, &scalar, validity)? {
4521 return Ok(None);
4522 }
4523 Ok(Some(scalar))
4524 }
4525
4526 fn resolve_dict_values_compression_metadata(
4527 field_metadata: &HashMap<String, String>,
4528 env_compression: Option<String>,
4529 env_compression_level: Option<String>,
4530 ) -> HashMap<String, String> {
4531 let mut metadata = HashMap::new();
4532
4533 let compression = field_metadata
4534 .get(DICT_VALUES_COMPRESSION_META_KEY)
4535 .cloned()
4536 .or(env_compression)
4537 .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
4538 metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
4539
4540 if let Some(compression_level) = field_metadata
4541 .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
4542 .cloned()
4543 .or(env_compression_level)
4544 {
4545 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
4546 }
4547
4548 metadata
4549 }
4550
4551 fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
4552 let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
4557 dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
4558 &field.metadata,
4559 env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
4560 env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
4561 );
4562 Ok(dict_values_field)
4563 }
4564
4565 #[allow(clippy::too_many_arguments)]
4566 fn encode_miniblock(
4567 column_idx: u32,
4568 field: &Field,
4569 compression_strategy: &dyn CompressionStrategy,
4570 data: DataBlock,
4571 repdef: crate::repdef::SerializedRepDefs,
4572 row_number: u64,
4573 dictionary_data: Option<DataBlock>,
4574 num_rows: u64,
4575 support_large_chunk: bool,
4576 ) -> Result<EncodedPage> {
4577 if let DataBlock::AllNull(_null_block) = data {
4578 unreachable!()
4581 }
4582
4583 let num_items = data.num_values();
4584
4585 let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
4586 let (compressed_data, value_encoding) = compressor.compress(data)?;
4587
4588 let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
4589
4590 let mut compressed_rep = repdef
4591 .rep_slicer()
4592 .map(|rep_slicer| {
4593 Self::compress_levels(
4594 rep_slicer,
4595 num_items,
4596 compression_strategy,
4597 &compressed_data.chunks,
4598 max_rep,
4599 )
4600 })
4601 .transpose()?;
4602
4603 let (rep_index, rep_index_depth) =
4604 match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
4605 Some(rep_index) => (Some(rep_index.clone()), 1),
4606 None => (None, 0),
4607 };
4608
4609 let mut compressed_def = repdef
4610 .def_slicer()
4611 .map(|def_slicer| {
4612 Self::compress_levels(
4613 def_slicer,
4614 num_items,
4615 compression_strategy,
4616 &compressed_data.chunks,
4617 0,
4618 )
4619 })
4620 .transpose()?;
4621
4622 let rep_data = compressed_rep
4628 .as_mut()
4629 .map(|cr| std::mem::take(&mut cr.data));
4630 let def_data = compressed_def
4631 .as_mut()
4632 .map(|cd| std::mem::take(&mut cd.data));
4633
4634 let serialized =
4635 Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?;
4636
4637 let mut data = Vec::with_capacity(4);
4639 data.push(serialized.metadata);
4640 data.push(serialized.data);
4641
4642 if let Some(dictionary_data) = dictionary_data {
4643 let num_dictionary_items = dictionary_data.num_values();
4644 let dict_values_field = Self::build_dict_values_compressor_field(field)?;
4645
4646 let (compressor, dictionary_encoding) = compression_strategy
4647 .create_block_compressor(&dict_values_field, &dictionary_data)?;
4648 let dictionary_buffer = compressor.compress(dictionary_data)?;
4649
4650 data.push(dictionary_buffer);
4651 if let Some(rep_index) = rep_index {
4652 data.push(rep_index);
4653 }
4654
4655 let description = ProtobufUtils21::miniblock_layout(
4656 compressed_rep.map(|cr| cr.compression),
4657 compressed_def.map(|cd| cd.compression),
4658 value_encoding,
4659 rep_index_depth,
4660 serialized.num_buffers,
4661 Some((dictionary_encoding, num_dictionary_items)),
4662 &repdef.def_meaning,
4663 num_items,
4664 support_large_chunk,
4665 );
4666 Ok(EncodedPage {
4667 num_rows,
4668 column_idx,
4669 data,
4670 description: PageEncoding::Structural(description),
4671 row_number,
4672 })
4673 } else {
4674 let description = ProtobufUtils21::miniblock_layout(
4675 compressed_rep.map(|cr| cr.compression),
4676 compressed_def.map(|cd| cd.compression),
4677 value_encoding,
4678 rep_index_depth,
4679 serialized.num_buffers,
4680 None,
4681 &repdef.def_meaning,
4682 num_items,
4683 support_large_chunk,
4684 );
4685
4686 if let Some(rep_index) = rep_index {
4687 let view = rep_index.borrow_to_typed_slice::<u64>();
4688 let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
4689 debug_assert_eq!(total, num_rows);
4690
4691 data.push(rep_index);
4692 }
4693
4694 Ok(EncodedPage {
4695 num_rows,
4696 column_idx,
4697 data,
4698 description: PageEncoding::Structural(description),
4699 row_number,
4700 })
4701 }
4702 }
4703
4704 fn serialize_full_zip_fixed(
4706 fixed: FixedWidthDataBlock,
4707 mut repdef: ControlWordIterator,
4708 num_values: u64,
4709 ) -> Result<SerializedFullZip> {
4710 if !fixed.bits_per_value.is_multiple_of(8) {
4711 return Err(Error::invalid_input_source(
4712 format!(
4713 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
4714 fixed.bits_per_value
4715 )
4716 .into(),
4717 ));
4718 }
4719
4720 let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
4721 let mut zipped_data = Vec::with_capacity(len);
4722
4723 let max_rep_index_val = if repdef.has_repetition() {
4724 len as u64
4725 } else {
4726 0
4728 };
4729 let mut rep_index_builder =
4730 BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
4731
4732 let bytes_per_value = fixed.bits_per_value as usize / 8;
4733 let mut offset = 0;
4734
4735 if bytes_per_value == 0 {
4736 while let Some(control) = repdef.append_next(&mut zipped_data) {
4738 if control.is_new_row {
4739 debug_assert!(offset <= len);
4741 unsafe { rep_index_builder.append(offset as u64) };
4743 }
4744 offset = zipped_data.len();
4745 }
4746 } else {
4747 let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
4749 while let Some(control) = repdef.append_next(&mut zipped_data) {
4750 if control.is_new_row {
4751 debug_assert!(offset <= len);
4753 unsafe { rep_index_builder.append(offset as u64) };
4755 }
4756 if control.is_visible {
4757 let value = data_iter.next().unwrap();
4758 zipped_data.extend_from_slice(value);
4759 }
4760 offset = zipped_data.len();
4761 }
4762 }
4763
4764 debug_assert_eq!(zipped_data.len(), len);
4765 unsafe {
4768 rep_index_builder.append(zipped_data.len() as u64);
4769 }
4770
4771 let zipped_data = LanceBuffer::from(zipped_data);
4772 let rep_index = rep_index_builder.into_data();
4773 let rep_index = if rep_index.is_empty() {
4774 None
4775 } else {
4776 Some(LanceBuffer::from(rep_index))
4777 };
4778 Ok(SerializedFullZip {
4779 values: zipped_data,
4780 repetition_index: rep_index,
4781 })
4782 }
4783
4784 fn serialize_full_zip_variable(
4788 variable: VariableWidthBlock,
4789 mut repdef: ControlWordIterator,
4790 num_items: u64,
4791 ) -> Result<SerializedFullZip> {
4792 let bytes_per_offset = variable.bits_per_offset as usize / 8;
4793 if !variable.bits_per_offset.is_multiple_of(8) {
4794 return Err(Error::invalid_input_source(
4795 format!(
4796 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
4797 variable.bits_per_offset
4798 )
4799 .into(),
4800 ));
4801 }
4802 let len = variable.data.len()
4803 + repdef.bytes_per_word() * num_items as usize
4804 + bytes_per_offset * variable.num_values as usize;
4805 let mut buf = Vec::with_capacity(len);
4806
4807 let max_rep_index_val = len as u64;
4808 let mut rep_index_builder =
4809 BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
4810
4811 match bytes_per_offset {
4813 4 => {
4814 let offs = variable.offsets.borrow_to_typed_slice::<u32>();
4815 let mut rep_offset = 0;
4816 let mut windows_iter = offs.as_ref().windows(2);
4817 while let Some(control) = repdef.append_next(&mut buf) {
4818 if control.is_new_row {
4819 debug_assert!(rep_offset <= len);
4821 unsafe { rep_index_builder.append(rep_offset as u64) };
4823 }
4824 if control.is_visible {
4825 let window = windows_iter.next().unwrap();
4826 if control.is_valid_item {
4827 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4828 buf.extend_from_slice(
4829 &variable.data[window[0] as usize..window[1] as usize],
4830 );
4831 }
4832 }
4833 rep_offset = buf.len();
4834 }
4835 }
4836 8 => {
4837 let offs = variable.offsets.borrow_to_typed_slice::<u64>();
4838 let mut rep_offset = 0;
4839 let mut windows_iter = offs.as_ref().windows(2);
4840 while let Some(control) = repdef.append_next(&mut buf) {
4841 if control.is_new_row {
4842 debug_assert!(rep_offset <= len);
4844 unsafe { rep_index_builder.append(rep_offset as u64) };
4846 }
4847 if control.is_visible {
4848 let window = windows_iter.next().unwrap();
4849 if control.is_valid_item {
4850 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4851 buf.extend_from_slice(
4852 &variable.data[window[0] as usize..window[1] as usize],
4853 );
4854 }
4855 }
4856 rep_offset = buf.len();
4857 }
4858 }
4859 _ => {
4860 return Err(Error::invalid_input_source(
4861 format!(
4862 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
4863 variable.bits_per_offset
4864 )
4865 .into(),
4866 ));
4867 }
4868 }
4869
4870 debug_assert!(buf.len() <= len);
4873 unsafe {
4876 rep_index_builder.append(buf.len() as u64);
4877 }
4878
4879 let zipped_data = LanceBuffer::from(buf);
4880 let rep_index = rep_index_builder.into_data();
4881 debug_assert!(!rep_index.is_empty());
4882 let rep_index = Some(LanceBuffer::from(rep_index));
4883 Ok(SerializedFullZip {
4884 values: zipped_data,
4885 repetition_index: rep_index,
4886 })
4887 }
4888
4889 fn serialize_full_zip(
4892 compressed_data: PerValueDataBlock,
4893 repdef: ControlWordIterator,
4894 num_items: u64,
4895 ) -> Result<SerializedFullZip> {
4896 match compressed_data {
4897 PerValueDataBlock::Fixed(fixed) => {
4898 Self::serialize_full_zip_fixed(fixed, repdef, num_items)
4899 }
4900 PerValueDataBlock::Variable(var) => {
4901 Self::serialize_full_zip_variable(var, repdef, num_items)
4902 }
4903 }
4904 }
4905
4906 fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
4907 debug_assert_eq!(fixed.bits_per_value, 1);
4908 let num_values = fixed.num_values as usize;
4909 let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
4910 let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
4911 FixedWidthDataBlock {
4912 data: LanceBuffer::from(expanded),
4913 bits_per_value: 8,
4914 num_values: fixed.num_values,
4915 block_info: BlockInfo::new(),
4916 }
4917 }
4918
4919 fn encode_full_zip(
4920 column_idx: u32,
4921 field: &Field,
4922 compression_strategy: &dyn CompressionStrategy,
4923 data: DataBlock,
4924 repdef: crate::repdef::SerializedRepDefs,
4925 row_number: u64,
4926 num_lists: u64,
4927 ) -> Result<EncodedPage> {
4928 let max_rep = repdef
4929 .repetition_levels
4930 .as_ref()
4931 .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
4932 let max_def = repdef
4933 .definition_levels
4934 .as_ref()
4935 .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
4936
4937 let (num_items, num_visible_items) =
4941 if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
4942 (rep_levels.len() as u64, data.num_values())
4945 } else {
4946 (data.num_values(), data.num_values())
4948 };
4949
4950 let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
4951
4952 let repdef_iter = build_control_word_iterator(
4953 repdef.repetition_levels.as_deref(),
4954 max_rep,
4955 repdef.definition_levels.as_deref(),
4956 max_def,
4957 max_visible_def,
4958 num_items as usize,
4959 );
4960 let bits_rep = repdef_iter.bits_rep();
4961 let bits_def = repdef_iter.bits_def();
4962
4963 let data = match data {
4965 DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
4966 DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
4967 }
4968 other => other,
4969 };
4970
4971 let compressor = compression_strategy.create_per_value(field, &data)?;
4972 let (compressed_data, value_encoding) = compressor.compress(data)?;
4973
4974 let description = match &compressed_data {
4975 PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
4976 bits_rep,
4977 bits_def,
4978 fixed.bits_per_value as u32,
4979 value_encoding,
4980 &repdef.def_meaning,
4981 num_items as u32,
4982 num_visible_items as u32,
4983 ),
4984 PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
4985 bits_rep,
4986 bits_def,
4987 variable.bits_per_offset as u32,
4988 value_encoding,
4989 &repdef.def_meaning,
4990 num_items as u32,
4991 num_visible_items as u32,
4992 ),
4993 };
4994
4995 let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
4996
4997 let data = if let Some(repindex) = zipped.repetition_index {
4998 vec![zipped.values, repindex]
4999 } else {
5000 vec![zipped.values]
5001 };
5002
5003 Ok(EncodedPage {
5004 num_rows: num_lists,
5005 column_idx,
5006 data,
5007 description: PageEncoding::Structural(description),
5008 row_number,
5009 })
5010 }
5011
5012 fn should_dictionary_encode(
5013 data_block: &DataBlock,
5014 field: &Field,
5015 version: LanceFileVersion,
5016 ) -> Option<DictEncodingBudget> {
5017 const DEFAULT_SAMPLE_SIZE: usize = 4096;
5018 const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5019
5020 match data_block {
5023 DataBlock::FixedWidth(fixed) => {
5024 if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
5025 return None;
5026 }
5027 if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5028 return None;
5029 }
5030 if fixed.bits_per_value % 8 != 0 {
5031 return None;
5032 }
5033 }
5034 DataBlock::VariableWidth(var) => {
5035 if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5036 return None;
5037 }
5038 }
5039 _ => return None,
5040 }
5041
5042 let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5044 .ok()
5045 .and_then(|val| val.parse().ok())
5046 .unwrap_or(100);
5047 if data_block.num_values() < too_small {
5048 return None;
5049 }
5050
5051 let num_values = data_block.num_values();
5052
5053 let divisor: u64 = field
5056 .metadata
5057 .get(DICT_DIVISOR_META_KEY)
5058 .and_then(|val| val.parse().ok())
5059 .or_else(|| {
5060 env::var("LANCE_ENCODING_DICT_DIVISOR")
5061 .ok()
5062 .and_then(|val| val.parse().ok())
5063 })
5064 .unwrap_or(DEFAULT_DICT_DIVISOR);
5065
5066 let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5067 .ok()
5068 .and_then(|val| val.parse().ok())
5069 .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5070
5071 let threshold_cardinality = num_values
5072 .checked_div(divisor.max(1))
5073 .unwrap_or(0)
5074 .min(max_cardinality);
5075 if threshold_cardinality == 0 {
5076 return None;
5077 }
5078
5079 let threshold_ratio = field
5081 .metadata
5082 .get(DICT_SIZE_RATIO_META_KEY)
5083 .and_then(|val| val.parse::<f64>().ok())
5084 .or_else(|| {
5085 env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5086 .ok()
5087 .and_then(|val| val.parse().ok())
5088 })
5089 .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5090
5091 if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5092 panic!(
5093 "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5094 threshold_ratio
5095 );
5096 }
5097
5098 let data_size = data_block.data_size();
5099 if data_size == 0 {
5100 return None;
5101 }
5102
5103 let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5104 let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5105
5106 if let Some(sample_unique_ratio) =
5109 Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
5110 {
5111 if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
5112 return None;
5113 }
5114
5115 let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
5116 if projected_cardinality > threshold_cardinality {
5117 return None;
5118 }
5119 }
5120
5121 let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5122 Some(DictEncodingBudget {
5123 max_dict_entries,
5124 max_encoded_size,
5125 })
5126 }
5127
5128 fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
5136 use std::collections::HashSet;
5137
5138 const NUM_SAMPLE_BLOCKS: usize = 32;
5139 const MIN_RELIABLE_SAMPLES: usize = 1024;
5140
5141 let num_values = usize::try_from(data_block.num_values()).ok()?;
5142 if num_values == 0 {
5143 return Some(None);
5144 }
5145
5146 let sample_count = num_values.min(max_samples).max(1);
5147 if sample_count < MIN_RELIABLE_SAMPLES {
5148 return Some(None);
5149 }
5150
5151 let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
5152 let samples_per_block = (sample_count / block_count).max(1);
5153 let mut indices = Vec::with_capacity(sample_count);
5154 for block_idx in 0..block_count {
5155 let block_start = block_idx * num_values / block_count;
5156 let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
5157 let block_len = next_block_start.saturating_sub(block_start);
5158 let samples_in_block = samples_per_block.min(block_len);
5159 indices.extend((0..samples_in_block).map(|offset| block_start + offset));
5160 }
5161
5162 if indices.len() < MIN_RELIABLE_SAMPLES {
5163 return Some(None);
5164 }
5165
5166 let ratio = match data_block {
5167 DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
5168 64 => {
5169 let values = fixed.data.borrow_to_typed_slice::<u64>();
5170 let values = values.as_ref();
5171 let mut unique: HashSet<u64> =
5172 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5173 for idx in indices.iter().copied() {
5174 unique.insert(values.get(idx).copied()?);
5175 }
5176 unique.len() as f64 / indices.len() as f64
5177 }
5178 128 => {
5179 let values = fixed.data.borrow_to_typed_slice::<u128>();
5180 let values = values.as_ref();
5181 let mut unique: HashSet<u128> =
5182 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5183 for idx in indices.iter().copied() {
5184 unique.insert(values.get(idx).copied()?);
5185 }
5186 unique.len() as f64 / indices.len() as f64
5187 }
5188 _ => return Some(None),
5189 },
5190 DataBlock::VariableWidth(var) => {
5191 use xxhash_rust::xxh3::xxh3_64;
5192
5193 let mut unique: HashSet<u64> =
5195 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
5196 match var.bits_per_offset {
5197 32 => {
5198 let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
5199 let offsets: &[u32] = offsets_ref.as_ref();
5200 for i in indices.iter().copied() {
5201 let start = usize::try_from(*offsets.get(i)?).ok()?;
5202 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5203 if start > end || end > var.data.len() {
5204 return None;
5205 }
5206 unique.insert(xxh3_64(&var.data[start..end]));
5207 }
5208 }
5209 64 => {
5210 let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
5211 let offsets: &[u64] = offsets_ref.as_ref();
5212 for i in indices.iter().copied() {
5213 let start = usize::try_from(*offsets.get(i)?).ok()?;
5214 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5215 if start > end || end > var.data.len() {
5216 return None;
5217 }
5218 unique.insert(xxh3_64(&var.data[start..end]));
5219 }
5220 }
5221 _ => return Some(None),
5222 }
5223 unique.len() as f64 / indices.len() as f64
5224 }
5225 _ => return Some(None),
5226 };
5227
5228 Some(Some(ratio))
5229 }
5230
5231 fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
5232 let repetition_levels = repdef
5233 .repetition_levels
5234 .as_ref()
5235 .map(|levels| levels[range.clone()].to_vec());
5236 let definition_levels = repdef
5237 .definition_levels
5238 .as_ref()
5239 .map(|levels| levels[range].to_vec());
5240 SerializedRepDefs::new_with_fixed_size_list_levels(
5241 repetition_levels,
5242 definition_levels,
5243 repdef.def_meaning.clone(),
5244 repdef.has_fixed_size_list_levels(),
5245 )
5246 }
5247
5248 fn slice_arrays(
5249 arrays: &[ArrayRef],
5250 value_start: u64,
5251 num_values: u64,
5252 ) -> Result<Vec<ArrayRef>> {
5253 if num_values == 0 {
5254 return Ok(Vec::new());
5255 }
5256
5257 let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
5258 Error::invalid_input(format!("Value start {} is too large", value_start))
5259 })?;
5260 let mut values_remaining = usize::try_from(num_values).map_err(|_| {
5261 Error::invalid_input(format!("Value count {} is too large", num_values))
5262 })?;
5263 let mut sliced = Vec::new();
5264
5265 for array in arrays {
5266 if values_to_skip >= array.len() {
5267 values_to_skip -= array.len();
5268 continue;
5269 }
5270
5271 let offset = values_to_skip;
5272 let len = (array.len() - offset).min(values_remaining);
5273 sliced.push(array.slice(offset, len));
5274 values_remaining -= len;
5275 values_to_skip = 0;
5276
5277 if values_remaining == 0 {
5278 break;
5279 }
5280 }
5281
5282 if values_remaining != 0 {
5283 return Err(Error::internal(format!(
5284 "Page split requested {} values starting at {}, but the page did not contain enough values",
5285 num_values, value_start
5286 )));
5287 }
5288
5289 Ok(sliced)
5290 }
5291
5292 fn split_structural_pages_for_miniblock_budget(
5293 arrays: Vec<ArrayRef>,
5294 repdef: SerializedRepDefs,
5295 plan: StructuralPagePlan,
5296 row_number: u64,
5297 num_rows: u64,
5298 ) -> Result<Vec<PrimitivePageData>> {
5299 if plan == StructuralPagePlan::Fits {
5300 return Ok(vec![PrimitivePageData {
5301 arrays,
5302 repdef,
5303 row_number,
5304 num_rows,
5305 unsplittable_miniblock_levels: None,
5306 }]);
5307 }
5308 if let StructuralPagePlan::UnsplittableOverBudget(num_levels) = plan {
5309 return Ok(vec![PrimitivePageData {
5310 arrays,
5311 repdef,
5312 row_number,
5313 num_rows,
5314 unsplittable_miniblock_levels: Some(num_levels),
5315 }]);
5316 }
5317
5318 let StructuralPagePlan::Split(splits) = plan else {
5319 unreachable!();
5320 };
5321
5322 let mut pages = Vec::with_capacity(splits.len());
5323 for split in splits {
5324 let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
5325 let repdef = Self::slice_repdef(&repdef, split.level_range);
5326 pages.push(PrimitivePageData {
5327 arrays,
5328 repdef,
5329 row_number: row_number + split.row_start,
5330 num_rows: split.num_rows,
5331 unsplittable_miniblock_levels: None,
5332 });
5333 }
5334 Ok(pages)
5335 }
5336
5337 fn encode_page(ctx: PrimitiveEncodeContext, page: PrimitivePageData) -> Result<EncodedPage> {
5338 let PrimitiveEncodeContext {
5339 column_idx,
5340 field,
5341 compression_strategy,
5342 encoding_metadata,
5343 support_large_chunk,
5344 version,
5345 is_simple_validity,
5346 has_repdef_info,
5347 } = ctx;
5348 let PrimitivePageData {
5349 arrays,
5350 repdef,
5351 row_number,
5352 num_rows,
5353 unsplittable_miniblock_levels,
5354 } = page;
5355 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5356
5357 if num_values == 0 {
5358 log::debug!(
5361 "Encoding column {} with {} items ({} rows) using complex-null layout",
5362 column_idx,
5363 num_values,
5364 num_rows
5365 );
5366 return Self::encode_complex_all_null(
5367 column_idx,
5368 repdef,
5369 row_number,
5370 num_rows,
5371 version,
5372 compression_strategy.as_ref(),
5373 );
5374 }
5375
5376 let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
5377 let all_null = leaf_validity
5378 .as_ref()
5379 .map(|validity| validity.count_set_bits() == 0)
5380 .unwrap_or(false);
5381
5382 if all_null {
5383 return if is_simple_validity {
5384 log::debug!(
5385 "Encoding column {} with {} items ({} rows) using simple-null layout",
5386 column_idx,
5387 num_values,
5388 num_rows
5389 );
5390 Self::encode_simple_all_null(column_idx, num_values, row_number)
5391 } else {
5392 log::debug!(
5393 "Encoding column {} with {} items ({} rows) using complex-null layout",
5394 column_idx,
5395 num_values,
5396 num_rows
5397 );
5398 Self::encode_complex_all_null(
5399 column_idx,
5400 repdef,
5401 row_number,
5402 num_rows,
5403 version,
5404 compression_strategy.as_ref(),
5405 )
5406 };
5407 }
5408
5409 if let DataType::Struct(fields) = &field.data_type()
5410 && fields.is_empty()
5411 {
5412 if has_repdef_info {
5413 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()));
5414 }
5415 return Self::encode_simple_all_null(column_idx, num_values, row_number);
5418 }
5419
5420 let data_block = DataBlock::from_arrays(&arrays, num_values);
5421
5422 if version.resolve() >= LanceFileVersion::V2_2
5423 && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())?
5424 {
5425 log::debug!(
5426 "Encoding column {} with {} items ({} rows) using constant layout",
5427 column_idx,
5428 num_values,
5429 num_rows
5430 );
5431 return constant::encode_constant_page(
5432 column_idx, scalar, repdef, row_number, num_rows,
5433 );
5434 }
5435
5436 if let Some(num_levels) = unsplittable_miniblock_levels {
5437 let requested_encoding = encoding_metadata
5438 .get(STRUCTURAL_ENCODING_META_KEY)
5439 .map(|requested| requested.to_lowercase());
5440 let fullzip_error = match &data_block {
5441 DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => {
5442 Some(format!(
5443 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
5444 fixed.bits_per_value
5445 ))
5446 }
5447 DataBlock::VariableWidth(variable)
5448 if !variable.bits_per_offset.is_multiple_of(8) =>
5449 {
5450 Some(format!(
5451 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
5452 variable.bits_per_offset
5453 ))
5454 }
5455 DataBlock::VariableWidth(variable)
5456 if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
5457 {
5458 Some(format!(
5459 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
5460 variable.bits_per_offset
5461 ))
5462 }
5463 DataBlock::Struct(struct_data_block)
5464 if !struct_data_block.has_variable_width_child() =>
5465 {
5466 Some(
5467 "Full-zip packed struct requires at least one variable-width child"
5468 .to_string(),
5469 )
5470 }
5471 DataBlock::Dictionary(_) => {
5472 Some("Full-zip does not encode dictionary data blocks directly".to_string())
5473 }
5474 DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
5475 Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
5476 Some(flat) => Some(format!(
5477 "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
5478 flat.bits_per_value
5479 )),
5480 None => Some(
5481 "Full-zip fixed-size-list capability requires a flat fixed-width child"
5482 .to_string(),
5483 ),
5484 },
5485 DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
5486 None
5487 }
5488 other => Some(format!(
5489 "Full-zip does not support value block type {}",
5490 other.name()
5491 )),
5492 };
5493 match requested_encoding.as_deref() {
5494 Some(STRUCTURAL_ENCODING_FULLZIP) => {
5495 if let Some(reason) = fullzip_error {
5496 return Err(Error::invalid_input_source(reason.into()));
5497 }
5498 return Self::encode_full_zip(
5499 column_idx,
5500 &field,
5501 compression_strategy.as_ref(),
5502 data_block,
5503 repdef,
5504 row_number,
5505 num_rows,
5506 );
5507 }
5508 Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
5509 if requested_encoding.is_none() && fullzip_error.is_none() {
5510 log::debug!(
5511 "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
5512 column_idx,
5513 num_values
5514 );
5515 return Self::encode_full_zip(
5516 column_idx,
5517 &field,
5518 compression_strategy.as_ref(),
5519 data_block,
5520 repdef,
5521 row_number,
5522 num_rows,
5523 );
5524 }
5525 return Err(Error::invalid_input_source(
5526 format!(
5527 "Mini-block cannot encode {} rep/def levels in one top-level row. \
5528 This usually means the row contains too much nested structure \
5529 for the current layout.",
5530 num_levels
5531 )
5532 .into(),
5533 ));
5534 }
5535 _ => {}
5536 }
5537 }
5538
5539 let requires_full_zip_packed_struct =
5540 if let DataBlock::Struct(ref struct_data_block) = data_block {
5541 struct_data_block.has_variable_width_child()
5542 } else {
5543 false
5544 };
5545
5546 if requires_full_zip_packed_struct {
5547 log::debug!(
5548 "Encoding column {} with {} items using full-zip packed struct layout",
5549 column_idx,
5550 num_values
5551 );
5552 return Self::encode_full_zip(
5553 column_idx,
5554 &field,
5555 compression_strategy.as_ref(),
5556 data_block,
5557 repdef,
5558 row_number,
5559 num_rows,
5560 );
5561 }
5562
5563 if let DataBlock::Dictionary(dict) = data_block {
5564 log::debug!(
5565 "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
5566 column_idx,
5567 num_values
5568 );
5569 let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
5570 indices_data_block.compute_stat();
5575 return Self::encode_miniblock(
5576 column_idx,
5577 &field,
5578 compression_strategy.as_ref(),
5579 indices_data_block,
5580 repdef,
5581 row_number,
5582 Some(dictionary_data_block),
5583 num_rows,
5584 support_large_chunk,
5585 );
5586 }
5587
5588 let dict_result = Self::should_dictionary_encode(&data_block, &field, version).and_then(|budget| {
5591 log::debug!(
5592 "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
5593 column_idx,
5594 num_values
5595 );
5596 dict::dictionary_encode(
5597 &data_block,
5598 budget.max_dict_entries,
5599 budget.max_encoded_size,
5600 )
5601 });
5602
5603 if let Some((indices_data_block, dictionary_data_block)) = dict_result {
5604 Self::encode_miniblock(
5605 column_idx,
5606 &field,
5607 compression_strategy.as_ref(),
5608 indices_data_block,
5609 repdef,
5610 row_number,
5611 Some(dictionary_data_block),
5612 num_rows,
5613 support_large_chunk,
5614 )
5615 } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
5616 log::debug!(
5617 "Encoding column {} with {} items using mini-block layout",
5618 column_idx,
5619 num_values
5620 );
5621 Self::encode_miniblock(
5622 column_idx,
5623 &field,
5624 compression_strategy.as_ref(),
5625 data_block,
5626 repdef,
5627 row_number,
5628 None,
5629 num_rows,
5630 support_large_chunk,
5631 )
5632 } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
5633 log::debug!(
5634 "Encoding column {} with {} items using full-zip layout",
5635 column_idx,
5636 num_values
5637 );
5638 Self::encode_full_zip(
5639 column_idx,
5640 &field,
5641 compression_strategy.as_ref(),
5642 data_block,
5643 repdef,
5644 row_number,
5645 num_rows,
5646 )
5647 } else {
5648 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()))
5649 }
5650 }
5651
5652 fn do_flush(
5654 &mut self,
5655 arrays: Vec<ArrayRef>,
5656 repdefs: Vec<RepDefBuilder>,
5657 row_number: u64,
5658 num_rows: u64,
5659 ) -> Result<Vec<EncodeTask>> {
5660 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5661 let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
5662 let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
5663 let (repdef, structural_plan) = RepDefBuilder::serialize_with_structural_plan(
5664 repdefs,
5665 miniblock::max_repdef_levels_per_chunk,
5666 num_rows,
5667 num_values,
5668 )?;
5669 let pages = Self::split_structural_pages_for_miniblock_budget(
5670 arrays,
5671 repdef,
5672 structural_plan,
5673 row_number,
5674 num_rows,
5675 )?;
5676
5677 let mut tasks = Vec::with_capacity(pages.len());
5678 let ctx = PrimitiveEncodeContext {
5679 column_idx: self.column_index,
5680 field: self.field.clone(),
5681 compression_strategy: self.compression_strategy.clone(),
5682 encoding_metadata: self.encoding_metadata.clone(),
5683 support_large_chunk: self.support_large_chunk,
5684 version: self.version,
5685 is_simple_validity,
5686 has_repdef_info,
5687 };
5688 for page in pages {
5689 let ctx = ctx.clone();
5690 let task = spawn_cpu(move || Self::encode_page(ctx, page)).boxed();
5691 tasks.push(task);
5692 }
5693 Ok(tasks)
5694 }
5695
5696 fn extract_validity_buf(
5697 array: Arc<dyn Array>,
5698 repdef: &mut RepDefBuilder,
5699 keep_original_array: bool,
5700 ) -> Result<Arc<dyn Array>> {
5701 if let Some(validity) = array.nulls() {
5702 if keep_original_array {
5703 repdef.add_validity_bitmap(validity.clone());
5704 } else {
5705 repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
5706 }
5707 let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
5708 Ok(make_array(data_no_nulls))
5709 } else {
5710 repdef.add_no_null(array.len());
5711 Ok(array)
5712 }
5713 }
5714
5715 fn extract_validity(
5716 mut array: Arc<dyn Array>,
5717 repdef: &mut RepDefBuilder,
5718 keep_original_array: bool,
5719 ) -> Result<Arc<dyn Array>> {
5720 match array.data_type() {
5721 DataType::Null => {
5722 repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
5723 Ok(array)
5724 }
5725 DataType::Dictionary(_, _) => {
5726 array = dict::normalize_dict_nulls(array)?;
5727 Self::extract_validity_buf(array, repdef, keep_original_array)
5728 }
5729 _ => Self::extract_validity_buf(array, repdef, keep_original_array),
5738 }
5739 }
5740}
5741
5742impl FieldEncoder for PrimitiveStructuralEncoder {
5743 fn maybe_encode(
5745 &mut self,
5746 array: ArrayRef,
5747 _external_buffers: &mut OutOfLineBuffers,
5748 mut repdef: RepDefBuilder,
5749 row_number: u64,
5750 num_rows: u64,
5751 ) -> Result<Vec<EncodeTask>> {
5752 let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
5753 self.accumulated_repdefs.push(repdef);
5754
5755 if let Some((arrays, row_number, num_rows)) =
5756 self.accumulation_queue.insert(array, row_number, num_rows)
5757 {
5758 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5759 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5760 } else {
5761 Ok(vec![])
5762 }
5763 }
5764
5765 fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
5767 if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
5768 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5769 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5770 } else {
5771 Ok(vec![])
5772 }
5773 }
5774
5775 fn num_columns(&self) -> u32 {
5776 1
5777 }
5778
5779 fn finish(
5780 &mut self,
5781 _external_buffers: &mut OutOfLineBuffers,
5782 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
5783 std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
5784 }
5785}
5786
5787#[cfg(test)]
5788#[allow(clippy::single_range_in_vec_init)]
5789mod tests {
5790 use super::{
5791 ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
5792 FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource,
5793 FullZipRepIndexDetails, FullZipScheduler, MiniBlockChunk, MiniBlockCompressed,
5794 MiniBlockRepIndex, PerValueDecompressor, PreambleAction, StructuralPageScheduler,
5795 VariableFullZipDecoder,
5796 };
5797 use crate::buffer::LanceBuffer;
5798 use crate::compression::DefaultDecompressionStrategy;
5799 use crate::constants::{
5800 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
5801 DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
5802 STRUCTURAL_ENCODING_MINIBLOCK,
5803 };
5804 use crate::data::BlockInfo;
5805 use crate::decoder::PageEncoding;
5806 use crate::encodings::logical::primitive::{
5807 ChunkDrainInstructions, PrimitiveStructuralEncoder,
5808 };
5809 use crate::format::ProtobufUtils21;
5810 use crate::format::pb21;
5811 use crate::format::pb21::compressive_encoding::Compression;
5812 use crate::repdef::build_control_word_iterator;
5813 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
5814 use crate::version::LanceFileVersion;
5815 use arrow_array::{ArrayRef, Int8Array, StringArray};
5816 use arrow_schema::DataType;
5817 use std::collections::HashMap;
5818 use std::{collections::VecDeque, sync::Arc};
5819
5820 #[test]
5821 fn test_is_narrow() {
5822 let int8_array = Int8Array::from(vec![1, 2, 3]);
5823 let array_ref: ArrayRef = Arc::new(int8_array);
5824 let block = DataBlock::from_array(array_ref);
5825
5826 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5827
5828 let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
5829 let block = DataBlock::from_array(string_array);
5830 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5831
5832 let string_array = StringArray::from(vec![
5833 Some("hello world".repeat(100)),
5834 Some("world".to_string()),
5835 ]);
5836 let block = DataBlock::from_array(string_array);
5837 assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
5838 }
5839
5840 #[test]
5841 fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
5842 let fixed = FixedWidthDataBlock {
5843 data: LanceBuffer::from(vec![0_u8]),
5844 bits_per_value: 1,
5845 num_values: 8,
5846 block_info: BlockInfo::new(),
5847 };
5848 let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
5849
5850 let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
5851 else {
5852 panic!("expected full-zip to reject 1-bit fixed-width values");
5853 };
5854 assert!(
5855 err.to_string().contains("byte aligned"),
5856 "unexpected error: {err}"
5857 );
5858 }
5859
5860 #[test]
5861 fn test_map_range() {
5862 let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
5865 let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
5866 let max_visible_def = 0;
5867 let total_items = 8;
5868 let max_rep = 1;
5869
5870 let check = |range, expected_item_range, expected_level_range| {
5871 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5872 range,
5873 rep.as_ref(),
5874 def.as_ref(),
5875 max_rep,
5876 max_visible_def,
5877 total_items,
5878 PreambleAction::Absent,
5879 );
5880 assert_eq!(item_range, expected_item_range);
5881 assert_eq!(level_range, expected_level_range);
5882 };
5883
5884 check(0..1, 0..3, 0..3);
5885 check(1..2, 3..5, 3..5);
5886 check(2..3, 5..5, 5..6);
5887 check(3..4, 5..8, 6..9);
5888 check(0..2, 0..5, 0..5);
5889 check(1..3, 3..5, 3..6);
5890 check(2..4, 5..8, 5..9);
5891 check(0..3, 0..5, 0..6);
5892 check(1..4, 3..8, 3..9);
5893 check(0..4, 0..8, 0..9);
5894
5895 let rep = Some(vec![1, 1, 0, 1]);
5898 let def = Some(vec![1, 0, 0, 0]);
5899 let max_visible_def = 0;
5900 let total_items = 3;
5901
5902 let check = |range, expected_item_range, expected_level_range| {
5903 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5904 range,
5905 rep.as_ref(),
5906 def.as_ref(),
5907 max_rep,
5908 max_visible_def,
5909 total_items,
5910 PreambleAction::Absent,
5911 );
5912 assert_eq!(item_range, expected_item_range);
5913 assert_eq!(level_range, expected_level_range);
5914 };
5915
5916 check(0..1, 0..0, 0..1);
5917 check(1..2, 0..2, 1..3);
5918 check(2..3, 2..3, 3..4);
5919 check(0..2, 0..2, 0..3);
5920 check(1..3, 0..3, 1..4);
5921 check(0..3, 0..3, 0..4);
5922
5923 let rep = Some(vec![1, 1, 0, 1]);
5926 let def = Some(vec![0, 0, 0, 1]);
5927 let max_visible_def = 0;
5928 let total_items = 3;
5929
5930 let check = |range, expected_item_range, expected_level_range| {
5931 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5932 range,
5933 rep.as_ref(),
5934 def.as_ref(),
5935 max_rep,
5936 max_visible_def,
5937 total_items,
5938 PreambleAction::Absent,
5939 );
5940 assert_eq!(item_range, expected_item_range);
5941 assert_eq!(level_range, expected_level_range);
5942 };
5943
5944 check(0..1, 0..1, 0..1);
5945 check(1..2, 1..3, 1..3);
5946 check(2..3, 3..3, 3..4);
5947 check(0..2, 0..3, 0..3);
5948 check(1..3, 1..3, 1..4);
5949 check(0..3, 0..3, 0..4);
5950
5951 let rep = Some(vec![1, 0, 1, 0, 1, 0]);
5954 let def: Option<&[u16]> = None;
5955 let max_visible_def = 0;
5956 let total_items = 6;
5957
5958 let check = |range, expected_item_range, expected_level_range| {
5959 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5960 range,
5961 rep.as_ref(),
5962 def.as_ref(),
5963 max_rep,
5964 max_visible_def,
5965 total_items,
5966 PreambleAction::Absent,
5967 );
5968 assert_eq!(item_range, expected_item_range);
5969 assert_eq!(level_range, expected_level_range);
5970 };
5971
5972 check(0..1, 0..2, 0..2);
5973 check(1..2, 2..4, 2..4);
5974 check(2..3, 4..6, 4..6);
5975 check(0..2, 0..4, 0..4);
5976 check(1..3, 2..6, 2..6);
5977 check(0..3, 0..6, 0..6);
5978
5979 let rep: Option<&[u16]> = None;
5982 let def = Some(vec![0, 0, 1, 0]);
5983 let max_visible_def = 1;
5984 let total_items = 4;
5985
5986 let check = |range, expected_item_range, expected_level_range| {
5987 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5988 range,
5989 rep.as_ref(),
5990 def.as_ref(),
5991 max_rep,
5992 max_visible_def,
5993 total_items,
5994 PreambleAction::Absent,
5995 );
5996 assert_eq!(item_range, expected_item_range);
5997 assert_eq!(level_range, expected_level_range);
5998 };
5999
6000 check(0..1, 0..1, 0..1);
6001 check(1..2, 1..2, 1..2);
6002 check(2..3, 2..3, 2..3);
6003 check(0..2, 0..2, 0..2);
6004 check(1..3, 1..3, 1..3);
6005 check(0..3, 0..3, 0..3);
6006
6007 let rep = Some(vec![0, 1, 0, 1]);
6012 let def = Some(vec![0, 0, 0, 1]);
6013 let max_visible_def = 0;
6014 let total_items = 3;
6015
6016 let check = |range, expected_item_range, expected_level_range| {
6017 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6018 range,
6019 rep.as_ref(),
6020 def.as_ref(),
6021 max_rep,
6022 max_visible_def,
6023 total_items,
6024 PreambleAction::Take,
6025 );
6026 assert_eq!(item_range, expected_item_range);
6027 assert_eq!(level_range, expected_level_range);
6028 };
6029
6030 check(0..1, 0..3, 0..3);
6032 check(0..2, 0..3, 0..4);
6033
6034 let check = |range, expected_item_range, expected_level_range| {
6035 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6036 range,
6037 rep.as_ref(),
6038 def.as_ref(),
6039 max_rep,
6040 max_visible_def,
6041 total_items,
6042 PreambleAction::Skip,
6043 );
6044 assert_eq!(item_range, expected_item_range);
6045 assert_eq!(level_range, expected_level_range);
6046 };
6047
6048 check(0..1, 1..3, 1..3);
6049 check(1..2, 3..3, 3..4);
6050 check(0..2, 1..3, 1..4);
6051
6052 let rep = Some(vec![0, 1, 1, 0]);
6057 let def = Some(vec![0, 1, 0, 0]);
6058 let max_visible_def = 0;
6059 let total_items = 4;
6060
6061 let check = |range, expected_item_range, expected_level_range| {
6062 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6063 range,
6064 rep.as_ref(),
6065 def.as_ref(),
6066 max_rep,
6067 max_visible_def,
6068 total_items,
6069 PreambleAction::Take,
6070 );
6071 assert_eq!(item_range, expected_item_range);
6072 assert_eq!(level_range, expected_level_range);
6073 };
6074
6075 check(0..1, 0..1, 0..2);
6077 check(0..2, 0..3, 0..4);
6078
6079 let check = |range, expected_item_range, expected_level_range| {
6080 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6081 range,
6082 rep.as_ref(),
6083 def.as_ref(),
6084 max_rep,
6085 max_visible_def,
6086 total_items,
6087 PreambleAction::Skip,
6088 );
6089 assert_eq!(item_range, expected_item_range);
6090 assert_eq!(level_range, expected_level_range);
6091 };
6092
6093 check(0..1, 1..1, 1..2);
6095 check(1..2, 1..3, 2..4);
6096 check(0..2, 1..3, 1..4);
6097
6098 let rep = Some(vec![0, 1, 0, 1]);
6101 let def: Option<Vec<u16>> = None;
6102 let max_visible_def = 0;
6103 let total_items = 4;
6104
6105 let check = |range, expected_item_range, expected_level_range| {
6106 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6107 range,
6108 rep.as_ref(),
6109 def.as_ref(),
6110 max_rep,
6111 max_visible_def,
6112 total_items,
6113 PreambleAction::Take,
6114 );
6115 assert_eq!(item_range, expected_item_range);
6116 assert_eq!(level_range, expected_level_range);
6117 };
6118
6119 check(0..1, 0..3, 0..3);
6121 check(0..2, 0..4, 0..4);
6122
6123 let check = |range, expected_item_range, expected_level_range| {
6124 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6125 range,
6126 rep.as_ref(),
6127 def.as_ref(),
6128 max_rep,
6129 max_visible_def,
6130 total_items,
6131 PreambleAction::Skip,
6132 );
6133 assert_eq!(item_range, expected_item_range);
6134 assert_eq!(level_range, expected_level_range);
6135 };
6136
6137 check(0..1, 1..3, 1..3);
6138 check(1..2, 3..4, 3..4);
6139 check(0..2, 1..4, 1..4);
6140
6141 let rep = Some(vec![2, 1, 2, 0, 1, 2]);
6145 let def = Some(vec![0, 1, 2, 0, 0, 0]);
6146 let max_rep = 2;
6147 let max_visible_def = 0;
6148 let total_items = 4;
6149
6150 let check = |range, expected_item_range, expected_level_range| {
6151 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6152 range,
6153 rep.as_ref(),
6154 def.as_ref(),
6155 max_rep,
6156 max_visible_def,
6157 total_items,
6158 PreambleAction::Absent,
6159 );
6160 assert_eq!(item_range, expected_item_range);
6161 assert_eq!(level_range, expected_level_range);
6162 };
6163
6164 check(0..3, 0..4, 0..6);
6165 check(0..1, 0..1, 0..2);
6166 check(1..2, 1..3, 2..5);
6167 check(2..3, 3..4, 5..6);
6168
6169 let rep = Some(vec![0, 0, 1, 0, 1, 1]);
6171 let def = Some(vec![0, 1, 0, 0, 0, 0]);
6172 let max_rep = 1;
6173 let max_visible_def = 0;
6174 let total_items = 5;
6175
6176 let check = |range, expected_item_range, expected_level_range| {
6177 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6178 range,
6179 rep.as_ref(),
6180 def.as_ref(),
6181 max_rep,
6182 max_visible_def,
6183 total_items,
6184 PreambleAction::Take,
6185 );
6186 assert_eq!(item_range, expected_item_range);
6187 assert_eq!(level_range, expected_level_range);
6188 };
6189
6190 check(0..0, 0..1, 0..2);
6191 check(0..1, 0..3, 0..4);
6192 check(0..2, 0..4, 0..5);
6193
6194 let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
6197 let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
6198 let max_rep = 1;
6199 let max_visible_def = 0;
6200 let total_items = 5;
6201
6202 let check = |range, expected_item_range, expected_level_range| {
6203 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6204 range,
6205 rep.as_ref(),
6206 def.as_ref(),
6207 max_rep,
6208 max_visible_def,
6209 total_items,
6210 PreambleAction::Skip,
6211 );
6212 assert_eq!(item_range, expected_item_range);
6213 assert_eq!(level_range, expected_level_range);
6214 };
6215
6216 check(2..3, 2..4, 5..7);
6217 }
6218
6219 #[test]
6220 fn test_slice_batch_data_and_rebase_offsets_u32() {
6221 let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
6222 let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
6223
6224 let (sliced_data, normalized_offsets) =
6225 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6226 .unwrap();
6227
6228 assert_eq!(sliced_data.as_ref(), b"6789ab");
6229 let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
6230 assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
6231 }
6232
6233 #[test]
6234 fn test_slice_batch_data_and_rebase_offsets_u64() {
6235 let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
6236 let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
6237
6238 let (sliced_data, normalized_offsets) =
6239 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
6240 .unwrap();
6241
6242 assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
6243 let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
6244 assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
6245 }
6246
6247 #[test]
6248 fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
6249 let data = LanceBuffer::copy_slice(b"abcd");
6250 let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
6251
6252 let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6253 .expect_err("offset end before start should error");
6254 assert!(err.to_string().contains("less than base"));
6255 }
6256
6257 #[test]
6258 fn test_schedule_instructions() {
6259 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6261 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6262 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6263
6264 let check = |user_ranges, expected_instructions| {
6265 let instructions =
6266 ChunkInstructions::schedule_instructions(&repetition_index, user_ranges);
6267 assert_eq!(instructions, expected_instructions);
6268 };
6269
6270 let expected_take_all = vec![
6272 ChunkInstructions {
6273 chunk_idx: 0,
6274 preamble: PreambleAction::Absent,
6275 rows_to_skip: 0,
6276 rows_to_take: 6,
6277 take_trailer: true,
6278 },
6279 ChunkInstructions {
6280 chunk_idx: 1,
6281 preamble: PreambleAction::Take,
6282 rows_to_skip: 0,
6283 rows_to_take: 2,
6284 take_trailer: false,
6285 },
6286 ChunkInstructions {
6287 chunk_idx: 2,
6288 preamble: PreambleAction::Absent,
6289 rows_to_skip: 0,
6290 rows_to_take: 5,
6291 take_trailer: true,
6292 },
6293 ChunkInstructions {
6294 chunk_idx: 3,
6295 preamble: PreambleAction::Take,
6296 rows_to_skip: 0,
6297 rows_to_take: 1,
6298 take_trailer: false,
6299 },
6300 ];
6301
6302 check(&[0..14], expected_take_all.clone());
6304
6305 check(
6307 &[
6308 0..1,
6309 1..2,
6310 2..3,
6311 3..4,
6312 4..5,
6313 5..6,
6314 6..7,
6315 7..8,
6316 8..9,
6317 9..10,
6318 10..11,
6319 11..12,
6320 12..13,
6321 13..14,
6322 ],
6323 expected_take_all,
6324 );
6325
6326 check(
6330 &[0..1, 3..4],
6331 vec![
6332 ChunkInstructions {
6333 chunk_idx: 0,
6334 preamble: PreambleAction::Absent,
6335 rows_to_skip: 0,
6336 rows_to_take: 1,
6337 take_trailer: false,
6338 },
6339 ChunkInstructions {
6340 chunk_idx: 0,
6341 preamble: PreambleAction::Absent,
6342 rows_to_skip: 3,
6343 rows_to_take: 1,
6344 take_trailer: false,
6345 },
6346 ],
6347 );
6348
6349 check(
6351 &[5..6],
6352 vec![
6353 ChunkInstructions {
6354 chunk_idx: 0,
6355 preamble: PreambleAction::Absent,
6356 rows_to_skip: 5,
6357 rows_to_take: 1,
6358 take_trailer: true,
6359 },
6360 ChunkInstructions {
6361 chunk_idx: 1,
6362 preamble: PreambleAction::Take,
6363 rows_to_skip: 0,
6364 rows_to_take: 0,
6365 take_trailer: false,
6366 },
6367 ],
6368 );
6369
6370 check(
6372 &[7..10],
6373 vec![
6374 ChunkInstructions {
6375 chunk_idx: 1,
6376 preamble: PreambleAction::Skip,
6377 rows_to_skip: 1,
6378 rows_to_take: 1,
6379 take_trailer: false,
6380 },
6381 ChunkInstructions {
6382 chunk_idx: 2,
6383 preamble: PreambleAction::Absent,
6384 rows_to_skip: 0,
6385 rows_to_take: 2,
6386 take_trailer: false,
6387 },
6388 ],
6389 );
6390 }
6391
6392 #[test]
6393 fn test_drain_instructions() {
6394 fn drain_from_instructions(
6395 instructions: &mut VecDeque<ChunkInstructions>,
6396 mut rows_desired: u64,
6397 need_preamble: &mut bool,
6398 skip_in_chunk: &mut u64,
6399 ) -> Vec<ChunkDrainInstructions> {
6400 let mut drain_instructions = Vec::with_capacity(instructions.len());
6402 while rows_desired > 0 || *need_preamble {
6403 let (next_instructions, consumed_chunk) = instructions
6404 .front()
6405 .unwrap()
6406 .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
6407 if consumed_chunk {
6408 instructions.pop_front();
6409 }
6410 drain_instructions.push(next_instructions);
6411 }
6412 drain_instructions
6413 }
6414
6415 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6417 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6418 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6419 let user_ranges = vec![1..7, 10..14];
6420
6421 let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6423
6424 let mut to_drain = VecDeque::from(scheduled.clone());
6425
6426 let mut need_preamble = false;
6429 let mut skip_in_chunk = 0;
6430
6431 let next_batch =
6432 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6433
6434 assert!(!need_preamble);
6435 assert_eq!(skip_in_chunk, 4);
6436 assert_eq!(
6437 next_batch,
6438 vec![ChunkDrainInstructions {
6439 chunk_instructions: scheduled[0].clone(),
6440 rows_to_take: 4,
6441 rows_to_skip: 0,
6442 preamble_action: PreambleAction::Absent,
6443 }]
6444 );
6445
6446 let next_batch =
6447 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6448
6449 assert!(!need_preamble);
6450 assert_eq!(skip_in_chunk, 2);
6451
6452 assert_eq!(
6453 next_batch,
6454 vec![
6455 ChunkDrainInstructions {
6456 chunk_instructions: scheduled[0].clone(),
6457 rows_to_take: 1,
6458 rows_to_skip: 4,
6459 preamble_action: PreambleAction::Absent,
6460 },
6461 ChunkDrainInstructions {
6462 chunk_instructions: scheduled[1].clone(),
6463 rows_to_take: 1,
6464 rows_to_skip: 0,
6465 preamble_action: PreambleAction::Take,
6466 },
6467 ChunkDrainInstructions {
6468 chunk_instructions: scheduled[2].clone(),
6469 rows_to_take: 2,
6470 rows_to_skip: 0,
6471 preamble_action: PreambleAction::Absent,
6472 }
6473 ]
6474 );
6475
6476 let next_batch =
6477 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6478
6479 assert!(!need_preamble);
6480 assert_eq!(skip_in_chunk, 0);
6481
6482 assert_eq!(
6483 next_batch,
6484 vec![
6485 ChunkDrainInstructions {
6486 chunk_instructions: scheduled[2].clone(),
6487 rows_to_take: 1,
6488 rows_to_skip: 2,
6489 preamble_action: PreambleAction::Absent,
6490 },
6491 ChunkDrainInstructions {
6492 chunk_instructions: scheduled[3].clone(),
6493 rows_to_take: 1,
6494 rows_to_skip: 0,
6495 preamble_action: PreambleAction::Take,
6496 },
6497 ]
6498 );
6499
6500 let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
6502 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6503 let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6504 let user_ranges = vec![0..28];
6505
6506 let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6508
6509 let mut to_drain = VecDeque::from(scheduled.clone());
6510
6511 let mut need_preamble = false;
6514 let mut skip_in_chunk = 0;
6515
6516 let next_batch =
6517 drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
6518
6519 assert_eq!(
6520 next_batch,
6521 vec![
6522 ChunkDrainInstructions {
6523 chunk_instructions: scheduled[0].clone(),
6524 rows_to_take: 6,
6525 rows_to_skip: 0,
6526 preamble_action: PreambleAction::Absent,
6527 },
6528 ChunkDrainInstructions {
6529 chunk_instructions: scheduled[1].clone(),
6530 rows_to_take: 1,
6531 rows_to_skip: 0,
6532 preamble_action: PreambleAction::Take,
6533 },
6534 ]
6535 );
6536
6537 assert!(!need_preamble);
6538 assert_eq!(skip_in_chunk, 1);
6539
6540 let next_batch =
6543 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6544
6545 assert_eq!(
6546 next_batch,
6547 vec![
6548 ChunkDrainInstructions {
6549 chunk_instructions: scheduled[1].clone(),
6550 rows_to_take: 2,
6551 rows_to_skip: 1,
6552 preamble_action: PreambleAction::Skip,
6553 },
6554 ChunkDrainInstructions {
6555 chunk_instructions: scheduled[2].clone(),
6556 rows_to_take: 0,
6557 rows_to_skip: 0,
6558 preamble_action: PreambleAction::Take,
6559 },
6560 ]
6561 );
6562
6563 assert!(!need_preamble);
6564 assert_eq!(skip_in_chunk, 0);
6565 }
6566
6567 #[tokio::test]
6568 async fn test_fullzip_initialize_is_lazy() {
6569 use futures::{FutureExt, future::BoxFuture};
6570 use std::ops::Range;
6571 use std::sync::Mutex;
6572
6573 #[derive(Debug, Clone)]
6574 struct RecordingScheduler {
6575 data: bytes::Bytes,
6576 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6577 }
6578
6579 impl RecordingScheduler {
6580 fn new(data: bytes::Bytes) -> Self {
6581 Self {
6582 data,
6583 requests: Arc::new(Mutex::new(Vec::new())),
6584 }
6585 }
6586
6587 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6588 self.requests.lock().unwrap().clone()
6589 }
6590 }
6591
6592 impl crate::EncodingsIo for RecordingScheduler {
6593 fn submit_request(
6594 &self,
6595 ranges: Vec<Range<u64>>,
6596 _priority: u64,
6597 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6598 self.requests.lock().unwrap().push(ranges.clone());
6599 let data = ranges
6600 .into_iter()
6601 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6602 .collect::<Vec<_>>();
6603 std::future::ready(Ok(data)).boxed()
6604 }
6605 }
6606
6607 #[derive(Debug)]
6608 struct TestFixedDecompressor;
6609
6610 impl FixedPerValueDecompressor for TestFixedDecompressor {
6611 fn decompress(
6612 &self,
6613 _data: FixedWidthDataBlock,
6614 _num_rows: u64,
6615 ) -> crate::Result<DataBlock> {
6616 unimplemented!("Test decompressor")
6617 }
6618
6619 fn bits_per_value(&self) -> u64 {
6620 32
6621 }
6622 }
6623
6624 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
6625 0;
6626 16 * 1024
6627 ])));
6628 let mut scheduler = FullZipScheduler {
6629 data_buf_position: 0,
6630 data_buf_size: 4096,
6631 rep_index: Some(FullZipRepIndexDetails {
6632 buf_position: 1000,
6633 bytes_per_value: 4,
6634 }),
6635 priority: 0,
6636 rows_in_page: 100,
6637 bits_per_offset: 32,
6638 details: Arc::new(FullZipDecodeDetails {
6639 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6640 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6641 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6642 max_rep: 0,
6643 max_visible_def: 0,
6644 }),
6645 cached_state: None,
6646 enable_cache: false,
6647 };
6648
6649 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6650 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6651
6652 assert!(
6653 cached_data
6654 .as_arc_any()
6655 .downcast_ref::<super::NoCachedPageData>()
6656 .is_some(),
6657 "FullZip initialize should not eagerly load repetition index data"
6658 );
6659 assert!(scheduler.cached_state.is_none());
6660 assert!(
6661 io.requests().is_empty(),
6662 "FullZip initialize should not issue any I/O"
6663 );
6664 }
6665
6666 #[tokio::test]
6667 async fn test_fullzip_read_source_slices_prefetched_page() {
6668 let page_start = 200_u64;
6669 let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
6670 let source = FullZipReadSource::PrefetchedPage {
6671 base_offset: page_start,
6672 data: page_data,
6673 };
6674 let ranges = vec![
6675 page_start..(page_start + 3),
6676 (page_start + 4)..(page_start + 8),
6677 ];
6678 let mut data = source.fetch(&ranges, 0).await.unwrap();
6679 assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
6680 assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
6681 }
6682
6683 #[tokio::test]
6684 async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
6685 use futures::{FutureExt, future::BoxFuture};
6686 use std::ops::Range;
6687 use std::sync::Mutex;
6688
6689 #[derive(Debug, Clone)]
6690 struct RecordingScheduler {
6691 data: bytes::Bytes,
6692 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6693 }
6694
6695 impl RecordingScheduler {
6696 fn new(data: bytes::Bytes) -> Self {
6697 Self {
6698 data,
6699 requests: Arc::new(Mutex::new(Vec::new())),
6700 }
6701 }
6702
6703 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6704 self.requests.lock().unwrap().clone()
6705 }
6706 }
6707
6708 impl crate::EncodingsIo for RecordingScheduler {
6709 fn submit_request(
6710 &self,
6711 ranges: Vec<Range<u64>>,
6712 _priority: u64,
6713 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6714 self.requests.lock().unwrap().push(ranges.clone());
6715 let data = ranges
6716 .into_iter()
6717 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6718 .collect::<Vec<_>>();
6719 std::future::ready(Ok(data)).boxed()
6720 }
6721 }
6722
6723 #[derive(Debug)]
6724 struct TestFixedDecompressor;
6725
6726 impl FixedPerValueDecompressor for TestFixedDecompressor {
6727 fn decompress(
6728 &self,
6729 _data: FixedWidthDataBlock,
6730 _num_rows: u64,
6731 ) -> crate::Result<DataBlock> {
6732 unimplemented!("Test decompressor")
6733 }
6734
6735 fn bits_per_value(&self) -> u64 {
6736 32
6737 }
6738 }
6739
6740 let rows_in_page = 100_u64;
6741 let bytes_per_value = 4_u64;
6742 let rep_start = 1000_u64;
6743 let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
6744 let mut data = vec![0_u8; 16 * 1024];
6745 data[rep_start as usize..rep_start as usize + rep_size].fill(7);
6746 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
6747
6748 let mut scheduler = FullZipScheduler {
6749 data_buf_position: 0,
6750 data_buf_size: 4096,
6751 rep_index: Some(FullZipRepIndexDetails {
6752 buf_position: rep_start,
6753 bytes_per_value,
6754 }),
6755 priority: 0,
6756 rows_in_page,
6757 bits_per_offset: 32,
6758 details: Arc::new(FullZipDecodeDetails {
6759 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6760 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6761 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6762 max_rep: 0,
6763 max_visible_def: 0,
6764 }),
6765 cached_state: None,
6766 enable_cache: true,
6767 };
6768
6769 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6770 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6771 assert!(
6772 cached_data
6773 .as_arc_any()
6774 .downcast_ref::<FullZipCacheableState>()
6775 .is_some()
6776 );
6777 assert!(scheduler.cached_state.is_some());
6778 assert_eq!(
6779 io.requests(),
6780 vec![vec![
6781 rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
6782 ]]
6783 );
6784 }
6785
6786 #[tokio::test]
6787 async fn test_fullzip_full_page_bypasses_rep_index_io() {
6788 use futures::{FutureExt, future::BoxFuture};
6789 use std::ops::Range;
6790 use std::sync::Mutex;
6791
6792 #[derive(Debug, Clone)]
6793 struct RecordingScheduler {
6794 data: bytes::Bytes,
6795 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6796 }
6797
6798 impl RecordingScheduler {
6799 fn new(data: bytes::Bytes) -> Self {
6800 Self {
6801 data,
6802 requests: Arc::new(Mutex::new(Vec::new())),
6803 }
6804 }
6805
6806 fn requests(&self) -> Vec<Vec<Range<u64>>> {
6807 self.requests.lock().unwrap().clone()
6808 }
6809 }
6810
6811 impl crate::EncodingsIo for RecordingScheduler {
6812 fn submit_request(
6813 &self,
6814 ranges: Vec<Range<u64>>,
6815 _priority: u64,
6816 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6817 self.requests.lock().unwrap().push(ranges.clone());
6818 let data = ranges
6819 .into_iter()
6820 .map(|range| self.data.slice(range.start as usize..range.end as usize))
6821 .collect::<Vec<_>>();
6822 std::future::ready(Ok(data)).boxed()
6823 }
6824 }
6825
6826 #[derive(Debug)]
6827 struct TestFixedDecompressor;
6828
6829 impl FixedPerValueDecompressor for TestFixedDecompressor {
6830 fn decompress(
6831 &self,
6832 _data: FixedWidthDataBlock,
6833 _num_rows: u64,
6834 ) -> crate::Result<DataBlock> {
6835 unimplemented!("Test decompressor")
6836 }
6837
6838 fn bits_per_value(&self) -> u64 {
6839 32
6840 }
6841 }
6842
6843 let rows_in_page = 100_u64;
6844 let data_start = 256_u64;
6845 let data_size = 500_u64;
6846 let rep_start = 4096_u64;
6847 let bytes_per_value = 4_u64;
6848
6849 let mut bytes = vec![0_u8; 16 * 1024];
6850 for i in 0..=rows_in_page {
6851 let offset = (i * 5) as u32;
6852 let pos = rep_start as usize + (i * bytes_per_value) as usize;
6853 bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
6854 }
6855 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
6856
6857 let scheduler = FullZipScheduler {
6858 data_buf_position: data_start,
6859 data_buf_size: data_size,
6860 rep_index: Some(FullZipRepIndexDetails {
6861 buf_position: rep_start,
6862 bytes_per_value,
6863 }),
6864 priority: 0,
6865 rows_in_page,
6866 bits_per_offset: 32,
6867 details: Arc::new(FullZipDecodeDetails {
6868 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6869 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6870 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6871 max_rep: 0,
6872 max_visible_def: 0,
6873 }),
6874 cached_state: None,
6875 enable_cache: false,
6876 };
6877
6878 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6879 let tasks = scheduler
6880 .schedule_ranges_rep(
6881 &[0..rows_in_page],
6882 &io_dyn,
6883 FullZipRepIndexDetails {
6884 buf_position: rep_start,
6885 bytes_per_value,
6886 },
6887 )
6888 .unwrap();
6889
6890 let requests = io.requests();
6891 assert_eq!(requests.len(), 1);
6892 assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
6893
6894 let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
6895 let requests_after_await = io.requests();
6896 assert_eq!(
6897 requests_after_await.len(),
6898 1,
6899 "full page path should not issue rep-index I/O"
6900 );
6901 }
6902
6903 #[tokio::test]
6905 async fn test_fuzz_issue_4492_empty_rep_values() {
6906 use lance_datagen::{RowCount, Seed, array, gen_batch};
6907
6908 let seed = 1823859942947654717u64;
6909 let num_rows = 2741usize;
6910
6911 let batch_gen = gen_batch().with_seed(Seed::from(seed));
6913 let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
6914 let list_generator = array::rand_list_any(base_generator, false);
6915
6916 let batch = batch_gen
6917 .anon_col(list_generator)
6918 .into_batch_rows(RowCount::from(num_rows as u64))
6919 .unwrap();
6920
6921 let list_array = batch.column(0).clone();
6922
6923 let mut metadata = HashMap::new();
6925 metadata.insert(
6926 STRUCTURAL_ENCODING_META_KEY.to_string(),
6927 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6928 );
6929
6930 let test_cases = TestCases::default()
6931 .with_min_file_version(LanceFileVersion::V2_1)
6932 .with_batch_size(100)
6933 .with_range(0..num_rows.min(500) as u64)
6934 .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
6935
6936 check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
6937 }
6938
6939 async fn test_minichunk_size_helper(
6940 string_data: Vec<Option<String>>,
6941 minichunk_size: u64,
6942 file_version: LanceFileVersion,
6943 ) {
6944 use crate::constants::MINICHUNK_SIZE_META_KEY;
6945 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6946 use arrow_array::{ArrayRef, StringArray};
6947 use std::sync::Arc;
6948
6949 let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
6950
6951 let mut metadata = HashMap::new();
6952 metadata.insert(
6953 MINICHUNK_SIZE_META_KEY.to_string(),
6954 minichunk_size.to_string(),
6955 );
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(file_version)
6963 .with_batch_size(1000);
6964
6965 check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
6966 }
6967
6968 #[tokio::test]
6969 async fn test_minichunk_size_roundtrip() {
6970 let mut string_data = Vec::new();
6972 for i in 0..100 {
6973 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6974 }
6975 test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await;
6977 }
6978
6979 #[tokio::test]
6980 async fn test_minichunk_size_128kb_v2_2() {
6981 let mut string_data = Vec::new();
6983 for i in 0..10000 {
6985 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6986 }
6987 test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6988 }
6989
6990 #[tokio::test]
6991 async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
6992 let mut string_data = Vec::new();
6993 for i in 0..10000 {
6995 string_data.push(Some(format!("t_{}", i)));
6996 }
6997 test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6998 }
6999
7000 #[tokio::test]
7001 async fn test_large_dictionary_general_compression() {
7002 use arrow_array::{ArrayRef, StringArray};
7003 use std::collections::HashMap;
7004 use std::sync::Arc;
7005
7006 let unique_values: Vec<String> = (0..100)
7009 .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
7010 .collect();
7011
7012 let repeated_strings: Vec<_> = unique_values
7014 .iter()
7015 .cycle()
7016 .take(100_000)
7017 .map(|s| Some(s.as_str()))
7018 .collect();
7019
7020 let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
7021
7022 let test_cases = TestCases::default()
7024 .with_min_file_version(LanceFileVersion::V2_2)
7025 .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
7026 assert_eq!(cols.len(), 1);
7027 let col = &cols[0];
7028
7029 if let Some(PageEncoding::Structural(page_layout)) =
7031 &col.final_pages.first().map(|p| &p.description)
7032 && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
7033 &page_layout.layout
7034 && let Some(dictionary_encoding) = &mini_block.dictionary
7035 {
7036 match dictionary_encoding.compression.as_ref() {
7037 Some(Compression::General(general)) => {
7038 let compression = general.compression.as_ref().unwrap();
7040 assert!(
7041 compression.scheme()
7042 == pb21::CompressionScheme::CompressionAlgorithmLz4
7043 || compression.scheme()
7044 == pb21::CompressionScheme::CompressionAlgorithmZstd,
7045 "Expected LZ4 or Zstd compression for large dictionary"
7046 );
7047 }
7048 _ => panic!("Expected General compression for large dictionary"),
7049 }
7050 }
7051 }));
7052
7053 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
7054 }
7055
7056 fn dictionary_encoding_from_page(
7057 page: &crate::encoder::EncodedPage,
7058 ) -> &crate::format::pb21::CompressiveEncoding {
7059 let PageEncoding::Structural(layout) = &page.description else {
7060 panic!("Expected structural page encoding");
7061 };
7062 let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
7063 else {
7064 panic!("Expected mini-block layout");
7065 };
7066 layout
7067 .dictionary
7068 .as_ref()
7069 .unwrap_or_else(|| panic!("Expected dictionary encoding"))
7070 }
7071
7072 async fn encode_variable_dict_page(
7073 metadata: HashMap<String, String>,
7074 ) -> crate::encoder::EncodedPage {
7075 use arrow_array::types::Int32Type;
7076 use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
7077
7078 let values = Arc::new(StringArray::from(
7079 (0..128)
7080 .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
7081 .collect::<Vec<_>>(),
7082 )) as ArrayRef;
7083 let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
7084 let dict_array =
7085 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
7086
7087 let field = arrow_schema::Field::new(
7088 "dict_col",
7089 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
7090 false,
7091 )
7092 .with_metadata(metadata);
7093
7094 encode_first_page(field, dict_array, LanceFileVersion::V2_2).await
7095 }
7096
7097 async fn encode_auto_fixed_dict_page(
7098 metadata: HashMap<String, String>,
7099 ) -> crate::encoder::EncodedPage {
7100 use arrow_array::{ArrayRef, Decimal128Array};
7101
7102 let values = (0..20_000)
7104 .map(|i| match i % 3 {
7105 0 => 10_i128,
7106 1 => 20_i128,
7107 _ => 30_i128,
7108 })
7109 .collect::<Vec<_>>();
7110 let decimal = Decimal128Array::from_iter_values(values)
7111 .with_precision_and_scale(38, 0)
7112 .unwrap();
7113 let decimal = Arc::new(decimal) as ArrayRef;
7114
7115 let mut field_metadata = metadata;
7116 field_metadata.insert(
7118 "lance-encoding:dict-size-ratio".to_string(),
7119 "0.99".to_string(),
7120 );
7121 let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
7122 .with_metadata(field_metadata);
7123
7124 encode_first_page(field, decimal, LanceFileVersion::V2_2).await
7125 }
7126
7127 #[tokio::test]
7128 async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
7129 let page = encode_variable_dict_page(HashMap::new()).await;
7130 let dictionary_encoding = dictionary_encoding_from_page(&page);
7131 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7132 panic!("Expected General compression for dictionary values");
7133 };
7134 let compression = general.compression.as_ref().unwrap();
7135 assert_eq!(
7136 compression.scheme(),
7137 pb21::CompressionScheme::CompressionAlgorithmLz4
7138 );
7139 }
7140
7141 #[tokio::test]
7142 async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
7143 let page = encode_auto_fixed_dict_page(HashMap::new()).await;
7144 let dictionary_encoding = dictionary_encoding_from_page(&page);
7145 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7146 panic!("Expected General compression for dictionary values");
7147 };
7148 let compression = general.compression.as_ref().unwrap();
7149 assert_eq!(
7150 compression.scheme(),
7151 pb21::CompressionScheme::CompressionAlgorithmLz4
7152 );
7153 }
7154
7155 #[tokio::test]
7156 async fn test_dict_values_general_compression_zstd() {
7157 let mut metadata = HashMap::new();
7158 metadata.insert(
7159 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7160 "zstd".to_string(),
7161 );
7162 let page = encode_variable_dict_page(metadata).await;
7163 let dictionary_encoding = dictionary_encoding_from_page(&page);
7164 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
7165 panic!("Expected General compression for dictionary values");
7166 };
7167 let compression = general.compression.as_ref().unwrap();
7168 assert_eq!(
7169 compression.scheme(),
7170 pb21::CompressionScheme::CompressionAlgorithmZstd
7171 );
7172 }
7173
7174 #[tokio::test]
7175 async fn test_dict_values_general_compression_none() {
7176 let mut metadata = HashMap::new();
7177 metadata.insert(
7178 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7179 "none".to_string(),
7180 );
7181 let page = encode_variable_dict_page(metadata).await;
7182 let dictionary_encoding = dictionary_encoding_from_page(&page);
7183 assert!(
7184 !matches!(
7185 dictionary_encoding.compression.as_ref(),
7186 Some(Compression::General(_))
7187 ),
7188 "Expected dictionary values to avoid General compression"
7189 );
7190 }
7191
7192 #[test]
7193 fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
7194 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7195 &HashMap::new(),
7196 None,
7197 None,
7198 );
7199 assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
7200 assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
7201 }
7202
7203 #[test]
7204 fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
7205 let field_metadata = HashMap::from([
7206 (
7207 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7208 "none".to_string(),
7209 ),
7210 (
7211 DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
7212 "7".to_string(),
7213 ),
7214 ]);
7215 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7216 &field_metadata,
7217 Some("zstd".to_string()),
7218 Some("3".to_string()),
7219 );
7220 assert_eq!(
7221 metadata.get(COMPRESSION_META_KEY),
7222 Some(&"none".to_string()),
7223 );
7224 assert_eq!(
7225 metadata.get(COMPRESSION_LEVEL_META_KEY),
7226 Some(&"7".to_string()),
7227 );
7228 }
7229
7230 #[test]
7231 fn test_resolve_dict_values_compression_metadata_env_fallback() {
7232 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7233 &HashMap::new(),
7234 Some("zstd".to_string()),
7235 Some("9".to_string()),
7236 );
7237 assert_eq!(
7238 metadata.get(COMPRESSION_META_KEY),
7239 Some(&"zstd".to_string()),
7240 );
7241 assert_eq!(
7242 metadata.get(COMPRESSION_LEVEL_META_KEY),
7243 Some(&"9".to_string()),
7244 );
7245 }
7246
7247 #[tokio::test]
7248 async fn test_dictionary_encode_int64() {
7249 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7250 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7251 use crate::version::LanceFileVersion;
7252 use arrow_array::{ArrayRef, Int64Array};
7253 use std::collections::HashMap;
7254 use std::sync::Arc;
7255
7256 let values = (0..1000)
7258 .map(|i| match i % 3 {
7259 0 => 10i64,
7260 1 => 20i64,
7261 _ => 30i64,
7262 })
7263 .collect::<Vec<_>>();
7264 let array = Arc::new(Int64Array::from(values)) as ArrayRef;
7265
7266 let mut metadata = HashMap::new();
7267 metadata.insert(
7268 STRUCTURAL_ENCODING_META_KEY.to_string(),
7269 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7270 );
7271 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7272
7273 let test_cases = TestCases::default()
7274 .with_min_file_version(LanceFileVersion::V2_2)
7275 .with_batch_size(1000)
7276 .with_range(0..1000)
7277 .with_indices(vec![0, 1, 10, 999])
7278 .with_expected_encoding("dictionary");
7279
7280 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7281 }
7282
7283 #[tokio::test]
7284 async fn test_dictionary_encode_float64() {
7285 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7286 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7287 use crate::version::LanceFileVersion;
7288 use arrow_array::{ArrayRef, Float64Array};
7289 use std::collections::HashMap;
7290 use std::sync::Arc;
7291
7292 let values = (0..1000)
7294 .map(|i| match i % 3 {
7295 0 => 0.1f64,
7296 1 => 0.2f64,
7297 _ => 0.3f64,
7298 })
7299 .collect::<Vec<_>>();
7300 let array = Arc::new(Float64Array::from(values)) as ArrayRef;
7301
7302 let mut metadata = HashMap::new();
7303 metadata.insert(
7304 STRUCTURAL_ENCODING_META_KEY.to_string(),
7305 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7306 );
7307 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7308
7309 let test_cases = TestCases::default()
7310 .with_min_file_version(LanceFileVersion::V2_2)
7311 .with_batch_size(1000)
7312 .with_range(0..1000)
7313 .with_indices(vec![0, 1, 10, 999])
7314 .with_expected_encoding("dictionary");
7315
7316 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7317 }
7318
7319 #[test]
7320 fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
7321 let rows = 10_000;
7322 let unique_values = 2_000;
7323
7324 let dictionary_encoding =
7325 ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
7326 let layout = pb21::MiniBlockLayout {
7327 rep_compression: None,
7328 def_compression: None,
7329 value_compression: Some(ProtobufUtils21::flat(64, None)),
7330 dictionary: Some(dictionary_encoding),
7331 num_dictionary_items: unique_values,
7332 layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
7333 num_buffers: 1,
7334 repetition_index_depth: 0,
7335 num_items: rows,
7336 has_large_chunk: false,
7337 };
7338
7339 let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
7340 let scheduler = super::MiniBlockScheduler::try_new(
7341 &buffer_offsets_and_sizes,
7342 0,
7343 rows,
7344 &layout,
7345 &DefaultDecompressionStrategy::default(),
7346 )
7347 .unwrap();
7348
7349 let dictionary = scheduler.dictionary.unwrap();
7350 assert_eq!(dictionary.num_dictionary_items, unique_values);
7351 assert_eq!(
7352 dictionary.dictionary_data_alignment,
7353 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
7354 );
7355 }
7356
7357 fn create_test_fixed_data_block(
7359 num_values: u64,
7360 cardinality: u64,
7361 bits_per_value: u64,
7362 ) -> DataBlock {
7363 assert!(cardinality > 0);
7364 assert!(cardinality <= num_values);
7365 let block_info = BlockInfo::default();
7366
7367 assert_eq!(bits_per_value % 8, 0);
7368 let data = match bits_per_value {
7369 32 => {
7370 let values = (0..num_values)
7371 .map(|i| (i % cardinality) as u32)
7372 .collect::<Vec<_>>();
7373 crate::buffer::LanceBuffer::reinterpret_vec(values)
7374 }
7375 64 => {
7376 let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
7377 crate::buffer::LanceBuffer::reinterpret_vec(values)
7378 }
7379 128 => {
7380 let values = (0..num_values)
7381 .map(|i| (i % cardinality) as u128)
7382 .collect::<Vec<_>>();
7383 crate::buffer::LanceBuffer::reinterpret_vec(values)
7384 }
7385 _ => unreachable!(),
7386 };
7387 DataBlock::FixedWidth(FixedWidthDataBlock {
7388 bits_per_value,
7389 data,
7390 num_values,
7391 block_info,
7392 })
7393 }
7394
7395 fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7397 use arrow_array::StringArray;
7398
7399 assert!(cardinality <= num_values && cardinality > 0);
7400
7401 let mut values = Vec::with_capacity(num_values as usize);
7402 for i in 0..num_values {
7403 values.push(format!("value_{:016}", i % cardinality));
7404 }
7405
7406 let array = StringArray::from(values);
7407 DataBlock::from_array(Arc::new(array) as ArrayRef)
7408 }
7409
7410 fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
7411 use arrow_array::StringArray;
7412
7413 assert!(cardinality <= num_values && cardinality > 0);
7414
7415 let mut values = Vec::with_capacity(num_values as usize);
7416 for i in 0..num_values {
7417 let value_idx = i * cardinality / num_values;
7418 values.push(format!("value_{:016}", value_idx));
7419 }
7420
7421 Arc::new(StringArray::from(values)) as ArrayRef
7422 }
7423
7424 fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7425 DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
7426 }
7427
7428 #[test]
7429 fn test_should_dictionary_encode() {
7430 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7431 use lance_core::datatypes::Field as LanceField;
7432
7433 let block = create_test_variable_width_block(1000, 10);
7435
7436 let mut metadata = HashMap::new();
7437 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7438 let arrow_field =
7439 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7440 let field = LanceField::try_from(&arrow_field).unwrap();
7441
7442 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7443 &block,
7444 &field,
7445 LanceFileVersion::V2_1,
7446 );
7447
7448 assert!(
7449 result.is_some(),
7450 "Should use dictionary encode based on size"
7451 );
7452 }
7453
7454 #[test]
7455 fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
7456 let sample_count: usize = 4096;
7457 let num_values: u64 = 200_000;
7458 let cardinality: u64 = 8_000;
7459 let run_length = num_values / cardinality;
7460 let stride = num_values as usize / sample_count;
7461 assert!(
7462 stride > run_length as usize,
7463 "test must construct the stride > run_length case"
7464 );
7465
7466 let block = create_sorted_variable_width_block(num_values, cardinality);
7467 let sample_unique_ratio =
7468 PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
7469
7470 assert!(
7471 sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
7472 "sorted low-cardinality data must not be classified as near-unique"
7473 );
7474 }
7475
7476 #[test]
7477 fn test_should_dictionary_encode_sorted_low_cardinality() {
7478 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7479 use lance_core::datatypes::Field as LanceField;
7480
7481 let block = create_sorted_variable_width_block(200_000, 8_000);
7482
7483 let mut metadata = HashMap::new();
7484 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7485 let arrow_field =
7486 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7487 let field = LanceField::try_from(&arrow_field).unwrap();
7488
7489 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7490 &block,
7491 &field,
7492 LanceFileVersion::V2_2,
7493 );
7494
7495 assert!(
7496 result.is_some(),
7497 "sorted low-cardinality data should reach dictionary encoding"
7498 );
7499 }
7500
7501 #[test]
7502 fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
7503 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7504 use lance_core::datatypes::Field as LanceField;
7505
7506 let num_values = 200_002;
7507 let cardinality = 100_001;
7508 let block = create_sorted_variable_width_block(num_values, cardinality);
7509
7510 let mut metadata = HashMap::new();
7511 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7512 let arrow_field =
7513 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7514 let field = LanceField::try_from(&arrow_field).unwrap();
7515
7516 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7517 &block,
7518 &field,
7519 LanceFileVersion::V2_2,
7520 );
7521
7522 assert!(
7523 result.is_none(),
7524 "sorted high-cardinality short runs should not trigger a full dictionary probe"
7525 );
7526 }
7527
7528 #[tokio::test]
7529 async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
7530 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7531
7532 let mut metadata = HashMap::new();
7533 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7534 let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7535 let array = create_sorted_string_array(200_000, 8_000);
7536
7537 let page = encode_first_page(field, array, LanceFileVersion::V2_2).await;
7538 let _ = dictionary_encoding_from_page(&page);
7539 }
7540
7541 #[test]
7542 fn test_should_not_dictionary_encode_unsupported_bits() {
7543 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7544 use lance_core::datatypes::Field as LanceField;
7545
7546 let block = create_test_fixed_data_block(1000, 1000, 32);
7547
7548 let mut metadata = HashMap::new();
7549 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7550 let arrow_field =
7551 arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
7552 let field = LanceField::try_from(&arrow_field).unwrap();
7553
7554 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7555 &block,
7556 &field,
7557 LanceFileVersion::V2_1,
7558 );
7559
7560 assert!(
7561 result.is_none(),
7562 "Should not use dictionary encode for unsupported bit width"
7563 );
7564 }
7565
7566 #[test]
7567 fn test_should_not_dictionary_encode_near_unique_sample() {
7568 use crate::constants::DICT_SIZE_RATIO_META_KEY;
7569 use lance_core::datatypes::Field as LanceField;
7570
7571 let num_values = 5000;
7572 let block = create_test_variable_width_block(num_values, num_values);
7573
7574 let mut metadata = HashMap::new();
7575 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
7576 let arrow_field =
7577 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7578 let field = LanceField::try_from(&arrow_field).unwrap();
7579
7580 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7581 &block,
7582 &field,
7583 LanceFileVersion::V2_1,
7584 );
7585
7586 assert!(
7587 result.is_none(),
7588 "Should not probe dictionary encoding for near-unique data"
7589 );
7590 }
7591
7592 #[test]
7593 fn test_v2_1_miniblock_serializes_log_num_values_15() {
7594 let miniblocks = MiniBlockCompressed {
7595 data: vec![LanceBuffer::from(vec![1_u8; 16])],
7596 chunks: vec![
7597 MiniBlockChunk {
7598 buffer_sizes: vec![8],
7599 log_num_values: 15,
7600 },
7601 MiniBlockChunk {
7602 buffer_sizes: vec![8],
7603 log_num_values: 0,
7604 },
7605 ],
7606 num_values: 32_769,
7607 };
7608
7609 let serialized =
7610 PrimitiveStructuralEncoder::serialize_miniblocks(miniblocks, None, None, false)
7611 .unwrap();
7612
7613 let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
7614 assert_eq!(chunk_metadata.len(), 2);
7615 assert_eq!(
7616 chunk_metadata[0] & 0x0F,
7617 15,
7618 "V2.1 metadata should use all 4 bits for log_num_values"
7619 );
7620 }
7621
7622 async fn encode_first_page(
7623 field: arrow_schema::Field,
7624 array: ArrayRef,
7625 version: LanceFileVersion,
7626 ) -> crate::encoder::EncodedPage {
7627 use crate::encoder::{
7628 ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
7629 default_encoding_strategy,
7630 };
7631 use crate::repdef::RepDefBuilder;
7632
7633 let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
7634 let encoding_strategy = default_encoding_strategy(version);
7635 let mut column_index_seq = ColumnIndexSequence::default();
7636 let encoding_options = EncodingOptions {
7637 cache_bytes_per_column: 1,
7638 max_page_bytes: 32 * 1024 * 1024,
7639 keep_original_array: true,
7640 buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
7641 version,
7642 };
7643
7644 let mut encoder = encoding_strategy
7645 .create_field_encoder(
7646 encoding_strategy.as_ref(),
7647 &lance_field,
7648 &mut column_index_seq,
7649 &encoding_options,
7650 )
7651 .unwrap();
7652
7653 let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
7654 let repdef = RepDefBuilder::default();
7655 let num_rows = array.len() as u64;
7656 let mut pages = Vec::new();
7657 for task in encoder
7658 .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
7659 .unwrap()
7660 {
7661 pages.push(task.await.unwrap());
7662 }
7663 for task in encoder.flush(&mut external_buffers).unwrap() {
7664 pages.push(task.await.unwrap());
7665 }
7666 pages.into_iter().next().unwrap()
7667 }
7668
7669 #[tokio::test]
7670 async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
7671 use crate::format::pb21::page_layout::Layout;
7672
7673 let val = vec![0xABu8; 33];
7674 let arr: ArrayRef = Arc::new(
7675 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7676 std::iter::repeat_n(Some(val.as_slice()), 256),
7677 33,
7678 )
7679 .unwrap(),
7680 );
7681 let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
7682 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7683
7684 let PageEncoding::Structural(layout) = &page.description else {
7685 panic!("Expected structural encoding");
7686 };
7687 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7688 panic!("Expected constant layout in slot 2");
7689 };
7690 assert!(layout.inline_value.is_none());
7691 assert_eq!(page.data.len(), 1);
7692
7693 let test_cases = TestCases::default()
7694 .with_min_file_version(LanceFileVersion::V2_2)
7695 .with_max_file_version(LanceFileVersion::V2_2)
7696 .with_page_sizes(vec![4096]);
7697 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7698 }
7699
7700 #[tokio::test]
7701 async fn test_constant_layout_out_of_line_utf8_v2_2() {
7702 use crate::format::pb21::page_layout::Layout;
7703
7704 let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
7705 std::iter::repeat_n("hello", 512),
7706 ));
7707 let field = arrow_schema::Field::new("c", DataType::Utf8, true);
7708 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7709
7710 let PageEncoding::Structural(layout) = &page.description else {
7711 panic!("Expected structural encoding");
7712 };
7713 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7714 panic!("Expected constant layout in slot 2");
7715 };
7716 assert!(layout.inline_value.is_none());
7717 assert_eq!(page.data.len(), 1);
7718
7719 let test_cases = TestCases::default()
7720 .with_min_file_version(LanceFileVersion::V2_2)
7721 .with_max_file_version(LanceFileVersion::V2_2)
7722 .with_page_sizes(vec![4096]);
7723 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7724 }
7725
7726 #[tokio::test]
7727 async fn test_constant_layout_nullable_item_v2_2() {
7728 use crate::format::pb21::page_layout::Layout;
7729
7730 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
7731 Some(7),
7732 None,
7733 Some(7),
7734 None,
7735 Some(7),
7736 ]));
7737 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7738 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7739
7740 let PageEncoding::Structural(layout) = &page.description else {
7741 panic!("Expected structural encoding");
7742 };
7743 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7744 panic!("Expected constant layout in slot 2");
7745 };
7746 assert!(layout.inline_value.is_some());
7747 assert_eq!(page.data.len(), 2);
7748
7749 let test_cases = TestCases::default()
7750 .with_min_file_version(LanceFileVersion::V2_2)
7751 .with_max_file_version(LanceFileVersion::V2_2)
7752 .with_page_sizes(vec![4096]);
7753 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7754 }
7755
7756 #[tokio::test]
7757 async fn test_constant_layout_list_repdef_v2_2() {
7758 use crate::format::pb21::page_layout::Layout;
7759 use arrow_array::builder::{Int32Builder, ListBuilder};
7760
7761 let mut builder = ListBuilder::new(Int32Builder::new());
7762 builder.values().append_value(7);
7763 builder.values().append_null();
7764 builder.values().append_value(7);
7765 builder.append(true);
7766
7767 builder.append(true);
7768
7769 builder.values().append_value(7);
7770 builder.append(true);
7771
7772 builder.append_null();
7773
7774 let arr: ArrayRef = Arc::new(builder.finish());
7775 let field = arrow_schema::Field::new(
7776 "c",
7777 DataType::List(Arc::new(arrow_schema::Field::new(
7778 "item",
7779 DataType::Int32,
7780 true,
7781 ))),
7782 true,
7783 );
7784 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7785
7786 let PageEncoding::Structural(layout) = &page.description else {
7787 panic!("Expected structural encoding");
7788 };
7789 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7790 panic!("Expected constant layout in slot 2");
7791 };
7792 assert!(layout.inline_value.is_some());
7793 assert_eq!(page.data.len(), 2);
7794
7795 let test_cases = TestCases::default()
7796 .with_min_file_version(LanceFileVersion::V2_2)
7797 .with_max_file_version(LanceFileVersion::V2_2)
7798 .with_page_sizes(vec![4096]);
7799 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7800 }
7801
7802 #[tokio::test]
7803 async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
7804 use crate::format::pb21::page_layout::Layout;
7805 use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
7806
7807 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
7808 for _ in 0..64 {
7809 builder.values().append_value(1);
7810 builder.values().append_null();
7811 builder.values().append_value(3);
7812 builder.append(true);
7813 }
7814 let arr: ArrayRef = Arc::new(builder.finish());
7815 let field = arrow_schema::Field::new(
7816 "c",
7817 DataType::FixedSizeList(
7818 Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
7819 3,
7820 ),
7821 true,
7822 );
7823 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7824
7825 if let PageEncoding::Structural(layout) = &page.description {
7826 assert!(
7827 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7828 "FixedSizeList should not use constant layout yet"
7829 );
7830 }
7831
7832 let test_cases = TestCases::default()
7833 .with_min_file_version(LanceFileVersion::V2_2)
7834 .with_max_file_version(LanceFileVersion::V2_2)
7835 .with_page_sizes(vec![4096]);
7836 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7837 }
7838
7839 #[tokio::test]
7840 async fn test_constant_layout_not_written_before_v2_2() {
7841 use crate::format::pb21::page_layout::Layout;
7842
7843 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
7844 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7845 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await;
7846
7847 let PageEncoding::Structural(layout) = &page.description else {
7848 return;
7849 };
7850 assert!(
7851 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7852 "Should not emit constant layout before v2.2"
7853 );
7854
7855 let test_cases = TestCases::default()
7856 .with_min_file_version(LanceFileVersion::V2_1)
7857 .with_max_file_version(LanceFileVersion::V2_1)
7858 .with_page_sizes(vec![4096]);
7859 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7860 }
7861
7862 #[tokio::test]
7863 async fn test_all_null_constant_layout_still_works_v2_2() {
7864 use crate::format::pb21::page_layout::Layout;
7865
7866 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
7867 let field = arrow_schema::Field::new("c", DataType::Int32, true);
7868 let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7869
7870 let PageEncoding::Structural(layout) = &page.description else {
7871 panic!("Expected structural encoding");
7872 };
7873 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7874 panic!("Expected layout in slot 2");
7875 };
7876 assert!(layout.inline_value.is_none());
7877 assert_eq!(page.data.len(), 0);
7878
7879 let test_cases = TestCases::default()
7880 .with_min_file_version(LanceFileVersion::V2_2)
7881 .with_max_file_version(LanceFileVersion::V2_2)
7882 .with_page_sizes(vec![4096]);
7883 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7884 }
7885
7886 #[test]
7887 fn test_encode_decode_complex_all_null_vals_roundtrip() {
7888 use crate::compression::{
7889 DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy,
7890 };
7891
7892 let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
7893
7894 let compression_strategy = DefaultCompressionStrategy::default();
7895 let decompression_strategy = DefaultDecompressionStrategy::default();
7896
7897 let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
7898 &values,
7899 &compression_strategy,
7900 )
7901 .unwrap();
7902
7903 let decompressor = decompression_strategy
7904 .create_block_decompressor(&encoding)
7905 .unwrap();
7906 let decompressed = decompressor
7907 .decompress(compressed_buf, values.len() as u64)
7908 .unwrap();
7909 let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
7910 assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
7911 assert_eq!(decompressed_fixed_width.bits_per_value, 16);
7912 let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
7913 assert_eq!(rep_result.as_ref(), values.as_ref());
7914 }
7915
7916 #[tokio::test]
7917 async fn test_complex_all_null_compression_gated_by_version() {
7918 use crate::format::pb21::page_layout::Layout;
7919 use arrow_array::ListArray;
7920
7921 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7922 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7923 );
7924 let arr: ArrayRef = Arc::new(list_array);
7925 let field = arrow_schema::Field::new(
7926 "c",
7927 DataType::List(Arc::new(arrow_schema::Field::new(
7928 "item",
7929 DataType::Int32,
7930 true,
7931 ))),
7932 true,
7933 );
7934
7935 let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await;
7936 let PageEncoding::Structural(layout_v21) = &page_v21.description else {
7937 panic!("Expected structural encoding");
7938 };
7939 let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
7940 panic!("Expected constant layout");
7941 };
7942 assert!(layout_v21.rep_compression.is_none());
7943 assert!(layout_v21.def_compression.is_none());
7944 assert_eq!(layout_v21.num_rep_values, 0);
7945 assert_eq!(layout_v21.num_def_values, 0);
7946
7947 let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await;
7948 let PageEncoding::Structural(layout_v22) = &page_v22.description else {
7949 panic!("Expected structural encoding");
7950 };
7951 let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
7952 panic!("Expected constant layout");
7953 };
7954 assert!(layout_v22.def_compression.is_some());
7955 assert!(layout_v22.num_def_values > 0);
7956 }
7957
7958 #[tokio::test]
7959 async fn test_complex_all_null_round_trip() {
7960 use arrow_array::ListArray;
7961
7962 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7963 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7964 );
7965
7966 let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2);
7967 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
7968 .await;
7969 }
7970
7971 #[tokio::test]
7973 async fn test_sparse_boolean_list_roundtrip() {
7974 use arrow_array::builder::{BooleanBuilder, ListBuilder};
7975
7976 let mut list_builder = ListBuilder::new(BooleanBuilder::new());
7977 for i in 0..1000i32 {
7978 if i % 64 == 0 {
7979 list_builder.values().append_value(i % 128 == 0);
7981 list_builder.append(true);
7982 } else {
7983 list_builder.append(false);
7984 }
7985 }
7986 let list_array = Arc::new(list_builder.finish());
7987
7988 let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
7989 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
7990 }
7991
7992 fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
7993 use crate::compression::VariablePerValueDecompressor;
7994 use crate::encodings::physical::binary::VariableDecoder;
7995 use crate::repdef::{ControlWordParser, DefinitionInterpretation};
7996 use std::sync::Arc;
7997 Arc::new(super::FullZipDecodeDetails {
7998 value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
7999 VariableDecoder::default(),
8000 )
8001 as Arc<dyn VariablePerValueDecompressor>),
8002 def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
8003 ctrl_word_parser: ControlWordParser::new(0, 0),
8004 max_rep: 0,
8005 max_visible_def: 0,
8006 })
8007 }
8008
8009 fn decode_variable_full_zip(
8010 buf: Vec<u8>,
8011 bits_per_offset: u8,
8012 ) -> lance_core::Result<super::VariableFullZipDecoder> {
8013 use std::collections::VecDeque;
8014 let mut data = VecDeque::new();
8015 data.push_back(crate::buffer::LanceBuffer::from(buf));
8016 super::VariableFullZipDecoder::new(
8017 truncated_tail_details(),
8018 data,
8019 1,
8020 bits_per_offset,
8021 bits_per_offset,
8022 )
8023 }
8024
8025 #[test]
8034 fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
8035 use lance_core::Error;
8036
8037 for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
8038 let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
8039 .expect_err("a truncated length prefix must not decode");
8040 assert!(
8041 matches!(err, Error::CorruptFile { .. }),
8042 "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
8043 bits,
8044 buf_len,
8045 err
8046 );
8047 let msg = err.to_string();
8048 assert!(
8049 msg.contains("truncated length prefix"),
8050 "error should say what is wrong, got: {msg}"
8051 );
8052 }
8053 }
8054}