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 STRUCTURAL_ENCODING_SPARSE,
19 },
20 data::DictionaryDataBlock,
21 encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler},
22 format::{
23 ProtobufUtils21,
24 pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression},
25 },
26};
27use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type};
28use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
29use arrow_schema::{DataType, Field as ArrowField};
30use bytes::Bytes;
31use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
32use itertools::Itertools;
33use lance_arrow::DataTypeExt;
34use lance_arrow::deepcopy::deep_copy_nulls;
35use lance_core::{
36 cache::{CacheKey, CacheKeySchema, Context, DeepSizeOf, KeyBuilder},
37 error::{Error, LanceOptionExt},
38 utils::bit::pad_bytes,
39};
40use log::{debug, trace};
41
42use crate::encodings::logical::primitive::miniblock::MiniBlockChunk;
43use crate::encodings::physical::rle::{RleDecompressor, RleRuns};
44use crate::utils::bytepack::ByteUnpacker;
45use crate::{
46 compression::{
47 BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
48 create_rle_decompressor,
49 },
50 data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
51 utils::bytepack::BytepackedIntegerEncoder,
52};
53use crate::{
54 compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
55 encodings::logical::primitive::fullzip::PerValueDataBlock,
56};
57use crate::{
58 encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext},
59 statistics::{ComputeStat, GetStat, Stat},
60};
61use crate::{
62 repdef::{
63 CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
64 MiniBlockRepDefBudget, NormalizedStructuralPlan, RepDefSlicer, SerializedRepDefs,
65 build_control_word_iterator,
66 },
67 utils::accumulation::AccumulationQueue,
68};
69use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
70
71use crate::constants::{
72 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
73 DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
74 DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
75 DICT_VALUES_COMPRESSION_META_KEY,
76};
77use crate::{
78 EncodingsIo,
79 buffer::LanceBuffer,
80 data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
81 decoder::{
82 ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
83 MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
84 StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
85 StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
86 },
87 encoder::{
88 EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
89 },
90 repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
91};
92
93pub mod blob;
94mod chunk_index;
95pub mod constant;
96pub mod dict;
97pub mod fullzip;
98mod layout;
99pub mod miniblock;
100pub(crate) mod sparse;
101
102use chunk_index::{ItemCounts, MiniBlockChunkIndex, PrefixSums, RowMapping, parse_nested_rep};
103
104const FILL_BYTE: u8 = 0xFE;
105const DEFAULT_DICT_DIVISOR: u64 = 2;
106const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
107const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
108const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
109
110struct PageLoadTask {
111 decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
112 num_rows: u64,
113}
114
115trait StructuralPageScheduler: std::fmt::Debug + Send {
118 fn initialize<'a>(
120 &'a mut self,
121 io: &Arc<dyn EncodingsIo>,
122 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
123 fn load(&mut self, data: &Arc<dyn CachedPageData>);
125 fn schedule_ranges(
134 &self,
135 ranges: &[Range<u64>],
136 io: &Arc<dyn EncodingsIo>,
137 ) -> Result<Vec<PageLoadTask>>;
138}
139
140#[derive(Debug)]
142struct ChunkMeta {
143 num_values: u64,
144 chunk_size_bytes: u64,
145 offset_bytes: u64,
146}
147
148#[derive(Debug, Clone)]
150struct DecodedMiniBlockChunk {
151 rep: Option<ScalarBuffer<u16>>,
152 def: Option<ScalarBuffer<u16>>,
153 values: DataBlock,
154}
155
156#[derive(Debug)]
164struct DecodeMiniBlockTask {
165 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
166 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
167 value_decompressor: Arc<dyn MiniBlockDecompressor>,
168 dictionary_data: Option<Arc<DataBlock>>,
169 def_meaning: Arc<[DefinitionInterpretation]>,
170 num_buffers: u64,
171 max_visible_level: u16,
172 instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
173 has_large_chunk: bool,
174}
175
176impl DecodeMiniBlockTask {
177 fn decode_levels(
178 rep_decompressor: &dyn BlockDecompressor,
179 levels: LanceBuffer,
180 num_levels: u16,
181 ) -> Result<ScalarBuffer<u16>> {
182 let rep = rep_decompressor.decompress(levels, num_levels as u64)?;
183 let rep = rep.as_fixed_width().unwrap();
184 debug_assert_eq!(rep.num_values, num_levels as u64);
185 debug_assert_eq!(rep.bits_per_value, 16);
186 Ok(rep.data.borrow_to_typed_slice::<u16>())
187 }
188
189 fn extend_levels(
196 range: Range<u64>,
197 levels: &mut Option<LevelBuffer>,
198 level_buf: &Option<impl AsRef<[u16]>>,
199 dest_offset: usize,
200 ) {
201 if let Some(level_buf) = level_buf {
202 if levels.is_none() {
203 let mut new_levels_vec =
206 LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
207 new_levels_vec.extend(iter::repeat_n(0, dest_offset));
208 *levels = Some(new_levels_vec);
209 }
210 levels.as_mut().unwrap().extend(
211 level_buf.as_ref()[range.start as usize..range.end as usize]
212 .iter()
213 .copied(),
214 );
215 } else if let Some(levels) = levels {
216 let num_values = (range.end - range.start) as usize;
217 levels.extend(iter::repeat_n(0, num_values));
220 }
221 }
222
223 fn map_range(
260 range: Range<u64>,
261 rep: Option<&impl AsRef<[u16]>>,
262 def: Option<&impl AsRef<[u16]>>,
263 max_rep: u16,
264 max_visible_def: u16,
265 total_items: u64,
268 preamble_action: PreambleAction,
269 ) -> (Range<u64>, Range<u64>) {
270 if let Some(rep) = rep {
271 let mut rep = rep.as_ref();
272 let mut items_in_preamble = 0_u64;
275 let first_row_start = match preamble_action {
276 PreambleAction::Skip | PreambleAction::Take => {
277 let first_row_start = if let Some(def) = def.as_ref() {
278 let mut first_row_start = None;
279 for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
280 if *rep == max_rep {
281 first_row_start = Some(idx as u64);
282 break;
283 }
284 if *def <= max_visible_def {
285 items_in_preamble += 1;
286 }
287 }
288 first_row_start
289 } else {
290 let first_row_start =
291 rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
292 items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
293 first_row_start
294 };
295 if first_row_start.is_none() {
298 assert!(preamble_action == PreambleAction::Take);
299 return (0..total_items, 0..rep.len() as u64);
300 }
301 let first_row_start = first_row_start.unwrap();
302 rep = &rep[first_row_start as usize..];
303 first_row_start
304 }
305 PreambleAction::Absent => {
306 debug_assert!(rep[0] == max_rep);
307 0
308 }
309 };
310
311 if range.start == range.end {
313 debug_assert!(preamble_action == PreambleAction::Take);
314 debug_assert!(items_in_preamble <= total_items);
315 return (0..items_in_preamble, 0..first_row_start);
316 }
317 assert!(range.start < range.end);
318
319 let mut rows_seen = 0;
320 let mut new_start = 0;
321 let mut new_levels_start = 0;
322
323 if let Some(def) = def {
324 let def = &def.as_ref()[first_row_start as usize..];
325
326 let mut lead_invis_seen = 0;
328
329 if range.start > 0 {
330 if def[0] > max_visible_def {
331 lead_invis_seen += 1;
332 }
333 for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
334 if *rep == max_rep {
335 rows_seen += 1;
336 if rows_seen == range.start {
337 new_start = idx as u64 + 1 - lead_invis_seen;
338 new_levels_start = idx as u64 + 1;
339 break;
340 }
341 }
342 if *def > max_visible_def {
343 lead_invis_seen += 1;
344 }
345 }
346 }
347
348 rows_seen += 1;
349
350 let mut new_end = u64::MAX;
351 let mut new_levels_end = rep.len() as u64;
352 let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
353 let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
354 for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
355 .iter()
356 .zip(&def[(new_levels_start + 1) as usize..])
357 .enumerate()
358 {
359 if *rep == max_rep {
360 rows_seen += 1;
361 if rows_seen == range.end + 1 {
362 new_end = idx as u64 + new_start + 1 - tail_invis_seen;
363 new_levels_end = idx as u64 + new_levels_start + 1;
364 break;
365 }
366 }
367 if *def > max_visible_def {
368 tail_invis_seen += 1;
369 }
370 }
371
372 if new_end == u64::MAX {
373 new_levels_end = rep.len() as u64;
374 let total_invis_seen = lead_invis_seen + tail_invis_seen;
375 new_end = rep.len() as u64 - total_invis_seen;
376 }
377
378 assert_ne!(new_end, u64::MAX);
379
380 if preamble_action == PreambleAction::Skip {
382 new_start += items_in_preamble;
383 new_end += items_in_preamble;
384 new_levels_start += first_row_start;
385 new_levels_end += first_row_start;
386 } else if preamble_action == PreambleAction::Take {
387 debug_assert_eq!(new_start, 0);
388 debug_assert_eq!(new_levels_start, 0);
389 new_end += items_in_preamble;
390 new_levels_end += first_row_start;
391 }
392
393 debug_assert!(new_end <= total_items);
394 (new_start..new_end, new_levels_start..new_levels_end)
395 } else {
396 if range.start > 0 {
402 for (idx, rep) in rep.iter().skip(1).enumerate() {
403 if *rep == max_rep {
404 rows_seen += 1;
405 if rows_seen == range.start {
406 new_start = idx as u64 + 1;
407 break;
408 }
409 }
410 }
411 }
412 let mut new_end = rep.len() as u64;
413 if range.end < total_items {
415 for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
416 if *rep == max_rep {
417 rows_seen += 1;
418 if rows_seen == range.end {
419 new_end = idx as u64 + new_start + 1;
420 break;
421 }
422 }
423 }
424 }
425
426 if preamble_action == PreambleAction::Skip {
428 new_start += first_row_start;
429 new_end += first_row_start;
430 } else if preamble_action == PreambleAction::Take {
431 debug_assert_eq!(new_start, 0);
432 new_end += first_row_start;
433 }
434
435 debug_assert!(new_end <= total_items);
436 (new_start..new_end, new_start..new_end)
437 }
438 } else {
439 (range.clone(), range)
442 }
443 }
444
445 fn read_buffer_sizes<const LARGE: bool>(
447 buf: &[u8],
448 offset: &mut usize,
449 num_buffers: u64,
450 ) -> Vec<u32> {
451 let read_size = if LARGE { 4 } else { 2 };
452 (0..num_buffers)
453 .map(|_| {
454 let bytes = &buf[*offset..*offset + read_size];
455 let size = if LARGE {
456 u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
457 } else {
458 u16::from_le_bytes([bytes[0], bytes[1]]) as u32
460 };
461 *offset += read_size;
462 size
463 })
464 .collect()
465 }
466
467 fn decode_miniblock_chunk(
469 &self,
470 buf: &LanceBuffer,
471 items_in_chunk: u64,
472 ) -> Result<DecodedMiniBlockChunk> {
473 let mut offset = 0;
474 let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
475 offset += 2;
476
477 let rep_size = if self.rep_decompressor.is_some() {
478 let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
479 offset += 2;
480 Some(rep_size)
481 } else {
482 None
483 };
484 let def_size = if self.def_decompressor.is_some() {
485 let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
486 offset += 2;
487 Some(def_size)
488 } else {
489 None
490 };
491
492 let buffer_sizes = if self.has_large_chunk {
493 Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
494 } else {
495 Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
496 };
497
498 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
499
500 let rep = rep_size.map(|rep_size| {
501 let rep = buf.slice_with_length(offset, rep_size as usize);
502 offset += rep_size as usize;
503 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
504 rep
505 });
506
507 let def = def_size.map(|def_size| {
508 let def = buf.slice_with_length(offset, def_size as usize);
509 offset += def_size as usize;
510 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
511 def
512 });
513
514 let buffers = buffer_sizes
515 .into_iter()
516 .map(|buf_size| {
517 let buf = buf.slice_with_length(offset, buf_size as usize);
518 offset += buf_size as usize;
519 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
520 buf
521 })
522 .collect::<Vec<_>>();
523
524 let values = self
525 .value_decompressor
526 .decompress(buffers, items_in_chunk)?;
527
528 let rep = rep
529 .map(|rep| {
530 Self::decode_levels(
531 self.rep_decompressor.as_ref().unwrap().as_ref(),
532 rep,
533 num_levels,
534 )
535 })
536 .transpose()?;
537 let def = def
538 .map(|def| {
539 Self::decode_levels(
540 self.def_decompressor.as_ref().unwrap().as_ref(),
541 def,
542 num_levels,
543 )
544 })
545 .transpose()?;
546
547 Ok(DecodedMiniBlockChunk { rep, def, values })
548 }
549}
550
551impl DecodePageTask for DecodeMiniBlockTask {
552 fn decode(self: Box<Self>) -> Result<DecodedPage> {
553 let mut repbuf: Option<LevelBuffer> = None;
555 let mut defbuf: Option<LevelBuffer> = None;
556
557 let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
558
559 let estimated_size_bytes = self
561 .instructions
562 .iter()
563 .map(|(_, chunk)| chunk.data.len())
564 .sum::<usize>()
565 * 2;
566 let mut data_builder =
567 DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
568
569 let mut level_offset = 0;
571
572 let needs_caching: Vec<bool> = self
574 .instructions
575 .windows(2)
576 .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
577 .chain(std::iter::once(false)) .collect();
579
580 let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
582
583 for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
585 let should_cache_this_chunk = needs_caching[idx];
586
587 let decoded_chunk = match &chunk_cache {
588 Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
589 cached_chunk.clone()
591 }
592 _ => {
593 let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
595
596 if should_cache_this_chunk {
598 chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
599 }
600 decoded
601 }
602 };
603
604 let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
605
606 let row_range_start =
608 instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
609 let row_range_end = row_range_start + instructions.rows_to_take;
610
611 let (item_range, level_range) = Self::map_range(
613 row_range_start..row_range_end,
614 rep.as_ref(),
615 def.as_ref(),
616 max_rep,
617 self.max_visible_level,
618 chunk.items_in_chunk,
619 instructions.preamble_action,
620 );
621 if item_range.end - item_range.start > chunk.items_in_chunk {
622 return Err(lance_core::Error::internal(format!(
623 "Item range {:?} is greater than chunk items in chunk {:?}",
624 item_range, chunk.items_in_chunk
625 )));
626 }
627
628 Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
630 Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
631 level_offset += (level_range.end - level_range.start) as usize;
632 data_builder.append(&values, item_range);
633 }
634
635 let mut data = data_builder.finish();
636
637 let unraveler =
638 RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
639
640 if let Some(dictionary) = &self.dictionary_data {
641 let DataBlock::FixedWidth(indices) = data else {
643 return Err(lance_core::Error::internal(format!(
644 "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
645 data
646 )));
647 };
648 data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
649 indices,
650 dictionary.as_ref().clone(),
651 ));
652 }
653
654 Ok(DecodedPage {
655 data,
656 repdef: unraveler,
657 })
658 }
659}
660
661#[derive(Debug)]
664struct LoadedChunk {
665 data: LanceBuffer,
666 items_in_chunk: u64,
667 byte_range: Range<u64>,
668 chunk_idx: usize,
669}
670
671impl Clone for LoadedChunk {
672 fn clone(&self) -> Self {
673 Self {
674 data: self.data.clone(),
676 items_in_chunk: self.items_in_chunk,
677 byte_range: self.byte_range.clone(),
678 chunk_idx: self.chunk_idx,
679 }
680 }
681}
682
683#[derive(Debug)]
686struct MiniBlockDecoder {
687 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
688 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
689 value_decompressor: Arc<dyn MiniBlockDecompressor>,
690 def_meaning: Arc<[DefinitionInterpretation]>,
691 loaded_chunks: VecDeque<LoadedChunk>,
692 instructions: VecDeque<ChunkInstructions>,
693 offset_in_current_chunk: u64,
694 num_rows: u64,
695 num_buffers: u64,
696 dictionary: Option<Arc<DataBlock>>,
697 has_large_chunk: bool,
698}
699
700impl StructuralPageDecoder for MiniBlockDecoder {
703 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
704 let mut items_desired = num_rows;
705 let mut need_preamble = false;
706 let mut skip_in_chunk = self.offset_in_current_chunk;
707 let mut drain_instructions = Vec::new();
708 while items_desired > 0 || need_preamble {
709 let (instructions, consumed) = self
710 .instructions
711 .front()
712 .unwrap()
713 .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
714
715 while self.loaded_chunks.front().unwrap().chunk_idx
716 != instructions.chunk_instructions.chunk_idx
717 {
718 self.loaded_chunks.pop_front();
719 }
720 drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
721 if consumed {
722 self.instructions.pop_front();
723 }
724 }
725 self.offset_in_current_chunk = skip_in_chunk;
728
729 let max_visible_level = self
730 .def_meaning
731 .iter()
732 .take_while(|l| !l.is_list())
733 .map(|l| l.num_def_levels())
734 .sum::<u16>();
735
736 Ok(Box::new(DecodeMiniBlockTask {
737 instructions: drain_instructions,
738 def_decompressor: self.def_decompressor.clone(),
739 rep_decompressor: self.rep_decompressor.clone(),
740 value_decompressor: self.value_decompressor.clone(),
741 dictionary_data: self.dictionary.clone(),
742 def_meaning: self.def_meaning.clone(),
743 num_buffers: self.num_buffers,
744 max_visible_level,
745 has_large_chunk: self.has_large_chunk,
746 }))
747 }
748
749 fn num_rows(&self) -> u64 {
750 self.num_rows
751 }
752}
753
754#[derive(Debug, Clone)]
758pub(crate) enum LevelCodec {
759 Uncompressed,
761 Rle(Arc<RleDecompressor>),
763 Block(Arc<dyn BlockDecompressor>),
766}
767
768impl LevelCodec {
769 fn try_new(
770 encoding: Option<&CompressiveEncoding>,
771 decompression_strategy: &dyn DecompressionStrategy,
772 ) -> Result<Self> {
773 match encoding {
774 None => Ok(Self::Uncompressed),
775 Some(encoding) => match encoding.compression.as_ref() {
776 Some(Compression::Rle(rle)) => Ok(Self::Rle(Arc::new(create_rle_decompressor(
777 rle,
778 decompression_strategy,
779 )?))),
780 _ => Ok(Self::Block(Arc::from(
781 decompression_strategy.create_block_decompressor(encoding)?,
782 ))),
783 },
784 }
785 }
786}
787
788#[derive(Debug)]
789enum RunEnds {
790 U16(Box<[u16]>),
791 U32(Box<[u32]>),
792 U64(Box<[u64]>),
793}
794
795impl RunEnds {
796 fn width_for(num_values: usize) -> usize {
797 if u16::try_from(num_values).is_ok() {
798 std::mem::size_of::<u16>()
799 } else if u32::try_from(num_values).is_ok() {
800 std::mem::size_of::<u32>()
801 } else {
802 std::mem::size_of::<u64>()
803 }
804 }
805
806 fn len(&self) -> usize {
807 match self {
808 Self::U16(ends) => ends.len(),
809 Self::U32(ends) => ends.len(),
810 Self::U64(ends) => ends.len(),
811 }
812 }
813
814 fn get(&self, run: usize) -> usize {
815 match self {
816 Self::U16(ends) => ends[run] as usize,
817 Self::U32(ends) => ends[run] as usize,
818 Self::U64(ends) => ends[run] as usize,
819 }
820 }
821
822 fn partition_point(&self, logical_index: usize) -> usize {
823 match self {
824 Self::U16(ends) => ends.partition_point(|&end| end as usize <= logical_index),
825 Self::U32(ends) => ends.partition_point(|&end| end as usize <= logical_index),
826 Self::U64(ends) => ends.partition_point(|&end| end as usize <= logical_index),
827 }
828 }
829
830 fn deep_size(&self) -> usize {
831 match self {
832 Self::U16(ends) => std::mem::size_of_val(ends.as_ref()),
833 Self::U32(ends) => std::mem::size_of_val(ends.as_ref()),
834 Self::U64(ends) => std::mem::size_of_val(ends.as_ref()),
835 }
836 }
837}
838
839enum RunEndsBuilder {
840 U16(Vec<u16>),
841 U32(Vec<u32>),
842 U64(Vec<u64>),
843}
844
845impl RunEndsBuilder {
846 fn with_capacity(num_values: usize, capacity: usize) -> Self {
847 if u16::try_from(num_values).is_ok() {
848 Self::U16(Vec::with_capacity(capacity))
849 } else if u32::try_from(num_values).is_ok() {
850 Self::U32(Vec::with_capacity(capacity))
851 } else {
852 Self::U64(Vec::with_capacity(capacity))
853 }
854 }
855
856 fn push(&mut self, end: usize) -> Result<()> {
857 match self {
858 Self::U16(ends) => ends.push(
859 u16::try_from(end)
860 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?,
861 ),
862 Self::U32(ends) => ends.push(
863 u32::try_from(end)
864 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?,
865 ),
866 Self::U64(ends) => ends.push(end as u64),
867 }
868 Ok(())
869 }
870
871 fn set_last(&mut self, end: usize) -> Result<()> {
872 match self {
873 Self::U16(ends) => {
874 let last = ends.last_mut().ok_or_else(|| {
875 Error::internal("Cannot extend an empty coalesced run buffer")
876 })?;
877 *last = u16::try_from(end)
878 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?;
879 }
880 Self::U32(ends) => {
881 let last = ends.last_mut().ok_or_else(|| {
882 Error::internal("Cannot extend an empty coalesced run buffer")
883 })?;
884 *last = u32::try_from(end)
885 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?;
886 }
887 Self::U64(ends) => {
888 let last = ends.last_mut().ok_or_else(|| {
889 Error::internal("Cannot extend an empty coalesced run buffer")
890 })?;
891 *last = end as u64;
892 }
893 }
894 Ok(())
895 }
896
897 fn finish(self) -> RunEnds {
898 match self {
899 Self::U16(ends) => RunEnds::U16(ends.into_boxed_slice()),
900 Self::U32(ends) => RunEnds::U32(ends.into_boxed_slice()),
901 Self::U64(ends) => RunEnds::U64(ends.into_boxed_slice()),
902 }
903 }
904}
905
906#[derive(Debug)]
907enum RunStorage {
908 Physical(RleRuns),
909 Coalesced { values: Box<[u16]>, ends: RunEnds },
910}
911
912impl RunStorage {
913 fn len(&self) -> usize {
914 match self {
915 Self::Physical(runs) => runs.num_values(),
916 Self::Coalesced { ends, .. } => ends.get(ends.len() - 1),
917 }
918 }
919
920 fn num_runs(&self) -> usize {
921 match self {
922 Self::Physical(runs) => runs.num_runs(),
923 Self::Coalesced { values, .. } => values.len(),
924 }
925 }
926
927 fn value(&self, run: usize) -> u16 {
928 match self {
929 Self::Physical(runs) => runs.value(run),
930 Self::Coalesced { values, .. } => values[run],
931 }
932 }
933
934 fn first_value_above(&self, max: u16) -> Option<(usize, u16)> {
935 (0..self.num_runs()).find_map(|run| {
936 let value = self.value(run);
937 (value > max).then_some((run, value))
938 })
939 }
940
941 fn seek(&self, position: &mut RunPosition, logical_index: usize) {
942 if logical_index >= self.len() {
943 *position = RunPosition {
944 run: self.num_runs(),
945 start: self.len(),
946 end: self.len(),
947 };
948 return;
949 }
950
951 match self {
952 Self::Physical(runs) => {
953 if position.run >= runs.num_runs()
954 || position.end == 0
955 || logical_index < position.start
956 {
957 *position = RunPosition {
958 run: 0,
959 start: 0,
960 end: runs.length(0),
961 };
962 }
963 while position.end <= logical_index {
964 self.advance(position);
965 }
966 }
967 Self::Coalesced { ends, .. } => {
968 if logical_index < position.start || logical_index >= position.end {
969 let run = ends.partition_point(logical_index);
970 *position = RunPosition {
971 run,
972 start: if run == 0 { 0 } else { ends.get(run - 1) },
973 end: ends.get(run),
974 };
975 }
976 }
977 }
978 }
979
980 fn advance(&self, position: &mut RunPosition) {
981 let next_run = position.run + 1;
982 if next_run >= self.num_runs() {
983 *position = RunPosition {
984 run: self.num_runs(),
985 start: self.len(),
986 end: self.len(),
987 };
988 return;
989 }
990
991 let start = position.end;
992 position.run = next_run;
993 position.start = start;
994 position.end = match self {
995 Self::Physical(runs) => start + runs.length(next_run),
996 Self::Coalesced { ends, .. } => ends.get(next_run),
997 };
998 }
999
1000 fn deep_size(&self) -> usize {
1001 match self {
1002 Self::Physical(runs) => runs.deep_size(),
1003 Self::Coalesced { values, ends } => {
1004 std::mem::size_of_val(values.as_ref()) + ends.deep_size()
1005 }
1006 }
1007 }
1008}
1009
1010#[derive(Debug, Clone)]
1016enum LazyLevels {
1017 Dense(ScalarBuffer<u16>),
1018 Runs(Arc<RunStorage>),
1019}
1020
1021#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1022enum LevelPlan {
1023 Physical,
1024 Coalesced,
1025 Dense,
1026}
1027
1028#[derive(Debug, Default, Clone, Copy)]
1029struct RunPosition {
1030 run: usize,
1031 start: usize,
1032 end: usize,
1033}
1034
1035#[derive(Debug, Default, Clone, Copy)]
1041struct LevelCursor {
1042 level: usize,
1044 row: u64,
1046 run: RunPosition,
1048}
1049
1050impl LazyLevels {
1051 fn from_rle_runs(runs: RleRuns) -> Result<Self> {
1052 let plan = Self::select_plan(&runs);
1053 match plan {
1054 LevelPlan::Physical => Ok(Self::Runs(Arc::new(RunStorage::Physical(
1055 runs.into_owned(),
1056 )))),
1057 LevelPlan::Coalesced => Self::build_coalesced(runs),
1058 LevelPlan::Dense => Self::build_dense(runs),
1059 }
1060 }
1061
1062 fn select_plan(runs: &RleRuns) -> LevelPlan {
1065 if runs.num_values() == 0 {
1066 return LevelPlan::Dense;
1067 }
1068
1069 let run_storage_size = std::mem::size_of::<RunStorage>() as u128;
1070 let physical_size = run_storage_size + runs.owned_size() as u128;
1071 let coalesced_size = run_storage_size
1072 + (runs.coalesced_runs() as u128)
1073 * (std::mem::size_of::<u16>() + RunEnds::width_for(runs.num_values())) as u128;
1074 let dense_size = (runs.num_values() as u128) * std::mem::size_of::<u16>() as u128;
1075 [
1076 (physical_size, runs.num_runs(), 0usize, LevelPlan::Physical),
1077 (
1078 coalesced_size,
1079 runs.coalesced_runs(),
1080 1usize,
1081 LevelPlan::Coalesced,
1082 ),
1083 (dense_size, runs.num_values(), 2usize, LevelPlan::Dense),
1084 ]
1085 .into_iter()
1086 .min_by_key(|(size, traversal, priority, _)| (*size, *traversal, *priority))
1087 .map(|(_, _, _, plan)| plan)
1088 .unwrap_or(LevelPlan::Dense)
1089 }
1090
1091 fn build_coalesced(runs: RleRuns) -> Result<Self> {
1092 let mut values = Vec::with_capacity(runs.coalesced_runs());
1093 let mut ends = RunEndsBuilder::with_capacity(runs.num_values(), runs.coalesced_runs());
1094 let mut logical_end = 0usize;
1095 for (value, length) in runs.iter() {
1096 logical_end = logical_end
1097 .checked_add(length)
1098 .ok_or_else(|| Error::internal("Validated RLE run length sum overflowed usize"))?;
1099 if values.last().copied() == Some(value) {
1100 ends.set_last(logical_end)?;
1101 } else {
1102 values.push(value);
1103 ends.push(logical_end)?;
1104 }
1105 }
1106 Ok(Self::Runs(Arc::new(RunStorage::Coalesced {
1107 values: values.into_boxed_slice(),
1108 ends: ends.finish(),
1109 })))
1110 }
1111
1112 fn build_dense(runs: RleRuns) -> Result<Self> {
1113 let mut values = Vec::new();
1114 values.try_reserve_exact(runs.num_values()).map_err(|_| {
1115 Error::internal(format!(
1116 "Cannot allocate {} dense repetition/definition levels",
1117 runs.num_values()
1118 ))
1119 })?;
1120 for (value, length) in runs.iter() {
1121 values.resize(values.len() + length, value);
1122 }
1123 Ok(Self::Dense(ScalarBuffer::from(values)))
1124 }
1125
1126 fn len(&self) -> usize {
1127 match self {
1128 Self::Dense(buf) => buf.len(),
1129 Self::Runs(runs) => runs.len(),
1130 }
1131 }
1132
1133 fn validate_max(&self, level_type: &str, max: u16) -> Result<()> {
1134 let invalid = match self {
1135 Self::Dense(levels) => levels
1136 .iter()
1137 .enumerate()
1138 .find_map(|(index, &value)| (value > max).then_some(("index", index, value))),
1139 Self::Runs(runs) => runs
1140 .first_value_above(max)
1141 .map(|(run, value)| ("run", run, value)),
1142 };
1143 if let Some((position_type, position, value)) = invalid {
1144 return Err(Error::invalid_input_source(
1145 format!(
1146 "Invalid {level_type} level {value} at {position_type} {position}: maximum is {max}"
1147 )
1148 .into(),
1149 ));
1150 }
1151 Ok(())
1152 }
1153
1154 fn seek_row_start(
1161 &self,
1162 cursor: &mut LevelCursor,
1163 target_row: u64,
1164 max_rep: u16,
1165 ) -> Result<usize> {
1166 let mut need = target_row.checked_sub(cursor.row).ok_or_else(|| {
1167 Error::internal(format!(
1168 "Complex all-null row ranges are not sorted: target row {target_row} follows {}",
1169 cursor.row
1170 ))
1171 })?;
1172 if need == 0 {
1173 return Ok(cursor.level);
1174 }
1175 match self {
1176 Self::Dense(buf) => {
1177 let mut level = cursor.level;
1178 while need > 0 {
1179 if level >= buf.len() {
1180 return Err(Error::internal(
1181 "Invalid complex all-null layout: repetition buffer too short",
1182 ));
1183 }
1184 if buf[level] != max_rep {
1185 return Err(Error::internal(
1186 "Invalid complex all-null layout: row did not start at max repetition level",
1187 ));
1188 }
1189 level += 1;
1190 while level < buf.len() && buf[level] != max_rep {
1191 level += 1;
1192 }
1193 need -= 1;
1194 }
1195 cursor.level = level;
1196 cursor.row = target_row;
1197 Ok(level)
1198 }
1199 Self::Runs(runs) => {
1200 let mut level = cursor.level;
1201 let mut run = cursor.run;
1202 runs.seek(&mut run, level);
1203 while need > 0 {
1204 if run.run >= runs.num_runs() {
1205 return Err(Error::internal(
1206 "Invalid complex all-null layout: repetition buffer too short",
1207 ));
1208 }
1209 if runs.value(run.run) != max_rep {
1210 return Err(Error::internal(
1211 "Invalid complex all-null layout: row did not start at max repetition level",
1212 ));
1213 }
1214 let avail = (run.end - level) as u64;
1215 if need < avail {
1216 level += need as usize;
1218 need = 0;
1219 } else {
1220 need -= avail;
1223 runs.advance(&mut run);
1224 while run.run < runs.num_runs() && runs.value(run.run) != max_rep {
1225 runs.advance(&mut run);
1226 }
1227 level = if run.run < runs.num_runs() {
1228 run.start
1229 } else {
1230 self.len()
1231 };
1232 }
1233 }
1234 cursor.level = level;
1235 cursor.row = target_row;
1236 cursor.run = run;
1237 Ok(level)
1238 }
1239 }
1240 }
1241
1242 fn count_le_cursor(
1248 &self,
1249 run_cursor: &mut RunPosition,
1250 range: Range<usize>,
1251 max: u16,
1252 ) -> (u64, RunPosition) {
1253 if range.is_empty() {
1254 return (0, *run_cursor);
1255 }
1256 match self {
1257 Self::Dense(buf) => (
1258 buf[range].iter().filter(|&&d| d <= max).count() as u64,
1259 RunPosition::default(),
1260 ),
1261 Self::Runs(runs) => {
1262 runs.seek(run_cursor, range.start);
1264 let start = *run_cursor;
1265 let mut count = 0u64;
1266 let mut current = *run_cursor;
1267 while current.run < runs.num_runs() && current.start < range.end {
1268 if runs.value(current.run) <= max {
1269 let lo = current.start.max(range.start);
1270 let hi = current.end.min(range.end);
1271 count += (hi - lo) as u64;
1272 }
1273 if current.end >= range.end {
1274 break;
1275 }
1276 runs.advance(&mut current);
1277 }
1278 *run_cursor = current;
1281 (count, start)
1282 }
1283 }
1284 }
1285
1286 fn extend_into(&self, range: Range<usize>, run: RunPosition, out: &mut Vec<u16>) {
1287 if range.is_empty() {
1288 return;
1289 }
1290 match self {
1291 Self::Dense(buf) => out.extend_from_slice(&buf[range]),
1292 Self::Runs(runs) => {
1293 let mut current = run;
1294 runs.seek(&mut current, range.start);
1295 while current.run < runs.num_runs() && current.start < range.end {
1296 let lo = current.start.max(range.start);
1297 let hi = current.end.min(range.end);
1298 if hi > lo {
1299 out.resize(out.len() + (hi - lo), runs.value(current.run));
1300 }
1301 runs.advance(&mut current);
1302 }
1303 }
1304 }
1305 }
1306
1307 #[cfg(test)]
1308 fn deep_size(&self) -> usize {
1309 self.deep_size_of_children(&mut Context::new())
1310 }
1311}
1312
1313impl DeepSizeOf for LazyLevels {
1314 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1315 match self {
1316 Self::Dense(buf) => buf.deep_size_of_children(ctx),
1317 Self::Runs(runs) => {
1318 let pointer = Arc::as_ptr(runs) as *const () as usize;
1319 if ctx.mark_seen(pointer) {
1320 std::mem::size_of_val(runs.as_ref()) + runs.deep_size()
1321 } else {
1322 0
1323 }
1324 }
1325 }
1326 }
1327}
1328
1329fn validate_complex_all_null_levels(
1330 rep: &Option<LazyLevels>,
1331 def: &Option<LazyLevels>,
1332 max_rep: u16,
1333 max_def: u16,
1334) -> Result<()> {
1335 if let Some(rep) = rep {
1336 rep.validate_max("repetition", max_rep)?;
1337 }
1338 if let Some(def) = def {
1339 def.validate_max("definition", max_def)?;
1340 }
1341 if let (Some(rep), Some(def)) = (rep, def)
1342 && rep.len() != def.len()
1343 {
1344 return Err(Error::invalid_input_source(
1345 format!(
1346 "Mismatched complex all-null level counts: repetition has {}, definition has {}",
1347 rep.len(),
1348 def.len()
1349 )
1350 .into(),
1351 ));
1352 }
1353 Ok(())
1354}
1355
1356fn expected_level_bytes(num_values: u64, level_type: &str) -> Result<usize> {
1357 usize::try_from(num_values)
1358 .ok()
1359 .and_then(|num_values| num_values.checked_mul(std::mem::size_of::<u16>()))
1360 .ok_or_else(|| {
1361 Error::invalid_input_source(
1362 format!("{level_type} level count {num_values} does not fit in memory").into(),
1363 )
1364 })
1365}
1366
1367fn dense_levels_from_block(
1368 decompressed: DataBlock,
1369 num_values: u64,
1370 level_type: &str,
1371) -> Result<LazyLevels> {
1372 let DataBlock::FixedWidth(block) = decompressed else {
1373 return Err(Error::invalid_input_source(
1374 format!("Expected fixed-width data block for {level_type} levels").into(),
1375 ));
1376 };
1377 if block.num_values != num_values {
1378 return Err(Error::invalid_input_source(
1379 format!(
1380 "Unexpected {level_type} level count after decompression: expected {num_values}, got {}",
1381 block.num_values
1382 )
1383 .into(),
1384 ));
1385 }
1386 if block.bits_per_value != 16 {
1387 return Err(Error::invalid_input_source(
1388 format!(
1389 "Unexpected {level_type} level bit width after decompression: expected 16, got {}",
1390 block.bits_per_value
1391 )
1392 .into(),
1393 ));
1394 }
1395 let expected_bytes = expected_level_bytes(num_values, level_type)?;
1396 if block.data.len() != expected_bytes {
1397 return Err(Error::invalid_input_source(
1398 format!(
1399 "Unexpected decompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1400 block.data.len()
1401 )
1402 .into(),
1403 ));
1404 }
1405 Ok(LazyLevels::Dense(block.data.borrow_to_typed_slice::<u16>()))
1406}
1407
1408#[derive(Debug)]
1409struct CachedComplexAllNullState {
1410 rep: Option<LazyLevels>,
1411 def: Option<LazyLevels>,
1412}
1413
1414impl DeepSizeOf for CachedComplexAllNullState {
1415 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1416 self.rep.deep_size_of_children(ctx) + self.def.deep_size_of_children(ctx)
1417 }
1418}
1419
1420impl CachedPageData for CachedComplexAllNullState {
1421 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1422 self
1423 }
1424}
1425
1426#[derive(Debug)]
1435pub struct ComplexAllNullScheduler {
1436 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1438 def_meaning: Arc<[DefinitionInterpretation]>,
1439 repdef: Option<Arc<CachedComplexAllNullState>>,
1440 max_rep: u16,
1441 max_def: u16,
1442 max_visible_level: u16,
1443 rep_codec: LevelCodec,
1444 def_codec: LevelCodec,
1445 num_rep_values: u64,
1446 num_def_values: u64,
1447}
1448
1449impl ComplexAllNullScheduler {
1450 pub(crate) fn new(
1451 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1452 def_meaning: Arc<[DefinitionInterpretation]>,
1453 rep_codec: LevelCodec,
1454 def_codec: LevelCodec,
1455 num_rep_values: u64,
1456 num_def_values: u64,
1457 ) -> Self {
1458 let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
1459 let max_def = def_meaning
1460 .iter()
1461 .map(|meaning| meaning.num_def_levels())
1462 .sum::<u16>();
1463 let max_visible_level = def_meaning
1464 .iter()
1465 .take_while(|l| !l.is_list())
1466 .map(|l| l.num_def_levels())
1467 .sum::<u16>();
1468 Self {
1469 buffer_offsets_and_sizes,
1470 def_meaning,
1471 repdef: None,
1472 max_rep,
1473 max_def,
1474 max_visible_level,
1475 rep_codec,
1476 def_codec,
1477 num_rep_values,
1478 num_def_values,
1479 }
1480 }
1481}
1482
1483impl StructuralPageScheduler for ComplexAllNullScheduler {
1484 fn initialize<'a>(
1485 &'a mut self,
1486 io: &Arc<dyn EncodingsIo>,
1487 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1488 let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
1490 let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
1491 let has_rep = rep_size > 0;
1492 let has_def = def_size > 0;
1493
1494 let mut reads = Vec::with_capacity(2);
1495 if has_rep {
1496 reads.push(rep_pos..rep_pos + rep_size);
1497 }
1498 if has_def {
1499 reads.push(def_pos..def_pos + def_size);
1500 }
1501
1502 let data = io.submit_request(reads, 0);
1503 let rep_codec = self.rep_codec.clone();
1504 let def_codec = self.def_codec.clone();
1505 let num_rep_values = self.num_rep_values;
1506 let num_def_values = self.num_def_values;
1507 let max_rep = self.max_rep;
1508 let max_def = self.max_def;
1509
1510 async move {
1511 let data = data.await?;
1512 let mut data_iter = data.into_iter();
1513
1514 let build_levels = |compressed_bytes: Bytes,
1517 codec: &LevelCodec,
1518 num_values: u64,
1519 level_type: &str|
1520 -> Result<LazyLevels> {
1521 match codec {
1522 LevelCodec::Uncompressed => {
1523 if num_values == 0 {
1524 if !compressed_bytes
1525 .len()
1526 .is_multiple_of(std::mem::size_of::<u16>())
1527 {
1528 return Err(Error::invalid_input_source(
1529 format!(
1530 "Unexpected uncompressed {level_type} level size: {} bytes is not divisible by {}",
1531 compressed_bytes.len(),
1532 std::mem::size_of::<u16>()
1533 )
1534 .into(),
1535 ));
1536 }
1537 } else {
1538 let expected_bytes = expected_level_bytes(num_values, level_type)?;
1539 if compressed_bytes.len() != expected_bytes {
1540 return Err(Error::invalid_input_source(
1541 format!(
1542 "Unexpected uncompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1543 compressed_bytes.len()
1544 )
1545 .into(),
1546 ));
1547 }
1548 }
1549 let buffer = LanceBuffer::from_bytes(compressed_bytes, 2);
1550 Ok(LazyLevels::Dense(buffer.borrow_to_typed_slice::<u16>()))
1551 }
1552 LevelCodec::Rle(decompressor) => {
1553 let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1554 let runs = decompressor.decode_u16_runs(frame, num_values)?;
1555 LazyLevels::from_rle_runs(runs)
1556 }
1557 LevelCodec::Block(decompressor) => {
1558 let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1559 let decompressed = decompressor.decompress(frame, num_values)?;
1560 dense_levels_from_block(decompressed, num_values, level_type)
1561 }
1562 }
1563 };
1564
1565 let rep = if has_rep {
1566 let rep = data_iter.next().unwrap();
1567 Some(build_levels(rep, &rep_codec, num_rep_values, "repetition")?)
1568 } else {
1569 None
1570 };
1571
1572 let def = if has_def {
1573 let def = data_iter.next().unwrap();
1574 Some(build_levels(def, &def_codec, num_def_values, "definition")?)
1575 } else {
1576 None
1577 };
1578
1579 validate_complex_all_null_levels(&rep, &def, max_rep, max_def)?;
1580 let repdef = Arc::new(CachedComplexAllNullState { rep, def });
1581
1582 self.repdef = Some(repdef.clone());
1583
1584 Ok(repdef as Arc<dyn CachedPageData>)
1585 }
1586 .boxed()
1587 }
1588
1589 fn load(&mut self, data: &Arc<dyn CachedPageData>) {
1590 self.repdef = Some(
1591 data.clone()
1592 .as_arc_any()
1593 .downcast::<CachedComplexAllNullState>()
1594 .unwrap(),
1595 );
1596 }
1597
1598 fn schedule_ranges(
1599 &self,
1600 ranges: &[Range<u64>],
1601 _io: &Arc<dyn EncodingsIo>,
1602 ) -> Result<Vec<PageLoadTask>> {
1603 let ranges = VecDeque::from_iter(ranges.iter().cloned());
1604 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1605 let decoder = Box::new(ComplexAllNullPageDecoder {
1606 ranges,
1607 rep: self.repdef.as_ref().unwrap().rep.clone(),
1608 def: self.repdef.as_ref().unwrap().def.clone(),
1609 num_rows,
1610 def_meaning: self.def_meaning.clone(),
1611 max_rep: self.max_rep,
1612 max_visible_level: self.max_visible_level,
1613 rep_cursor: LevelCursor::default(),
1614 def_run_cursor: RunPosition::default(),
1615 }) as Box<dyn StructuralPageDecoder>;
1616 let page_load_task = PageLoadTask {
1617 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1618 num_rows,
1619 };
1620 Ok(vec![page_load_task])
1621 }
1622}
1623
1624#[derive(Debug)]
1625pub struct ComplexAllNullPageDecoder {
1626 ranges: VecDeque<Range<u64>>,
1627 rep: Option<LazyLevels>,
1628 def: Option<LazyLevels>,
1629 num_rows: u64,
1630 def_meaning: Arc<[DefinitionInterpretation]>,
1631 max_rep: u16,
1632 max_visible_level: u16,
1633 rep_cursor: LevelCursor,
1635 def_run_cursor: RunPosition,
1637}
1638
1639impl ComplexAllNullPageDecoder {
1640 fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
1641 let mut rows_desired = num_rows;
1642 let mut ranges = Vec::with_capacity(self.ranges.len());
1643 while rows_desired > 0 {
1644 let front = self.ranges.front_mut().unwrap();
1645 let avail = front.end - front.start;
1646 if avail > rows_desired {
1647 ranges.push(front.start..front.start + rows_desired);
1648 front.start += rows_desired;
1649 rows_desired = 0;
1650 } else {
1651 ranges.push(self.ranges.pop_front().unwrap());
1652 rows_desired -= avail;
1653 }
1654 }
1655 ranges
1656 }
1657
1658 fn seek_row_start(&mut self, target_row: u64) -> Result<usize> {
1661 match &self.rep {
1662 Some(rep) => rep.seek_row_start(&mut self.rep_cursor, target_row, self.max_rep),
1663 None => {
1664 self.rep_cursor.row = target_row;
1666 self.rep_cursor.level = target_row as usize;
1667 Ok(target_row as usize)
1668 }
1669 }
1670 }
1671
1672 fn count_visible(&mut self, levels: Range<usize>) -> Result<(u64, RunPosition)> {
1675 match &self.def {
1676 Some(def) => {
1677 if levels.end > def.len() {
1678 return Err(Error::internal(
1679 "Invalid complex all-null layout: definition buffer too short",
1680 ));
1681 }
1682 Ok(def.count_le_cursor(&mut self.def_run_cursor, levels, self.max_visible_level))
1683 }
1684 None => Ok(((levels.end - levels.start) as u64, RunPosition::default())),
1685 }
1686 }
1687}
1688
1689impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1690 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1691 let drained_ranges = self.drain_ranges(num_rows);
1692 let mut level_slices: Vec<LevelSlice> = Vec::with_capacity(drained_ranges.len());
1693 let mut visible_items_total = 0;
1694
1695 for range in drained_ranges {
1700 let level_start = self.seek_row_start(range.start)?;
1701 let rep_run = self.rep_cursor.run;
1702 let level_end = self.seek_row_start(range.end)?;
1703 let (visible_items, def_run) = self.count_visible(level_start..level_end)?;
1704 visible_items_total += visible_items;
1705 if let Some(last) = level_slices.last_mut()
1706 && last.range.end == level_start
1707 {
1708 last.range.end = level_end;
1709 } else {
1710 level_slices.push(LevelSlice {
1711 range: level_start..level_end,
1712 rep_run,
1713 def_run,
1714 });
1715 }
1716 }
1717
1718 Ok(Box::new(DecodeComplexAllNullTask {
1719 level_slices,
1720 visible_items_total,
1721 rep: self.rep.clone(),
1722 def: self.def.clone(),
1723 def_meaning: self.def_meaning.clone(),
1724 max_visible_level: self.max_visible_level,
1725 }))
1726 }
1727
1728 fn num_rows(&self) -> u64 {
1729 self.num_rows
1730 }
1731}
1732
1733#[derive(Debug, Clone)]
1736struct LevelSlice {
1737 range: Range<usize>,
1738 rep_run: RunPosition,
1739 def_run: RunPosition,
1740}
1741
1742#[derive(Clone, Copy)]
1743enum LevelKind {
1744 Repetition,
1745 Definition,
1746}
1747
1748impl LevelSlice {
1749 fn run(&self, kind: LevelKind) -> RunPosition {
1750 match kind {
1751 LevelKind::Repetition => self.rep_run,
1752 LevelKind::Definition => self.def_run,
1753 }
1754 }
1755}
1756
1757#[derive(Debug)]
1758pub struct DecodeComplexAllNullTask {
1759 level_slices: Vec<LevelSlice>,
1760 visible_items_total: u64,
1761 rep: Option<LazyLevels>,
1762 def: Option<LazyLevels>,
1763 def_meaning: Arc<[DefinitionInterpretation]>,
1764 max_visible_level: u16,
1765}
1766
1767impl DecodeComplexAllNullTask {
1768 fn decode_level(&self, levels: &Option<LazyLevels>, kind: LevelKind) -> Option<Vec<u16>> {
1769 levels.as_ref().map(|levels| {
1770 let num_levels = self
1771 .level_slices
1772 .iter()
1773 .map(|slice| slice.range.end - slice.range.start)
1774 .sum();
1775 let mut referenced_levels = Vec::with_capacity(num_levels);
1776 for slice in &self.level_slices {
1777 levels.extend_into(slice.range.clone(), slice.run(kind), &mut referenced_levels);
1778 }
1779 referenced_levels
1780 })
1781 }
1782}
1783
1784impl DecodePageTask for DecodeComplexAllNullTask {
1785 fn decode(self: Box<Self>) -> Result<DecodedPage> {
1786 let rep = self.decode_level(&self.rep, LevelKind::Repetition);
1787 let def = self.decode_level(&self.def, LevelKind::Definition);
1788
1789 let num_values = if let Some(def) = &def {
1793 def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
1794 } else {
1795 self.visible_items_total
1796 };
1797
1798 let data = DataBlock::AllNull(AllNullDataBlock { num_values });
1799 let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
1800 Ok(DecodedPage {
1801 data,
1802 repdef: unraveler,
1803 })
1804 }
1805}
1806
1807#[derive(Debug, Default)]
1812pub struct SimpleAllNullScheduler {}
1813
1814impl StructuralPageScheduler for SimpleAllNullScheduler {
1815 fn initialize<'a>(
1816 &'a mut self,
1817 _io: &Arc<dyn EncodingsIo>,
1818 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1819 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
1820 }
1821
1822 fn load(&mut self, _cache: &Arc<dyn CachedPageData>) {}
1823
1824 fn schedule_ranges(
1825 &self,
1826 ranges: &[Range<u64>],
1827 _io: &Arc<dyn EncodingsIo>,
1828 ) -> Result<Vec<PageLoadTask>> {
1829 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1830 let decoder =
1831 Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
1832 let page_load_task = PageLoadTask {
1833 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1834 num_rows,
1835 };
1836 Ok(vec![page_load_task])
1837 }
1838}
1839
1840#[derive(Debug)]
1843struct SimpleAllNullDecodePageTask {
1844 num_values: u64,
1845}
1846impl DecodePageTask for SimpleAllNullDecodePageTask {
1847 fn decode(self: Box<Self>) -> Result<DecodedPage> {
1848 let unraveler = RepDefUnraveler::new(
1849 None,
1850 Some(vec![1; self.num_values as usize]),
1851 Arc::new([DefinitionInterpretation::NullableItem]),
1852 self.num_values,
1853 );
1854 Ok(DecodedPage {
1855 data: DataBlock::AllNull(AllNullDataBlock {
1856 num_values: self.num_values,
1857 }),
1858 repdef: unraveler,
1859 })
1860 }
1861}
1862
1863#[derive(Debug)]
1864pub struct SimpleAllNullPageDecoder {
1865 num_rows: u64,
1866}
1867
1868impl StructuralPageDecoder for SimpleAllNullPageDecoder {
1869 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1870 Ok(Box::new(SimpleAllNullDecodePageTask {
1871 num_values: num_rows,
1872 }))
1873 }
1874
1875 fn num_rows(&self) -> u64 {
1876 self.num_rows
1877 }
1878}
1879
1880#[derive(Debug, Clone)]
1881struct MiniBlockSchedulerDictionary {
1882 dictionary_decompressor: Arc<dyn BlockDecompressor>,
1884 dictionary_buf_position_and_size: (u64, u64),
1885 dictionary_data_alignment: u64,
1886 num_dictionary_items: u64,
1887}
1888
1889#[derive(Debug)]
1891struct MiniBlockCacheableState {
1892 chunk_index: MiniBlockChunkIndex,
1894 dictionary: Option<Arc<DataBlock>>,
1896}
1897
1898impl DeepSizeOf for MiniBlockCacheableState {
1899 fn deep_size_of_children(&self, context: &mut Context) -> usize {
1900 self.chunk_index.deep_size_of_children(context)
1901 + self
1902 .dictionary
1903 .as_ref()
1904 .map(|dict| dict.data_size() as usize)
1905 .unwrap_or(0)
1906 }
1907}
1908
1909impl CachedPageData for MiniBlockCacheableState {
1910 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1911 self
1912 }
1913}
1914
1915#[derive(Debug)]
1942pub struct MiniBlockScheduler {
1943 buffer_offsets_and_sizes: Vec<(u64, u64)>,
1945 priority: u64,
1946 items_in_page: u64,
1947 repetition_index_depth: u16,
1948 num_buffers: u64,
1949 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
1950 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
1951 value_decompressor: Arc<dyn MiniBlockDecompressor>,
1952 def_meaning: Arc<[DefinitionInterpretation]>,
1953 dictionary: Option<MiniBlockSchedulerDictionary>,
1954 page_meta: Option<Arc<MiniBlockCacheableState>>,
1956 has_large_chunk: bool,
1957}
1958
1959impl MiniBlockScheduler {
1960 fn try_new(
1961 buffer_offsets_and_sizes: &[(u64, u64)],
1962 priority: u64,
1963 items_in_page: u64,
1964 layout: &pb21::MiniBlockLayout,
1965 decompressors: &dyn DecompressionStrategy,
1966 ) -> Result<Self> {
1967 let rep_decompressor = layout
1968 .rep_compression
1969 .as_ref()
1970 .map(|rep_compression| {
1971 decompressors
1972 .create_block_decompressor(rep_compression)
1973 .map(Arc::from)
1974 })
1975 .transpose()?;
1976 let def_decompressor = layout
1977 .def_compression
1978 .as_ref()
1979 .map(|def_compression| {
1980 decompressors
1981 .create_block_decompressor(def_compression)
1982 .map(Arc::from)
1983 })
1984 .transpose()?;
1985 let def_meaning = layout
1986 .layers
1987 .iter()
1988 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
1989 .collect::<Vec<_>>();
1990 let value_decompressor = decompressors.create_miniblock_decompressor(
1991 layout.value_compression.as_ref().unwrap(),
1992 decompressors,
1993 )?;
1994
1995 let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
1996 let num_dictionary_items = layout.num_dictionary_items;
1997 let dictionary_decompressor = decompressors
1998 .create_block_decompressor(dictionary_encoding)?
1999 .into();
2000 let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
2001 {
2002 Compression::Variable(_) => 4,
2003 Compression::Flat(_) => 16,
2004 Compression::General(_) => 1,
2005 Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
2006 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
2007 }
2008 _ => {
2009 return Err(Error::invalid_input_source(
2010 format!(
2011 "Unsupported mini-block dictionary encoding: {:?}",
2012 dictionary_encoding.compression.as_ref().unwrap()
2013 )
2014 .into(),
2015 ));
2016 }
2017 };
2018 Some(MiniBlockSchedulerDictionary {
2019 dictionary_decompressor,
2020 dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
2021 dictionary_data_alignment,
2022 num_dictionary_items,
2023 })
2024 } else {
2025 None
2026 };
2027
2028 Ok(Self {
2029 buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
2030 rep_decompressor,
2031 def_decompressor,
2032 value_decompressor: value_decompressor.into(),
2033 repetition_index_depth: layout.repetition_index_depth as u16,
2034 num_buffers: layout.num_buffers,
2035 priority,
2036 items_in_page,
2037 dictionary,
2038 def_meaning: def_meaning.into(),
2039 page_meta: None,
2040 has_large_chunk: layout.has_large_chunk,
2041 })
2042 }
2043
2044 fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
2045 let chunk_index = &self.page_meta.as_ref().unwrap().chunk_index;
2046 chunk_indices
2047 .iter()
2048 .map(|&chunk_idx| LoadedChunk {
2049 byte_range: chunk_index.byte_range(chunk_idx),
2050 items_in_chunk: chunk_index.items_in_chunk(chunk_idx),
2051 chunk_idx,
2052 data: LanceBuffer::empty(),
2053 })
2054 .collect()
2055 }
2056}
2057
2058#[derive(Debug, PartialEq, Eq, Clone, Copy)]
2059enum PreambleAction {
2060 Take,
2061 Skip,
2062 Absent,
2063}
2064
2065#[derive(Clone, Debug, PartialEq, Eq)]
2100struct ChunkInstructions {
2101 chunk_idx: usize,
2103 preamble: PreambleAction,
2109 rows_to_skip: u64,
2113 rows_to_take: u64,
2116 take_trailer: bool,
2123}
2124
2125#[derive(Debug, PartialEq, Eq)]
2143struct ChunkDrainInstructions {
2144 chunk_instructions: ChunkInstructions,
2145 rows_to_skip: u64,
2146 rows_to_take: u64,
2147 preamble_action: PreambleAction,
2148}
2149
2150impl ChunkInstructions {
2151 fn schedule_instructions(
2157 chunk_index: &MiniBlockChunkIndex,
2158 user_ranges: &[Range<u64>],
2159 ) -> Vec<Self> {
2160 let num_chunks = chunk_index.num_chunks();
2163 let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
2167
2168 for user_range in user_ranges {
2169 let mut rows_needed = user_range.end - user_range.start;
2170 let mut need_preamble = false;
2171
2172 let mut block_index = chunk_index.find_chunk(user_range.start);
2175
2176 let mut to_skip = user_range.start - chunk_index.first_row(block_index);
2177
2178 while rows_needed > 0 || need_preamble {
2179 if block_index >= num_chunks {
2181 log::warn!(
2182 "schedule_instructions inconsistency: block_index >= num_chunks, exiting early"
2183 );
2184 break;
2185 }
2186
2187 let starts_including_trailer = chunk_index.rows_in_chunk(block_index);
2188 let has_preamble = chunk_index.has_preamble(block_index);
2189 let has_trailer = chunk_index.has_trailer(block_index);
2190 let rows_avail = starts_including_trailer.saturating_sub(to_skip);
2191
2192 if rows_avail == 0 && to_skip == 0 {
2196 if has_preamble && need_preamble {
2198 chunk_instructions.push(Self {
2199 chunk_idx: block_index,
2200 preamble: PreambleAction::Take,
2201 rows_to_skip: 0,
2202 rows_to_take: 0,
2203 take_trailer: has_trailer,
2207 });
2208 if starts_including_trailer > 0 || block_index == num_chunks - 1 {
2212 need_preamble = false;
2213 }
2214 }
2215 block_index += 1;
2217 continue;
2218 }
2219
2220 if rows_avail == 0 && to_skip > 0 {
2224 to_skip -= starts_including_trailer;
2227 block_index += 1;
2228 continue;
2229 }
2230
2231 let rows_to_take = rows_avail.min(rows_needed);
2232 rows_needed -= rows_to_take;
2233
2234 let mut take_trailer = false;
2235 let preamble = if has_preamble {
2236 if need_preamble {
2237 PreambleAction::Take
2238 } else {
2239 PreambleAction::Skip
2240 }
2241 } else {
2242 PreambleAction::Absent
2243 };
2244
2245 if rows_to_take == rows_avail && has_trailer {
2247 take_trailer = true;
2248 need_preamble = true;
2249 } else {
2250 need_preamble = false;
2251 };
2252
2253 chunk_instructions.push(Self {
2254 preamble,
2255 chunk_idx: block_index,
2256 rows_to_skip: to_skip,
2257 rows_to_take,
2258 take_trailer,
2259 });
2260
2261 to_skip = 0;
2262 block_index += 1;
2263 }
2264 }
2265
2266 if user_ranges.len() > 1 {
2270 let mut write = 0;
2274 for read in 1..chunk_instructions.len() {
2275 let merges = {
2276 let last = &chunk_instructions[write];
2277 let candidate = &chunk_instructions[read];
2278 last.chunk_idx == candidate.chunk_idx
2279 && last.rows_to_take + last.rows_to_skip == candidate.rows_to_skip
2280 };
2281 if merges {
2282 let rows_to_take = chunk_instructions[read].rows_to_take;
2283 let take_trailer = chunk_instructions[read].take_trailer;
2284 let last = &mut chunk_instructions[write];
2285 last.rows_to_take += rows_to_take;
2286 last.take_trailer |= take_trailer;
2287 } else {
2288 write += 1;
2289 if write != read {
2290 chunk_instructions.swap(write, read);
2291 }
2292 }
2293 }
2294 chunk_instructions.truncate(write + 1);
2295 }
2296 chunk_instructions
2297 }
2298
2299 fn drain_from_instruction(
2300 &self,
2301 rows_desired: &mut u64,
2302 need_preamble: &mut bool,
2303 skip_in_chunk: &mut u64,
2304 ) -> (ChunkDrainInstructions, bool) {
2305 debug_assert!(!*need_preamble || *skip_in_chunk == 0);
2307 let rows_avail = self.rows_to_take - *skip_in_chunk;
2308 let has_preamble = self.preamble != PreambleAction::Absent;
2309 let preamble_action = match (*need_preamble, has_preamble) {
2310 (true, true) => PreambleAction::Take,
2311 (true, false) => panic!("Need preamble but there isn't one"),
2312 (false, true) => PreambleAction::Skip,
2313 (false, false) => PreambleAction::Absent,
2314 };
2315
2316 let rows_taking = if *rows_desired >= rows_avail {
2319 *need_preamble = self.take_trailer;
2327 rows_avail
2328 } else {
2329 *need_preamble = false;
2332 *rows_desired
2333 };
2334 let rows_skipped = *skip_in_chunk;
2335
2336 let consumed_chunk = if *rows_desired >= rows_avail {
2338 *rows_desired -= rows_avail;
2339 *skip_in_chunk = 0;
2340 true
2341 } else {
2342 *skip_in_chunk += *rows_desired;
2343 *rows_desired = 0;
2344 false
2345 };
2346
2347 (
2348 ChunkDrainInstructions {
2349 chunk_instructions: self.clone(),
2350 rows_to_skip: rows_skipped,
2351 rows_to_take: rows_taking,
2352 preamble_action,
2353 },
2354 consumed_chunk,
2355 )
2356 }
2357}
2358
2359enum Words {
2360 U16(ScalarBuffer<u16>),
2361 U32(ScalarBuffer<u32>),
2362}
2363
2364struct WordsIter<'a> {
2365 iter: Box<dyn Iterator<Item = u32> + 'a>,
2366}
2367
2368impl Words {
2369 pub fn len(&self) -> usize {
2370 match self {
2371 Self::U16(b) => b.len(),
2372 Self::U32(b) => b.len(),
2373 }
2374 }
2375
2376 pub fn iter(&self) -> WordsIter<'_> {
2377 match self {
2378 Self::U16(buf) => WordsIter {
2379 iter: Box::new(buf.iter().map(|&x| x as u32)),
2380 },
2381 Self::U32(buf) => WordsIter {
2382 iter: Box::new(buf.iter().copied()),
2383 },
2384 }
2385 }
2386
2387 pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
2388 let bytes_per_value = if has_large_chunk { 4 } else { 2 };
2389 assert_eq!(bytes.len() % bytes_per_value, 0);
2390 let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
2391 if has_large_chunk {
2392 Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
2393 } else {
2394 Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
2395 }
2396 }
2397}
2398
2399impl<'a> Iterator for WordsIter<'a> {
2400 type Item = u32;
2401
2402 fn next(&mut self) -> Option<Self::Item> {
2403 self.iter.next()
2404 }
2405}
2406
2407struct FlatValueCounts {
2412 logs: Vec<u8>,
2413 uniform: bool,
2414 values_per_chunk: u64,
2415 last_chunk_values: u64,
2416}
2417
2418fn analyze_value_counts(words: &Words, items_in_page: u64) -> FlatValueCounts {
2419 let num_chunks = words.len();
2420 let logs = words.iter().map(|w| (w & 0x0F) as u8).collect::<Vec<_>>();
2421 let mut counted = 0u64;
2422 for &log in logs.iter().take(num_chunks.saturating_sub(1)) {
2423 debug_assert!(log > 0);
2425 counted += 1u64 << log;
2426 }
2427 let last_chunk_values = items_in_page - counted;
2428 if let Some(&last_log) = logs.last() {
2429 debug_assert!(last_log == 0 || (1u64 << last_log) == last_chunk_values);
2430 }
2431 let uniform = num_chunks <= 1 || logs[..num_chunks - 1].iter().all(|&log| log == logs[0]);
2432 let values_per_chunk = if num_chunks <= 1 {
2435 items_in_page.max(1)
2436 } else {
2437 1u64 << logs[0]
2438 };
2439 FlatValueCounts {
2440 logs,
2441 uniform,
2442 values_per_chunk,
2443 last_chunk_values,
2444 }
2445}
2446
2447fn flat_value_counts_iter(logs: &[u8], items_in_page: u64) -> impl Iterator<Item = u64> + '_ {
2450 let num_chunks = logs.len();
2451 let mut counted = 0u64;
2452 (0..num_chunks).map(move |i| {
2453 let count = if i + 1 < num_chunks {
2454 1u64 << logs[i]
2455 } else {
2456 items_in_page - counted
2457 };
2458 counted += count;
2459 count
2460 })
2461}
2462
2463fn build_chunk_index(
2468 words: &Words,
2469 items_in_page: u64,
2470 base: u64,
2471 data_buf_size: u64,
2472 rep_index_bytes: Option<&[u8]>,
2473 repetition_index_depth: u16,
2474) -> MiniBlockChunkIndex {
2475 let num_chunks = words.len();
2476 let byte_starts = PrefixSums::from_deltas(
2479 words
2480 .iter()
2481 .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64),
2482 num_chunks,
2483 data_buf_size,
2484 );
2485
2486 let rows = if let Some(rep_index_data) = rep_index_bytes {
2489 assert!(rep_index_data.len() % 8 == 0);
2490 let stride = repetition_index_depth as usize + 1;
2491 let (row_starts, has_trailer) = parse_nested_rep(rep_index_data, stride);
2492 let value_counts = analyze_value_counts(words, items_in_page);
2493 let item_counts = if value_counts.uniform {
2494 ItemCounts::Uniform {
2495 values_per_chunk: value_counts.values_per_chunk,
2496 last_chunk_values: value_counts.last_chunk_values,
2497 }
2498 } else {
2499 ItemCounts::PerChunkLog {
2500 logs: value_counts.logs,
2501 last_chunk_values: value_counts.last_chunk_values,
2502 }
2503 };
2504 RowMapping::Nested {
2505 row_starts,
2506 has_trailer,
2507 item_counts,
2508 }
2509 } else {
2510 let value_counts = analyze_value_counts(words, items_in_page);
2511 if value_counts.uniform {
2512 RowMapping::UniformFlat {
2513 values_per_chunk: value_counts.values_per_chunk,
2514 last_chunk_values: value_counts.last_chunk_values,
2515 num_chunks,
2516 }
2517 } else {
2518 let value_starts = PrefixSums::from_deltas(
2519 flat_value_counts_iter(&value_counts.logs, items_in_page),
2520 num_chunks,
2521 items_in_page,
2522 );
2523 RowMapping::Flat { value_starts }
2524 }
2525 };
2526
2527 MiniBlockChunkIndex::new(base, byte_starts, rows)
2528}
2529
2530impl StructuralPageScheduler for MiniBlockScheduler {
2531 fn initialize<'a>(
2532 &'a mut self,
2533 io: &Arc<dyn EncodingsIo>,
2534 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2535 let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
2539 let base = self.buffer_offsets_and_sizes[1].0;
2540 let data_buf_size = self.buffer_offsets_and_sizes[1].1;
2541 let mut bufs_needed = 1;
2542 if self.dictionary.is_some() {
2543 bufs_needed += 1;
2544 }
2545 if self.repetition_index_depth > 0 {
2546 bufs_needed += 1;
2547 }
2548 let mut required_ranges = Vec::with_capacity(bufs_needed);
2549 required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size);
2550 if let Some(ref dictionary) = self.dictionary {
2551 required_ranges.push(
2552 dictionary.dictionary_buf_position_and_size.0
2553 ..dictionary.dictionary_buf_position_and_size.0
2554 + dictionary.dictionary_buf_position_and_size.1,
2555 );
2556 }
2557 if self.repetition_index_depth > 0 {
2558 let (rep_index_pos, rep_index_size) = self.buffer_offsets_and_sizes.last().unwrap();
2559 required_ranges.push(*rep_index_pos..*rep_index_pos + *rep_index_size);
2560 }
2561 let io_req = io.submit_request(required_ranges, 0);
2562
2563 async move {
2564 let mut buffers = io_req.await?.into_iter().fuse();
2565 let meta_bytes = buffers.next().unwrap();
2566 let dictionary_bytes = self.dictionary.as_ref().and_then(|_| buffers.next());
2567 let rep_index_bytes = buffers.next();
2568
2569 let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
2570 let chunk_index = build_chunk_index(
2571 &words,
2572 self.items_in_page,
2573 base,
2574 data_buf_size,
2575 rep_index_bytes.as_deref(),
2576 self.repetition_index_depth,
2577 );
2578
2579 let dictionary = if let Some(ref mut dictionary) = self.dictionary {
2581 let dictionary_data = dictionary_bytes.unwrap();
2582 Some(Arc::new(dictionary.dictionary_decompressor.decompress(
2583 LanceBuffer::from_bytes(dictionary_data, dictionary.dictionary_data_alignment),
2584 dictionary.num_dictionary_items,
2585 )?))
2586 } else {
2587 None
2588 };
2589
2590 let page_meta = Arc::new(MiniBlockCacheableState {
2591 chunk_index,
2592 dictionary,
2593 });
2594 self.page_meta = Some(page_meta.clone());
2595 Ok(page_meta as Arc<dyn CachedPageData>)
2596 }
2597 .boxed()
2598 }
2599
2600 fn load(&mut self, data: &Arc<dyn CachedPageData>) {
2601 self.page_meta = Some(
2602 data.clone()
2603 .as_arc_any()
2604 .downcast::<MiniBlockCacheableState>()
2605 .unwrap(),
2606 );
2607 }
2608
2609 fn schedule_ranges(
2610 &self,
2611 ranges: &[Range<u64>],
2612 io: &Arc<dyn EncodingsIo>,
2613 ) -> Result<Vec<PageLoadTask>> {
2614 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2615
2616 let page_meta = self.page_meta.as_ref().unwrap();
2617
2618 let chunk_instructions =
2619 ChunkInstructions::schedule_instructions(&page_meta.chunk_index, ranges);
2620
2621 debug_assert_eq!(
2622 num_rows,
2623 chunk_instructions
2624 .iter()
2625 .map(|ci| ci.rows_to_take)
2626 .sum::<u64>()
2627 );
2628
2629 let chunks_needed = chunk_instructions
2630 .iter()
2631 .map(|ci| ci.chunk_idx)
2632 .unique()
2633 .collect::<Vec<_>>();
2634
2635 let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
2636 let chunk_ranges = loaded_chunks
2637 .iter()
2638 .map(|c| c.byte_range.clone())
2639 .collect::<Vec<_>>();
2640 let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
2641
2642 let rep_decompressor = self.rep_decompressor.clone();
2643 let def_decompressor = self.def_decompressor.clone();
2644 let value_decompressor = self.value_decompressor.clone();
2645 let num_buffers = self.num_buffers;
2646 let has_large_chunk = self.has_large_chunk;
2647 let dictionary = page_meta
2648 .dictionary
2649 .as_ref()
2650 .map(|dictionary| dictionary.clone());
2651 let def_meaning = self.def_meaning.clone();
2652
2653 let res = async move {
2654 let loaded_chunk_data = loaded_chunk_data.await?;
2655 for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
2656 loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
2657 }
2658
2659 Ok(Box::new(MiniBlockDecoder {
2660 rep_decompressor,
2661 def_decompressor,
2662 value_decompressor,
2663 def_meaning,
2664 loaded_chunks: VecDeque::from_iter(loaded_chunks),
2665 instructions: VecDeque::from(chunk_instructions),
2666 offset_in_current_chunk: 0,
2667 dictionary,
2668 num_rows,
2669 num_buffers,
2670 has_large_chunk,
2671 }) as Box<dyn StructuralPageDecoder>)
2672 }
2673 .boxed();
2674 let page_load_task = PageLoadTask {
2675 decoder_fut: res,
2676 num_rows,
2677 };
2678 Ok(vec![page_load_task])
2679 }
2680}
2681
2682#[derive(Debug, Clone, Copy)]
2683struct FullZipRepIndexDetails {
2684 buf_position: u64,
2685 bytes_per_value: u64, }
2687
2688#[derive(Debug)]
2689enum PerValueDecompressor {
2690 Fixed(Arc<dyn FixedPerValueDecompressor>),
2691 Variable(Arc<dyn VariablePerValueDecompressor>),
2692}
2693
2694#[derive(Debug)]
2695struct FullZipDecodeDetails {
2696 value_decompressor: PerValueDecompressor,
2697 def_meaning: Arc<[DefinitionInterpretation]>,
2698 ctrl_word_parser: ControlWordParser,
2699 max_rep: u16,
2700 max_visible_def: u16,
2701}
2702
2703#[derive(Debug, Clone)]
2715enum FullZipReadSource {
2716 Remote(Arc<dyn EncodingsIo>),
2718 PrefetchedPage { base_offset: u64, data: LanceBuffer },
2720}
2721
2722impl FullZipReadSource {
2723 fn fetch(
2727 &self,
2728 ranges: &[Range<u64>],
2729 priority: u64,
2730 ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
2731 match self {
2732 Self::Remote(io) => {
2733 let io = io.clone();
2734 let ranges = ranges.to_vec();
2735 async move {
2736 let data = io.submit_request(ranges, priority).await?;
2737 Ok(data
2738 .into_iter()
2739 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2740 .collect::<VecDeque<_>>())
2741 }
2742 .boxed()
2743 }
2744 Self::PrefetchedPage { base_offset, data } => {
2745 let base_offset = *base_offset;
2746 let data = data.clone();
2747 let page_end = base_offset + data.len() as u64;
2748 std::future::ready(
2749 ranges
2750 .iter()
2751 .map(|range| {
2752 if range.start > range.end
2753 || range.start < base_offset
2754 || range.end > page_end
2755 {
2756 return Err(Error::internal(format!(
2757 "Requested range {:?} is outside page range {}..{}",
2758 range, base_offset, page_end
2759 )));
2760 }
2761 let start = (range.start - base_offset) as usize;
2762 let len = (range.end - range.start) as usize;
2763 Ok(data.slice_with_length(start, len))
2764 })
2765 .collect::<Result<VecDeque<_>>>(),
2766 )
2767 .boxed()
2768 }
2769 }
2770 }
2771}
2772
2773#[derive(Debug)]
2781pub struct FullZipScheduler {
2782 data_buf_position: u64,
2783 data_buf_size: u64,
2784 rep_index: Option<FullZipRepIndexDetails>,
2785 priority: u64,
2786 rows_in_page: u64,
2787 bits_per_offset: u8,
2788 details: Arc<FullZipDecodeDetails>,
2789 cached_state: Option<Arc<FullZipCacheableState>>,
2791 enable_cache: bool,
2793}
2794
2795impl FullZipScheduler {
2796 fn try_new(
2797 buffer_offsets_and_sizes: &[(u64, u64)],
2798 priority: u64,
2799 rows_in_page: u64,
2800 layout: &pb21::FullZipLayout,
2801 decompressors: &dyn DecompressionStrategy,
2802 ) -> Result<Self> {
2803 let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
2804 let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
2805 let num_reps = rows_in_page + 1;
2806 let bytes_per_rep = len / num_reps;
2807 debug_assert_eq!(len % num_reps, 0);
2808 debug_assert!(
2809 bytes_per_rep == 1
2810 || bytes_per_rep == 2
2811 || bytes_per_rep == 4
2812 || bytes_per_rep == 8
2813 );
2814 FullZipRepIndexDetails {
2815 buf_position: *pos,
2816 bytes_per_value: bytes_per_rep,
2817 }
2818 });
2819
2820 let value_decompressor = match layout.details {
2821 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
2822 let decompressor = decompressors.create_fixed_per_value_decompressor(
2823 layout.value_compression.as_ref().unwrap(),
2824 )?;
2825 PerValueDecompressor::Fixed(decompressor.into())
2826 }
2827 Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
2828 let decompressor = decompressors.create_variable_per_value_decompressor(
2829 layout.value_compression.as_ref().unwrap(),
2830 )?;
2831 PerValueDecompressor::Variable(decompressor.into())
2832 }
2833 None => {
2834 panic!("Full-zip layout must have a `details` field");
2835 }
2836 };
2837 let ctrl_word_parser = ControlWordParser::new(
2838 layout.bits_rep.try_into().unwrap(),
2839 layout.bits_def.try_into().unwrap(),
2840 );
2841 let def_meaning = layout
2842 .layers
2843 .iter()
2844 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2845 .collect::<Vec<_>>();
2846
2847 let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
2848 let max_visible_def = def_meaning
2849 .iter()
2850 .filter(|d| !d.is_list())
2851 .map(|d| d.num_def_levels())
2852 .sum();
2853
2854 let bits_per_offset = match layout.details {
2855 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
2856 Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
2857 bits_per_offset as u8
2858 }
2859 None => panic!("Full-zip layout must have a `details` field"),
2860 };
2861
2862 let details = Arc::new(FullZipDecodeDetails {
2863 value_decompressor,
2864 def_meaning: def_meaning.into(),
2865 ctrl_word_parser,
2866 max_rep,
2867 max_visible_def,
2868 });
2869 Ok(Self {
2870 data_buf_position,
2871 data_buf_size,
2872 rep_index,
2873 details,
2874 priority,
2875 rows_in_page,
2876 bits_per_offset,
2877 cached_state: None,
2878 enable_cache: false,
2879 })
2880 }
2881
2882 fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
2883 if ranges.is_empty() {
2884 return false;
2885 }
2886 let mut expected_start = 0;
2887 for range in ranges {
2888 if range.start != expected_start || range.end > rows_in_page || range.end < range.start
2889 {
2890 return false;
2891 }
2892 expected_start = range.end;
2893 }
2894 expected_start == rows_in_page
2895 }
2896
2897 fn create_page_load_task(
2898 io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
2899 num_rows: u64,
2900 details: Arc<FullZipDecodeDetails>,
2901 bits_per_offset: u8,
2902 ) -> PageLoadTask {
2903 let load_task = async move {
2904 let buffers = io_future.await?;
2905 let data = buffers
2906 .into_iter()
2907 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2908 .collect::<VecDeque<_>>();
2909 Self::create_decoder(details, data, num_rows, bits_per_offset)
2910 }
2911 .boxed();
2912 PageLoadTask {
2913 decoder_fut: load_task,
2914 num_rows,
2915 }
2916 }
2917
2918 fn create_decoder(
2920 details: Arc<FullZipDecodeDetails>,
2921 data: VecDeque<LanceBuffer>,
2922 num_rows: u64,
2923 bits_per_offset: u8,
2924 ) -> Result<Box<dyn StructuralPageDecoder>> {
2925 match &details.value_decompressor {
2926 PerValueDecompressor::Fixed(decompressor) => {
2927 let bits_per_value = decompressor.bits_per_value();
2928 if bits_per_value % 8 != 0 {
2929 return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
2930 }
2931 let bytes_per_value = bits_per_value / 8;
2932 let total_bytes_per_value =
2933 bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
2934 if total_bytes_per_value == 0 {
2935 return Err(lance_core::Error::internal(
2936 "Invalid encoding: per-row byte width must be greater than 0",
2937 ));
2938 }
2939 Ok(Box::new(FixedFullZipDecoder {
2940 details,
2941 data,
2942 num_rows,
2943 offset_in_current: 0,
2944 bytes_per_value: bytes_per_value as usize,
2945 total_bytes_per_value,
2946 }) as Box<dyn StructuralPageDecoder>)
2947 }
2948 PerValueDecompressor::Variable(_decompressor) => {
2949 Ok(Box::new(VariableFullZipDecoder::new(
2950 details,
2951 data,
2952 num_rows,
2953 bits_per_offset,
2954 bits_per_offset,
2955 )?))
2956 }
2957 }
2958 }
2959
2960 fn extract_byte_ranges_from_pairs(
2963 buffer: LanceBuffer,
2964 bytes_per_value: u64,
2965 data_buf_position: u64,
2966 ) -> Vec<Range<u64>> {
2967 ByteUnpacker::new(buffer, bytes_per_value as usize)
2968 .chunks(2)
2969 .into_iter()
2970 .map(|mut c| {
2971 let start = c.next().unwrap() + data_buf_position;
2972 let end = c.next().unwrap() + data_buf_position;
2973 start..end
2974 })
2975 .collect::<Vec<_>>()
2976 }
2977
2978 fn extract_byte_ranges_from_cached(
2981 buffer: &LanceBuffer,
2982 ranges: &[Range<u64>],
2983 bytes_per_value: u64,
2984 data_buf_position: u64,
2985 ) -> Vec<Range<u64>> {
2986 ranges
2987 .iter()
2988 .map(|r| {
2989 let start_offset = (r.start * bytes_per_value) as usize;
2990 let end_offset = (r.end * bytes_per_value) as usize;
2991
2992 let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
2993 let start_val =
2994 ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
2995 .next()
2996 .unwrap();
2997
2998 let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
2999 let end_val =
3000 ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
3001 .next()
3002 .unwrap();
3003
3004 (data_buf_position + start_val)..(data_buf_position + end_val)
3005 })
3006 .collect()
3007 }
3008
3009 fn compute_rep_index_ranges(
3011 ranges: &[Range<u64>],
3012 rep_index: &FullZipRepIndexDetails,
3013 ) -> Vec<Range<u64>> {
3014 ranges
3015 .iter()
3016 .flat_map(|r| {
3017 let first_val_start =
3018 rep_index.buf_position + (r.start * rep_index.bytes_per_value);
3019 let first_val_end = first_val_start + rep_index.bytes_per_value;
3020 let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
3021 let last_val_end = last_val_start + rep_index.bytes_per_value;
3022 [first_val_start..first_val_end, last_val_start..last_val_end]
3023 })
3024 .collect()
3025 }
3026
3027 fn schedule_ranges_rep(
3029 &self,
3030 ranges: &[Range<u64>],
3031 io: &Arc<dyn EncodingsIo>,
3032 rep_index: FullZipRepIndexDetails,
3033 ) -> Result<Vec<PageLoadTask>> {
3034 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3035 let data_buf_position = self.data_buf_position;
3036 let priority = self.priority;
3037 let details = self.details.clone();
3038 let bits_per_offset = self.bits_per_offset;
3039
3040 if Self::covers_entire_page(ranges, self.rows_in_page) {
3041 let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
3042 let page_data = io.submit_single(full_range.clone(), priority);
3043 let load_task = async move {
3044 let page_data = page_data.await?;
3045 let source = FullZipReadSource::PrefetchedPage {
3046 base_offset: full_range.start,
3047 data: LanceBuffer::from_bytes(page_data, 1),
3048 };
3049 let read_ranges = vec![full_range];
3050 let data = source.fetch(&read_ranges, priority).await?;
3051 Self::create_decoder(details, data, num_rows, bits_per_offset)
3052 }
3053 .boxed();
3054 let page_load_task = PageLoadTask {
3055 decoder_fut: load_task,
3056 num_rows,
3057 };
3058 return Ok(vec![page_load_task]);
3059 }
3060
3061 if let Some(cached_state) = &self.cached_state {
3062 let byte_ranges = Self::extract_byte_ranges_from_cached(
3063 &cached_state.rep_index_buffer,
3064 ranges,
3065 rep_index.bytes_per_value,
3066 data_buf_position,
3067 );
3068 let io_future = io.submit_request(byte_ranges, priority);
3069 let page_load_task =
3070 Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
3071 return Ok(vec![page_load_task]);
3072 }
3073
3074 let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
3075 let rep_data = io.submit_request(rep_ranges, priority);
3076 let io_clone = io.clone();
3077 let load_task = async move {
3078 let rep_data = rep_data.await?;
3079 let rep_buffer = LanceBuffer::concat(
3080 &rep_data
3081 .into_iter()
3082 .map(|d| LanceBuffer::from_bytes(d, 1))
3083 .collect::<Vec<_>>(),
3084 );
3085 let byte_ranges = Self::extract_byte_ranges_from_pairs(
3086 rep_buffer,
3087 rep_index.bytes_per_value,
3088 data_buf_position,
3089 );
3090 let source = FullZipReadSource::Remote(io_clone);
3091 let data = source.fetch(&byte_ranges, priority).await?;
3092 Self::create_decoder(details, data, num_rows, bits_per_offset)
3093 }
3094 .boxed();
3095 let page_load_task = PageLoadTask {
3096 decoder_fut: load_task,
3097 num_rows,
3098 };
3099 Ok(vec![page_load_task])
3100 }
3101
3102 fn schedule_ranges_simple(
3106 &self,
3107 ranges: &[Range<u64>],
3108 io: &Arc<dyn EncodingsIo>,
3109 ) -> Result<Vec<PageLoadTask>> {
3110 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3112
3113 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
3114 unreachable!()
3115 };
3116
3117 let bits_per_value = decompressor.bits_per_value();
3119 if !bits_per_value.is_multiple_of(8) {
3120 return Err(Error::invalid_input_source(
3121 format!(
3122 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
3123 bits_per_value
3124 )
3125 .into(),
3126 ));
3127 }
3128 let bytes_per_value = bits_per_value / 8;
3129 let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
3130 let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
3131 let byte_ranges = ranges
3132 .iter()
3133 .map(|r| {
3134 debug_assert!(r.end <= self.rows_in_page);
3135 let start = self.data_buf_position + r.start * total_bytes_per_value;
3136 let end = self.data_buf_position + r.end * total_bytes_per_value;
3137 start..end
3138 })
3139 .collect::<Vec<_>>();
3140
3141 let io_future = io.submit_request(byte_ranges, self.priority);
3142 let page_load_task = Self::create_page_load_task(
3143 io_future,
3144 num_rows,
3145 self.details.clone(),
3146 self.bits_per_offset,
3147 );
3148 Ok(vec![page_load_task])
3149 }
3150}
3151
3152#[derive(Debug)]
3154struct FullZipCacheableState {
3155 rep_index_buffer: LanceBuffer,
3157}
3158
3159impl DeepSizeOf for FullZipCacheableState {
3160 fn deep_size_of_children(&self, _context: &mut Context) -> usize {
3161 self.rep_index_buffer.len()
3162 }
3163}
3164
3165impl CachedPageData for FullZipCacheableState {
3166 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3167 self
3168 }
3169}
3170
3171impl StructuralPageScheduler for FullZipScheduler {
3172 fn initialize<'a>(
3173 &'a mut self,
3174 io: &Arc<dyn EncodingsIo>,
3175 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
3176 if self.enable_cache
3177 && let Some(rep_index) = self.rep_index
3178 {
3179 let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
3180 let rep_index_range = rep_index.buf_position..(rep_index.buf_position + total_size);
3181 let io_clone = io.clone();
3182 return async move {
3183 let rep_index_data = io_clone.submit_request(vec![rep_index_range], 0).await?;
3184 let state = Arc::new(FullZipCacheableState {
3185 rep_index_buffer: LanceBuffer::from_bytes(rep_index_data[0].clone(), 1),
3186 });
3187 self.cached_state = Some(state.clone());
3188 Ok(state as Arc<dyn CachedPageData>)
3189 }
3190 .boxed();
3191 }
3192 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
3193 }
3194
3195 fn load(&mut self, cache: &Arc<dyn CachedPageData>) {
3199 if let Ok(cached_state) = cache
3201 .clone()
3202 .as_arc_any()
3203 .downcast::<FullZipCacheableState>()
3204 {
3205 self.cached_state = Some(cached_state);
3207 }
3208 }
3209
3210 fn schedule_ranges(
3211 &self,
3212 ranges: &[Range<u64>],
3213 io: &Arc<dyn EncodingsIo>,
3214 ) -> Result<Vec<PageLoadTask>> {
3215 if let Some(rep_index) = self.rep_index {
3216 self.schedule_ranges_rep(ranges, io, rep_index)
3217 } else {
3218 self.schedule_ranges_simple(ranges, io)
3219 }
3220 }
3221}
3222
3223#[derive(Debug)]
3231struct FixedFullZipDecoder {
3232 details: Arc<FullZipDecodeDetails>,
3233 data: VecDeque<LanceBuffer>,
3234 offset_in_current: usize,
3235 bytes_per_value: usize,
3236 total_bytes_per_value: usize,
3237 num_rows: u64,
3238}
3239
3240impl FixedFullZipDecoder {
3241 fn slice_next_task(&mut self, num_rows: u64) -> FullZipDecodeTaskItem {
3242 debug_assert!(num_rows > 0);
3243 let cur_buf = self.data.front_mut().unwrap();
3244 let start = self.offset_in_current;
3245 if self.details.ctrl_word_parser.has_rep() {
3246 let mut rows_started = 0;
3249 let mut num_items = 0;
3252 while self.offset_in_current < cur_buf.len() {
3253 let control = self.details.ctrl_word_parser.parse_desc(
3254 &cur_buf[self.offset_in_current..],
3255 self.details.max_rep,
3256 self.details.max_visible_def,
3257 );
3258 if control.is_new_row {
3259 if rows_started == num_rows {
3260 break;
3261 }
3262 rows_started += 1;
3263 }
3264 num_items += 1;
3265 if control.is_visible {
3266 self.offset_in_current += self.total_bytes_per_value;
3267 } else {
3268 self.offset_in_current += self.details.ctrl_word_parser.bytes_per_word();
3269 }
3270 }
3271
3272 let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
3273 if self.offset_in_current == cur_buf.len() {
3274 self.data.pop_front();
3275 self.offset_in_current = 0;
3276 }
3277
3278 FullZipDecodeTaskItem {
3279 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3280 data: task_slice,
3281 bits_per_value: self.bytes_per_value as u64 * 8,
3282 num_values: num_items,
3283 block_info: BlockInfo::new(),
3284 }),
3285 rows_in_buf: rows_started,
3286 }
3287 } else {
3288 let cur_buf = self.data.front_mut().unwrap();
3291 let bytes_avail = cur_buf.len() - self.offset_in_current;
3292 let offset_in_cur = self.offset_in_current;
3293
3294 let bytes_needed = num_rows as usize * self.total_bytes_per_value;
3295 let mut rows_taken = num_rows;
3296 let task_slice = if bytes_needed >= bytes_avail {
3297 self.offset_in_current = 0;
3298 rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
3299 self.data
3300 .pop_front()
3301 .unwrap()
3302 .slice_with_length(offset_in_cur, bytes_avail)
3303 } else {
3304 self.offset_in_current += bytes_needed;
3305 cur_buf.slice_with_length(offset_in_cur, bytes_needed)
3306 };
3307 FullZipDecodeTaskItem {
3308 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3309 data: task_slice,
3310 bits_per_value: self.bytes_per_value as u64 * 8,
3311 num_values: rows_taken,
3312 block_info: BlockInfo::new(),
3313 }),
3314 rows_in_buf: rows_taken,
3315 }
3316 }
3317 }
3318}
3319
3320impl StructuralPageDecoder for FixedFullZipDecoder {
3321 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3322 let mut task_data = Vec::with_capacity(self.data.len());
3323 let mut remaining = num_rows;
3324 while remaining > 0 {
3325 let task_item = self.slice_next_task(remaining);
3326 remaining -= task_item.rows_in_buf;
3327 task_data.push(task_item);
3328 }
3329 Ok(Box::new(FixedFullZipDecodeTask {
3330 details: self.details.clone(),
3331 data: task_data,
3332 bytes_per_value: self.bytes_per_value,
3333 num_rows: num_rows as usize,
3334 }))
3335 }
3336
3337 fn num_rows(&self) -> u64 {
3338 self.num_rows
3339 }
3340}
3341
3342#[derive(Debug)]
3347struct VariableFullZipDecoder {
3348 details: Arc<FullZipDecodeDetails>,
3349 decompressor: Arc<dyn VariablePerValueDecompressor>,
3350 data: LanceBuffer,
3351 offsets: LanceBuffer,
3352 rep: ScalarBuffer<u16>,
3353 def: ScalarBuffer<u16>,
3354 repdef_starts: Vec<usize>,
3355 data_starts: Vec<usize>,
3356 offset_starts: Vec<usize>,
3357 visible_item_counts: Vec<u64>,
3358 bits_per_offset: u8,
3359 current_idx: usize,
3360 num_rows: u64,
3361}
3362
3363impl VariableFullZipDecoder {
3364 fn new(
3365 details: Arc<FullZipDecodeDetails>,
3366 data: VecDeque<LanceBuffer>,
3367 num_rows: u64,
3368 in_bits_per_length: u8,
3369 out_bits_per_offset: u8,
3370 ) -> Result<Self> {
3371 let decompressor = match details.value_decompressor {
3372 PerValueDecompressor::Variable(ref d) => d.clone(),
3373 _ => unreachable!(),
3374 };
3375
3376 assert_eq!(in_bits_per_length % 8, 0);
3377 assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
3378
3379 let mut decoder = Self {
3380 details,
3381 decompressor,
3382 data: LanceBuffer::empty(),
3383 offsets: LanceBuffer::empty(),
3384 rep: LanceBuffer::empty().borrow_to_typed_slice(),
3385 def: LanceBuffer::empty().borrow_to_typed_slice(),
3386 bits_per_offset: out_bits_per_offset,
3387 repdef_starts: Vec::with_capacity(num_rows as usize + 1),
3388 data_starts: Vec::with_capacity(num_rows as usize + 1),
3389 offset_starts: Vec::with_capacity(num_rows as usize + 1),
3390 visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
3391 current_idx: 0,
3392 num_rows,
3393 };
3394
3395 decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
3416
3417 Ok(decoder)
3418 }
3419
3420 fn slice_batch_data_and_rebase_offsets_typed<T>(
3421 data: &LanceBuffer,
3422 offsets: &LanceBuffer,
3423 ) -> Result<(LanceBuffer, LanceBuffer)>
3424 where
3425 T: arrow_buffer::ArrowNativeType
3426 + Copy
3427 + PartialOrd
3428 + std::ops::Sub<Output = T>
3429 + std::fmt::Display
3430 + TryInto<usize>,
3431 {
3432 let offsets_slice = offsets.borrow_to_typed_slice::<T>();
3433 let offsets_slice = offsets_slice.as_ref();
3434 if offsets_slice.is_empty() {
3435 return Err(Error::internal(
3436 "Variable offsets cannot be empty".to_string(),
3437 ));
3438 }
3439
3440 let base = offsets_slice[0];
3441 let end = *offsets_slice.last().unwrap();
3442 if end < base {
3443 return Err(Error::internal(format!(
3444 "Invalid variable offsets: end ({end}) is less than base ({base})"
3445 )));
3446 }
3447
3448 let data_start = base.try_into().map_err(|_| {
3449 Error::internal(format!("Variable offset ({base}) does not fit into usize"))
3450 })?;
3451 let data_end = end.try_into().map_err(|_| {
3452 Error::internal(format!("Variable offset ({end}) does not fit into usize"))
3453 })?;
3454 if data_end > data.len() {
3455 return Err(Error::internal(format!(
3456 "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
3457 data.len()
3458 )));
3459 }
3460
3461 let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
3462 for &offset in offsets_slice {
3463 if offset < base {
3464 return Err(Error::internal(format!(
3465 "Invalid variable offsets: offset ({offset}) is less than base ({base})"
3466 )));
3467 }
3468 rebased_offsets.push(offset - base);
3469 }
3470
3471 let sliced_data = data.slice_with_length(data_start, data_end - data_start);
3472 let sliced_data = LanceBuffer::copy_slice(&sliced_data);
3474 let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
3475 Ok((sliced_data, rebased_offsets))
3476 }
3477
3478 fn slice_batch_data_and_rebase_offsets(
3479 data: &LanceBuffer,
3480 offsets: &LanceBuffer,
3481 bits_per_offset: u8,
3482 ) -> Result<(LanceBuffer, LanceBuffer)> {
3483 match bits_per_offset {
3484 32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
3485 64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
3486 _ => Err(Error::internal(format!(
3487 "Unsupported bits_per_offset={bits_per_offset}"
3488 ))),
3489 }
3490 }
3491
3492 fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
3499 let width = bits_per_offset as usize / 8;
3500 if data.len() < width {
3501 return Err(Error::corrupt_file_named(
3502 "variable_full_zip",
3503 format!(
3504 "truncated length prefix: {} byte(s) remain in the page buffer but a \
3505 {}-bit length prefix requires {}",
3506 data.len(),
3507 bits_per_offset,
3508 width
3509 ),
3510 ));
3511 }
3512 Ok(match bits_per_offset {
3513 8 => data[0] as u64,
3514 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
3515 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
3516 64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
3517 _ => unreachable!(),
3518 })
3519 }
3520
3521 fn unzip(
3522 &mut self,
3523 data: VecDeque<LanceBuffer>,
3524 in_bits_per_length: u8,
3525 out_bits_per_offset: u8,
3526 num_rows: u64,
3527 ) -> Result<()> {
3528 let mut rep = Vec::with_capacity(num_rows as usize);
3530 let mut def = Vec::with_capacity(num_rows as usize);
3531 let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
3532
3533 let bytes_per_offset = out_bits_per_offset as usize / 8;
3536 let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
3537 let mut offsets_data = Vec::with_capacity(bytes_offsets);
3538
3539 let bytes_per_length = in_bits_per_length as usize / 8;
3540 let bytes_lengths = bytes_per_length * num_rows as usize;
3541
3542 let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
3543 let mut unzipped_data =
3546 Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
3547
3548 let mut current_offset = 0_u64;
3549 let mut visible_item_count = 0_u64;
3550 for databuf in data.into_iter() {
3551 let mut databuf = databuf.as_ref();
3552 while !databuf.is_empty() {
3553 let data_start = unzipped_data.len();
3554 let offset_start = offsets_data.len();
3555 let repdef_start = rep.len().max(def.len());
3558 let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
3560 databuf,
3561 self.details.max_rep,
3562 self.details.max_visible_def,
3563 );
3564 self.details
3565 .ctrl_word_parser
3566 .parse(databuf, &mut rep, &mut def);
3567 databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
3568
3569 if ctrl_desc.is_new_row {
3570 self.repdef_starts.push(repdef_start);
3571 self.data_starts.push(data_start);
3572 self.offset_starts.push(offset_start);
3573 self.visible_item_counts.push(visible_item_count);
3574 }
3575 if ctrl_desc.is_visible {
3576 visible_item_count += 1;
3577 if ctrl_desc.is_valid_item {
3578 let length = Self::parse_length(databuf, in_bits_per_length)?;
3579 match out_bits_per_offset {
3580 32 => offsets_data
3581 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3582 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3583 _ => unreachable!(),
3584 };
3585 databuf = &databuf[bytes_per_offset..];
3586 unzipped_data.extend_from_slice(&databuf[..length as usize]);
3587 databuf = &databuf[length as usize..];
3588 current_offset += length;
3589 } else {
3590 match out_bits_per_offset {
3592 32 => offsets_data
3593 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3594 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3595 _ => unreachable!(),
3596 }
3597 }
3598 }
3599 }
3600 }
3601 self.repdef_starts.push(rep.len().max(def.len()));
3602 self.data_starts.push(unzipped_data.len());
3603 self.offset_starts.push(offsets_data.len());
3604 self.visible_item_counts.push(visible_item_count);
3605 match out_bits_per_offset {
3606 32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
3607 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3608 _ => unreachable!(),
3609 };
3610 self.rep = ScalarBuffer::from(rep);
3611 self.def = ScalarBuffer::from(def);
3612 self.data = LanceBuffer::from(unzipped_data);
3613 self.offsets = LanceBuffer::from(offsets_data);
3614 Ok(())
3615 }
3616}
3617
3618impl StructuralPageDecoder for VariableFullZipDecoder {
3619 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3620 let start = self.current_idx;
3621 let end = start + num_rows as usize;
3622
3623 let offset_start = self.offset_starts[start];
3624 let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
3625 let offsets = self
3626 .offsets
3627 .slice_with_length(offset_start, offset_end - offset_start);
3628 let (data, offsets) =
3630 Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
3631
3632 let repdef_start = self.repdef_starts[start];
3633 let repdef_end = self.repdef_starts[end];
3634 let rep = if self.rep.is_empty() {
3635 self.rep.clone()
3636 } else {
3637 self.rep.slice(repdef_start, repdef_end - repdef_start)
3638 };
3639 let def = if self.def.is_empty() {
3640 self.def.clone()
3641 } else {
3642 self.def.slice(repdef_start, repdef_end - repdef_start)
3643 };
3644
3645 let visible_item_counts_start = self.visible_item_counts[start];
3646 let visible_item_counts_end = self.visible_item_counts[end];
3647 let num_visible_items = visible_item_counts_end - visible_item_counts_start;
3648
3649 self.current_idx += num_rows as usize;
3650
3651 Ok(Box::new(VariableFullZipDecodeTask {
3652 details: self.details.clone(),
3653 decompressor: self.decompressor.clone(),
3654 data,
3655 offsets,
3656 bits_per_offset: self.bits_per_offset,
3657 num_visible_items,
3658 rep,
3659 def,
3660 }))
3661 }
3662
3663 fn num_rows(&self) -> u64 {
3664 self.num_rows
3665 }
3666}
3667
3668#[derive(Debug)]
3669struct VariableFullZipDecodeTask {
3670 details: Arc<FullZipDecodeDetails>,
3671 decompressor: Arc<dyn VariablePerValueDecompressor>,
3672 data: LanceBuffer,
3673 offsets: LanceBuffer,
3674 bits_per_offset: u8,
3675 num_visible_items: u64,
3676 rep: ScalarBuffer<u16>,
3677 def: ScalarBuffer<u16>,
3678}
3679
3680impl DecodePageTask for VariableFullZipDecodeTask {
3681 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3682 let block = VariableWidthBlock {
3683 data: self.data,
3684 offsets: self.offsets,
3685 bits_per_offset: self.bits_per_offset,
3686 num_values: self.num_visible_items,
3687 block_info: BlockInfo::new(),
3688 };
3689 let decomopressed = self.decompressor.decompress(block)?;
3690 let rep = if self.rep.is_empty() {
3691 None
3692 } else {
3693 Some(self.rep.to_vec())
3694 };
3695 let def = if self.def.is_empty() {
3696 None
3697 } else {
3698 Some(self.def.to_vec())
3699 };
3700 let unraveler = RepDefUnraveler::new(
3701 rep,
3702 def,
3703 self.details.def_meaning.clone(),
3704 self.num_visible_items,
3705 );
3706 Ok(DecodedPage {
3707 data: decomopressed,
3708 repdef: unraveler,
3709 })
3710 }
3711}
3712
3713#[derive(Debug)]
3714struct FullZipDecodeTaskItem {
3715 data: PerValueDataBlock,
3716 rows_in_buf: u64,
3717}
3718
3719#[derive(Debug)]
3722struct FixedFullZipDecodeTask {
3723 details: Arc<FullZipDecodeDetails>,
3724 data: Vec<FullZipDecodeTaskItem>,
3725 num_rows: usize,
3726 bytes_per_value: usize,
3727}
3728
3729impl DecodePageTask for FixedFullZipDecodeTask {
3730 fn decode(self: Box<Self>) -> Result<DecodedPage> {
3731 let estimated_size_bytes = self
3733 .data
3734 .iter()
3735 .map(|task_item| task_item.data.data_size() as usize)
3736 .sum::<usize>()
3737 * 2;
3738 let mut data_builder =
3739 DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3740
3741 if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3742 for task_item in self.data.into_iter() {
3746 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3747 unreachable!()
3748 };
3749 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3750 else {
3751 unreachable!()
3752 };
3753 debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3754 let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3755 data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3756 }
3757
3758 let unraveler = RepDefUnraveler::new(
3759 None,
3760 None,
3761 self.details.def_meaning.clone(),
3762 self.num_rows as u64,
3763 );
3764
3765 Ok(DecodedPage {
3766 data: data_builder.finish(),
3767 repdef: unraveler,
3768 })
3769 } else {
3770 let mut rep = Vec::with_capacity(self.num_rows);
3772 let mut def = Vec::with_capacity(self.num_rows);
3773
3774 for task_item in self.data.into_iter() {
3775 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3776 unreachable!()
3777 };
3778 let mut buf_slice = fixed_data.data.as_ref();
3779 let num_values = fixed_data.num_values as usize;
3780 let mut values = Vec::with_capacity(
3783 fixed_data.data.len()
3784 - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3785 );
3786 let mut visible_items = 0;
3787 for _ in 0..num_values {
3788 self.details
3790 .ctrl_word_parser
3791 .parse(buf_slice, &mut rep, &mut def);
3792 buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3793
3794 let is_visible = def
3795 .last()
3796 .map(|d| *d <= self.details.max_visible_def)
3797 .unwrap_or(true);
3798 if is_visible {
3799 values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3801 buf_slice = &buf_slice[self.bytes_per_value..];
3802 visible_items += 1;
3803 }
3804 }
3805
3806 let values_buf = LanceBuffer::from(values);
3808 let fixed_data = FixedWidthDataBlock {
3809 bits_per_value: self.bytes_per_value as u64 * 8,
3810 block_info: BlockInfo::new(),
3811 data: values_buf,
3812 num_values: visible_items,
3813 };
3814 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3815 else {
3816 unreachable!()
3817 };
3818 let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3819 data_builder.append(&decompressed, 0..visible_items);
3820 }
3821
3822 let repetition = if rep.is_empty() { None } else { Some(rep) };
3823 let definition = if def.is_empty() { None } else { Some(def) };
3824
3825 let unraveler = RepDefUnraveler::new(
3826 repetition,
3827 definition,
3828 self.details.def_meaning.clone(),
3829 self.num_rows as u64,
3830 );
3831 let data = data_builder.finish();
3832
3833 Ok(DecodedPage {
3834 data,
3835 repdef: unraveler,
3836 })
3837 }
3838 }
3839}
3840
3841#[derive(Debug)]
3842struct StructuralPrimitiveFieldSchedulingJob<'a> {
3843 scheduler: &'a StructuralPrimitiveFieldScheduler,
3844 ranges: Vec<Range<u64>>,
3845 page_idx: usize,
3846 range_idx: usize,
3847 global_row_offset: u64,
3848}
3849
3850impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3851 pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3852 Self {
3853 scheduler,
3854 ranges,
3855 page_idx: 0,
3856 range_idx: 0,
3857 global_row_offset: 0,
3858 }
3859 }
3860}
3861
3862impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3863 fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3864 if self.range_idx >= self.ranges.len() {
3865 return Ok(Vec::new());
3866 }
3867 let mut range = self.ranges[self.range_idx].clone();
3869 let priority = range.start;
3870
3871 let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3872 trace!(
3873 "Current range is {:?} and current page has {} rows",
3874 range, cur_page.num_rows
3875 );
3876 while cur_page.num_rows + self.global_row_offset <= range.start {
3878 self.global_row_offset += cur_page.num_rows;
3879 self.page_idx += 1;
3880 trace!("Skipping entire page of {} rows", cur_page.num_rows);
3881 cur_page = &self.scheduler.page_schedulers[self.page_idx];
3882 }
3883
3884 let mut ranges_in_page = Vec::new();
3888 while cur_page.num_rows + self.global_row_offset > range.start {
3889 range.start = range.start.max(self.global_row_offset);
3890 let start_in_page = range.start - self.global_row_offset;
3891 let end_in_page = start_in_page + (range.end - range.start);
3892 let end_in_page = end_in_page.min(cur_page.num_rows);
3893 let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3894
3895 ranges_in_page.push(start_in_page..end_in_page);
3896 if last_in_range {
3897 self.range_idx += 1;
3898 if self.range_idx == self.ranges.len() {
3899 break;
3900 }
3901 range = self.ranges[self.range_idx].clone();
3902 } else {
3903 break;
3904 }
3905 }
3906
3907 trace!(
3908 "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3909 ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3910 ranges_in_page.len(),
3911 cur_page.num_rows,
3912 priority,
3913 self.scheduler.column_index,
3914 cur_page.page_index,
3915 );
3916
3917 self.global_row_offset += cur_page.num_rows;
3918 self.page_idx += 1;
3919
3920 let page_decoders = cur_page
3921 .scheduler
3922 .schedule_ranges(&ranges_in_page, context.io())?;
3923
3924 let cur_path = context.current_path();
3925 page_decoders
3926 .into_iter()
3927 .map(|page_load_task| {
3928 let cur_path = cur_path.clone();
3929 let page_decoder = page_load_task.decoder_fut;
3930 let unloaded_page = async move {
3931 let page_decoder = page_decoder.await?;
3932 Ok(LoadedPageShard {
3933 decoder: page_decoder,
3934 path: cur_path,
3935 })
3936 }
3937 .boxed();
3938 Ok(ScheduledScanLine {
3939 decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3940 rows_scheduled: page_load_task.num_rows,
3941 })
3942 })
3943 .collect::<Result<Vec<_>>>()
3944 }
3945}
3946
3947#[derive(Debug)]
3948struct PageInfoAndScheduler {
3949 page_index: usize,
3950 num_rows: u64,
3951 scheduler: Box<dyn StructuralPageScheduler>,
3952}
3953
3954#[derive(Debug)]
3959pub struct StructuralPrimitiveFieldScheduler {
3960 page_schedulers: Vec<PageInfoAndScheduler>,
3961 column_index: u32,
3962 view_tag: String,
3968}
3969
3970impl StructuralPrimitiveFieldScheduler {
3971 pub fn try_new(
3972 column_info: &ColumnInfo,
3973 decompressors: &dyn DecompressionStrategy,
3974 cache_repetition_index: bool,
3975 target_field: &Field,
3976 ) -> Result<Self> {
3977 let page_schedulers = column_info
3978 .page_infos
3979 .iter()
3980 .enumerate()
3981 .map(|(page_index, page_info)| {
3982 Self::page_info_to_scheduler(
3983 page_info,
3984 page_index,
3985 decompressors,
3986 cache_repetition_index,
3987 target_field,
3988 )
3989 })
3990 .collect::<Result<Vec<_>>>()?;
3991 Ok(Self {
3992 page_schedulers,
3993 column_index: column_info.index,
3994 view_tag: format!("{:?}", target_field.data_type()),
3995 })
3996 }
3997
3998 fn page_layout_to_scheduler(
3999 page_info: &PageInfo,
4000 page_layout: &PageLayout,
4001 decompressors: &dyn DecompressionStrategy,
4002 cache_repetition_index: bool,
4003 target_field: &Field,
4004 ) -> Result<Box<dyn StructuralPageScheduler>> {
4005 use pb21::page_layout::Layout;
4006 Ok(match page_layout.layout.as_ref().expect_ok()? {
4007 Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
4008 &page_info.buffer_offsets_and_sizes,
4009 page_info.priority,
4010 mini_block.num_items,
4011 mini_block,
4012 decompressors,
4013 )?),
4014 Layout::SparseLayout(sparse_layout) => {
4015 Box::new(sparse::SparseStructuralScheduler::try_new(
4016 &page_info.buffer_offsets_and_sizes,
4017 page_info.priority,
4018 page_info.num_rows,
4019 target_field.data_type(),
4020 sparse_layout,
4021 decompressors,
4022 )?)
4023 }
4024 Layout::FullZipLayout(full_zip) => {
4025 let mut scheduler = FullZipScheduler::try_new(
4026 &page_info.buffer_offsets_and_sizes,
4027 page_info.priority,
4028 page_info.num_rows,
4029 full_zip,
4030 decompressors,
4031 )?;
4032 scheduler.enable_cache = cache_repetition_index;
4033 Box::new(scheduler)
4034 }
4035 Layout::ConstantLayout(constant_layout) => {
4036 let def_meaning = constant_layout
4037 .layers
4038 .iter()
4039 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4040 .collect::<Vec<_>>();
4041 let has_scalar_value = constant_layout.inline_value.is_some()
4042 || page_info.buffer_offsets_and_sizes.len() == 1
4043 || page_info.buffer_offsets_and_sizes.len() == 3;
4044 if has_scalar_value {
4045 Box::new(constant::ConstantPageScheduler::try_new(
4046 page_info.buffer_offsets_and_sizes.clone(),
4047 constant_layout.inline_value.clone(),
4048 target_field.data_type(),
4049 def_meaning.into(),
4050 )?) as Box<dyn StructuralPageScheduler>
4051 } else if def_meaning.len() == 1
4052 && def_meaning[0] == DefinitionInterpretation::NullableItem
4053 {
4054 Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
4055 } else {
4056 let rep_codec = LevelCodec::try_new(
4059 constant_layout.rep_compression.as_ref(),
4060 decompressors,
4061 )?;
4062 let def_codec = LevelCodec::try_new(
4063 constant_layout.def_compression.as_ref(),
4064 decompressors,
4065 )?;
4066
4067 Box::new(ComplexAllNullScheduler::new(
4068 page_info.buffer_offsets_and_sizes.clone(),
4069 def_meaning.into(),
4070 rep_codec,
4071 def_codec,
4072 constant_layout.num_rep_values,
4073 constant_layout.num_def_values,
4074 )) as Box<dyn StructuralPageScheduler>
4075 }
4076 }
4077 Layout::BlobLayout(blob) => {
4078 let inner_scheduler = Self::page_layout_to_scheduler(
4079 page_info,
4080 blob.inner_layout.as_ref().expect_ok()?.as_ref(),
4081 decompressors,
4082 cache_repetition_index,
4083 target_field,
4084 )?;
4085 let def_meaning = blob
4086 .layers
4087 .iter()
4088 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4089 .collect::<Vec<_>>();
4090 if matches!(target_field.data_type(), DataType::Struct(_)) {
4091 Box::new(BlobDescriptionPageScheduler::new(
4093 inner_scheduler,
4094 def_meaning.into(),
4095 ))
4096 } else {
4097 Box::new(BlobPageScheduler::new(
4099 inner_scheduler,
4100 page_info.priority,
4101 page_info.num_rows,
4102 def_meaning.into(),
4103 ))
4104 }
4105 }
4106 })
4107 }
4108
4109 fn page_info_to_scheduler(
4110 page_info: &PageInfo,
4111 page_index: usize,
4112 decompressors: &dyn DecompressionStrategy,
4113 cache_repetition_index: bool,
4114 target_field: &Field,
4115 ) -> Result<PageInfoAndScheduler> {
4116 let page_layout = page_info.encoding.as_structural();
4117 let scheduler = Self::page_layout_to_scheduler(
4118 page_info,
4119 page_layout,
4120 decompressors,
4121 cache_repetition_index,
4122 target_field,
4123 )?;
4124 Ok(PageInfoAndScheduler {
4125 page_index,
4126 num_rows: page_info.num_rows,
4127 scheduler,
4128 })
4129 }
4130}
4131
4132pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
4133 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
4134}
4135
4136pub struct NoCachedPageData;
4137
4138impl DeepSizeOf for NoCachedPageData {
4139 fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
4140 0
4141 }
4142}
4143impl CachedPageData for NoCachedPageData {
4144 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
4145 self
4146 }
4147}
4148
4149pub struct CachedFieldData {
4150 pages: Vec<Arc<dyn CachedPageData>>,
4151}
4152
4153impl DeepSizeOf for CachedFieldData {
4154 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
4155 self.pages.deep_size_of_children(ctx)
4156 }
4157}
4158
4159#[derive(Debug, Clone)]
4169pub struct FieldDataCacheKey {
4170 pub column_index: u32,
4171 pub view_tag: String,
4172}
4173
4174impl CacheKey for FieldDataCacheKey {
4175 type ValueType = CachedFieldData;
4176
4177 fn key(&self) -> std::borrow::Cow<'_, str> {
4178 format!("{}:{}", self.column_index, self.view_tag).into()
4179 }
4180
4181 fn type_name() -> &'static str {
4182 "FieldData"
4183 }
4184
4185 fn schema() -> CacheKeySchema {
4186 CacheKeySchema::new("lance.encoding.logical.primitive.field-data-key", 1)
4187 }
4188
4189 fn write_key(&self, builder: &mut KeyBuilder) {
4190 builder.write_u32(self.column_index);
4191 builder.write_str(&self.view_tag);
4192 }
4193}
4194
4195impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
4196 fn initialize<'a>(
4197 &'a mut self,
4198 _filter: &'a FilterExpression,
4199 context: &'a SchedulerContext,
4200 ) -> BoxFuture<'a, Result<()>> {
4201 let cache_key = FieldDataCacheKey {
4202 column_index: self.column_index,
4203 view_tag: self.view_tag.clone(),
4204 };
4205 let cache = context.cache().clone();
4206
4207 async move {
4208 if let Some(cached_data) = cache.get_with_key(&cache_key).await {
4209 self.page_schedulers
4210 .iter_mut()
4211 .zip(cached_data.pages.iter())
4212 .for_each(|(page_scheduler, cached_data)| {
4213 page_scheduler.scheduler.load(cached_data);
4214 });
4215 return Ok(());
4216 }
4217
4218 let page_data = self
4219 .page_schedulers
4220 .iter_mut()
4221 .map(|s| s.scheduler.initialize(context.io()))
4222 .collect::<FuturesOrdered<_>>();
4223
4224 let page_data = page_data.try_collect::<Vec<_>>().await?;
4225 let cached_data = Arc::new(CachedFieldData { pages: page_data });
4226 cache.insert_with_key(&cache_key, cached_data).await;
4227 Ok(())
4228 }
4229 .boxed()
4230 }
4231
4232 fn schedule_ranges<'a>(
4233 &'a self,
4234 ranges: &[Range<u64>],
4235 _filter: &FilterExpression,
4236 ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
4237 let ranges = ranges.to_vec();
4238 Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
4239 self, ranges,
4240 )))
4241 }
4242}
4243
4244#[derive(Debug)]
4247pub struct StructuralCompositeDecodeArrayTask {
4248 tasks: Vec<Box<dyn DecodePageTask>>,
4249 should_validate: bool,
4250 data_type: DataType,
4251}
4252
4253impl StructuralCompositeDecodeArrayTask {
4254 fn restore_validity(
4255 array: Arc<dyn Array>,
4256 unraveler: &mut CompositeRepDefUnraveler,
4257 ) -> Result<Arc<dyn Array>> {
4258 let validity = unraveler.unravel_validity(array.len())?;
4259 let Some(validity) = validity else {
4260 return Ok(array);
4261 };
4262 if array.data_type() == &DataType::Null {
4263 return Ok(array);
4265 }
4266 if validity.len() != array.len() {
4267 return Err(Error::invalid_input_source(
4268 format!(
4269 "Structural validity has {} entries for an array with {} values",
4270 validity.len(),
4271 array.len()
4272 )
4273 .into(),
4274 ));
4275 }
4276 Ok(make_array(unsafe {
4279 array
4280 .to_data()
4281 .into_builder()
4282 .nulls(Some(validity))
4283 .build_unchecked()
4284 }))
4285 }
4286}
4287
4288impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
4289 fn decode(self: Box<Self>) -> Result<DecodedArray> {
4290 let mut arrays = Vec::with_capacity(self.tasks.len());
4291 let mut unravelers = Vec::with_capacity(self.tasks.len());
4292 let mut data_size = 0u64;
4293 for task in self.tasks {
4294 let decoded = task.decode()?;
4295 data_size += decoded.data.data_size();
4296 unravelers.push(decoded.repdef);
4297
4298 let array = make_array(
4299 decoded
4300 .data
4301 .into_arrow(self.data_type.clone(), self.should_validate)?,
4302 );
4303
4304 arrays.push(array);
4305 }
4306 let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
4307 let array = arrow_select::concat::concat(&array_refs)?;
4308 let mut repdef = CompositeRepDefUnraveler::new(unravelers);
4309
4310 let array = Self::restore_validity(array, &mut repdef)?;
4311
4312 Ok(DecodedArray {
4313 array,
4314 repdef,
4315 data_size,
4316 })
4317 }
4318}
4319
4320#[derive(Debug)]
4321pub struct StructuralPrimitiveFieldDecoder {
4322 field: Arc<ArrowField>,
4323 page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
4324 should_validate: bool,
4325 rows_drained_in_current: u64,
4326}
4327
4328impl StructuralPrimitiveFieldDecoder {
4329 pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
4330 Self {
4331 field: field.clone(),
4332 page_decoders: VecDeque::new(),
4333 should_validate,
4334 rows_drained_in_current: 0,
4335 }
4336 }
4337}
4338
4339impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
4340 fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
4341 assert!(child.path.is_empty());
4342 self.page_decoders.push_back(child.decoder);
4343 Ok(())
4344 }
4345
4346 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
4347 let mut remaining = num_rows;
4348 let mut tasks = Vec::new();
4349 while remaining > 0 {
4350 let queued_pages = self.page_decoders.len();
4351 let Some(cur_page) = self.page_decoders.front_mut() else {
4352 return Err(Error::internal(format!(
4353 "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})",
4354 self.field.name(),
4355 self.field.data_type(),
4356 num_rows,
4357 remaining,
4358 self.rows_drained_in_current,
4359 queued_pages
4360 )));
4361 };
4362 let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
4363 let to_take = num_in_page.min(remaining);
4364
4365 let task = cur_page.drain(to_take)?;
4366 tasks.push(task);
4367
4368 if to_take == num_in_page {
4369 self.page_decoders.pop_front();
4370 self.rows_drained_in_current = 0;
4371 } else {
4372 self.rows_drained_in_current += to_take;
4373 }
4374
4375 remaining -= to_take;
4376 }
4377 Ok(Box::new(StructuralCompositeDecodeArrayTask {
4378 tasks,
4379 should_validate: self.should_validate,
4380 data_type: self.field.data_type().clone(),
4381 }))
4382 }
4383
4384 fn data_type(&self) -> &DataType {
4385 self.field.data_type()
4386 }
4387}
4388
4389struct SerializedFullZip {
4391 values: LanceBuffer,
4393 repetition_index: Option<LanceBuffer>,
4395}
4396
4397const MINIBLOCK_ALIGNMENT: usize = 8;
4412
4413#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4440pub(crate) enum MiniblockChunkSize {
4441 U16,
4442 U32,
4443}
4444
4445#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4446enum ComplexNullEncoding {
4447 RawLevels,
4448 CompressedLevels,
4449}
4450
4451#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4452enum FixedWidthDictionaryEncoding {
4453 Exclude64Bit,
4454 Include64Bit,
4455}
4456
4457trait PrimitivePageEncodingBehavior: Send + Sync + Debug {
4458 fn validate_field(&self, _field: &Field, _metadata: &HashMap<String, String>) -> Result<()> {
4459 Ok(())
4460 }
4461
4462 fn try_plan_pages(
4463 &self,
4464 _ctx: &PrimitivePlanContext<'_>,
4465 _arrays: &[ArrayRef],
4466 _normalized: &NormalizedStructuralPlan,
4467 _row_number: u64,
4468 _num_rows: u64,
4469 _num_values: u64,
4470 ) -> Result<Option<Vec<PrimitivePageData>>> {
4471 Ok(None)
4472 }
4473
4474 fn try_encode_page(
4475 &self,
4476 _ctx: &PrimitiveEncodeContext,
4477 page: PrimitivePageData,
4478 ) -> Result<PrimitiveEncodeAttempt> {
4479 Ok(PrimitiveEncodeAttempt::Unhandled(page))
4480 }
4481}
4482
4483#[derive(Debug, Clone)]
4486pub struct PrimitivePageEncoding {
4487 behavior: Arc<dyn PrimitivePageEncodingBehavior>,
4488}
4489
4490impl PrimitivePageEncoding {
4491 pub fn reject_sparse() -> Self {
4493 Self {
4494 behavior: Arc::new(RejectSparsePrimitiveEncoding),
4495 }
4496 }
4497
4498 pub fn constant() -> Self {
4500 Self {
4501 behavior: Arc::new(ConstantPrimitiveEncoding),
4502 }
4503 }
4504
4505 pub fn sparse(compression: Arc<dyn CompressionStrategy>) -> Self {
4507 Self {
4508 behavior: Arc::new(SparsePrimitiveEncoding { compression }),
4509 }
4510 }
4511
4512 pub fn dense_u16(compression: Arc<dyn CompressionStrategy>) -> Self {
4514 Self {
4515 behavior: Arc::new(DenseU16PrimitiveEncoding { compression }),
4516 }
4517 }
4518
4519 pub fn dense_u32(compression: Arc<dyn CompressionStrategy>) -> Self {
4521 Self {
4522 behavior: Arc::new(DenseU32PrimitiveEncoding { compression }),
4523 }
4524 }
4525}
4526
4527#[derive(Debug)]
4528struct RejectSparsePrimitiveEncoding;
4529
4530#[derive(Debug)]
4531struct ConstantPrimitiveEncoding;
4532
4533#[derive(Debug)]
4534struct SparsePrimitiveEncoding {
4535 compression: Arc<dyn CompressionStrategy>,
4536}
4537
4538#[derive(Debug)]
4539struct DenseU16PrimitiveEncoding {
4540 compression: Arc<dyn CompressionStrategy>,
4541}
4542
4543#[derive(Debug)]
4544struct DenseU32PrimitiveEncoding {
4545 compression: Arc<dyn CompressionStrategy>,
4546}
4547
4548pub struct PrimitiveStructuralEncoder {
4549 accumulation_queue: AccumulationQueue,
4551
4552 keep_original_array: bool,
4553 accumulated_repdefs: Vec<RepDefBuilder>,
4554 page_encodings: Arc<[PrimitivePageEncoding]>,
4555 column_index: u32,
4556 field: Field,
4557 encoding_metadata: Arc<HashMap<String, String>>,
4558}
4559
4560struct CompressedLevelsChunk {
4561 data: LanceBuffer,
4562 num_levels: u16,
4563}
4564
4565struct CompressedLevels {
4566 data: Vec<CompressedLevelsChunk>,
4567 compression: CompressiveEncoding,
4568 rep_index: Option<LanceBuffer>,
4569}
4570
4571struct SerializedMiniBlockPage {
4572 num_buffers: u64,
4573 data: LanceBuffer,
4574 metadata: LanceBuffer,
4575}
4576
4577#[derive(Debug, Clone, Copy)]
4578struct DictEncodingBudget {
4579 max_dict_entries: u32,
4580 max_encoded_size: usize,
4581}
4582
4583enum PrimitivePageStructure {
4584 Dense {
4585 repdef: SerializedRepDefs,
4586 single_row_miniblock_repdef_levels: Option<u64>,
4587 },
4588 Sparse {
4589 plan: sparse::SparseStructuralPlan,
4590 prepared_values: Option<sparse::writer::PreparedSparseValues>,
4591 },
4592}
4593
4594struct PrimitivePageData {
4596 arrays: Vec<ArrayRef>,
4598 structure: PrimitivePageStructure,
4600 row_number: u64,
4602 num_rows: u64,
4604}
4605
4606struct PrimitivePlanContext<'a> {
4607 column_idx: u32,
4608 field: &'a Field,
4609 encoding_metadata: &'a HashMap<String, String>,
4610}
4611
4612enum PrimitiveEncodeAttempt {
4613 Encoded(EncodedPage),
4614 Unhandled(PrimitivePageData),
4615}
4616
4617#[derive(Clone)]
4622struct PrimitiveEncodeContext {
4623 column_idx: u32,
4625 field: Field,
4626 encoding_metadata: Arc<HashMap<String, String>>,
4627 is_simple_validity: bool,
4628 has_repdef_info: bool,
4629}
4630
4631impl PrimitiveStructuralEncoder {
4632 pub fn try_new(
4633 options: &EncodingOptions,
4634 page_encodings: Arc<[PrimitivePageEncoding]>,
4635 column_index: u32,
4636 field: Field,
4637 encoding_metadata: Arc<HashMap<String, String>>,
4638 ) -> Result<Self> {
4639 for page_encoding in page_encodings.iter() {
4640 page_encoding
4641 .behavior
4642 .validate_field(&field, &encoding_metadata)?;
4643 }
4644 Ok(Self {
4645 accumulation_queue: AccumulationQueue::new(
4646 options.cache_bytes_per_column,
4647 column_index,
4648 options.keep_original_array,
4649 ),
4650 keep_original_array: options.keep_original_array,
4651 accumulated_repdefs: Vec::new(),
4652 column_index,
4653 page_encodings,
4654 field,
4655 encoding_metadata,
4656 })
4657 }
4658
4659 fn encode_page(
4660 page_encodings: &[PrimitivePageEncoding],
4661 ctx: &PrimitiveEncodeContext,
4662 mut page: PrimitivePageData,
4663 ) -> Result<EncodedPage> {
4664 for page_encoding in page_encodings {
4665 match page_encoding.behavior.try_encode_page(ctx, page)? {
4666 PrimitiveEncodeAttempt::Encoded(page) => return Ok(page),
4667 PrimitiveEncodeAttempt::Unhandled(unhandled) => page = unhandled,
4668 }
4669 }
4670 Err(Error::invalid_input_source(
4671 format!(
4672 "No primitive page encoding atom supports field '{}'",
4673 ctx.field.name
4674 )
4675 .into(),
4676 ))
4677 }
4678
4679 fn is_narrow(data_block: &DataBlock) -> bool {
4687 const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
4688
4689 if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
4690 let max_len_array = max_len_array
4691 .as_any()
4692 .downcast_ref::<PrimitiveArray<UInt64Type>>()
4693 .unwrap();
4694 if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
4695 return true;
4696 }
4697 }
4698 false
4699 }
4700
4701 fn prefers_miniblock(
4702 data_block: &DataBlock,
4703 encoding_metadata: &HashMap<String, String>,
4704 ) -> bool {
4705 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
4707 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
4708 }
4709 Self::is_narrow(data_block)
4711 }
4712
4713 fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
4714 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
4718 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
4719 }
4720 true
4721 }
4722
4723 fn serialize_miniblocks(
4770 miniblocks: MiniBlockCompressed,
4771 rep: Option<Vec<CompressedLevelsChunk>>,
4772 def: Option<Vec<CompressedLevelsChunk>>,
4773 miniblock_chunk_size: MiniblockChunkSize,
4774 ) -> Result<SerializedMiniBlockPage> {
4775 let bytes_rep = rep
4776 .as_ref()
4777 .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
4778 .unwrap_or(0);
4779 let bytes_def = def
4780 .as_ref()
4781 .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
4782 .unwrap_or(0);
4783 let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
4784 let mut num_buffers = miniblocks.data.len();
4785 if rep.is_some() {
4786 num_buffers += 1;
4787 }
4788 if def.is_some() {
4789 num_buffers += 1;
4790 }
4791 let max_extra = 9 * num_buffers;
4793 let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
4794 let chunk_size_bytes = match miniblock_chunk_size {
4795 MiniblockChunkSize::U16 => 2,
4796 MiniblockChunkSize::U32 => 4,
4797 };
4798 let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
4799
4800 let mut rep_iter = rep.map(|r| r.into_iter());
4801 let mut def_iter = def.map(|d| d.into_iter());
4802
4803 let mut buffer_offsets = vec![0; miniblocks.data.len()];
4804 for chunk in miniblocks.chunks {
4805 let start_pos = data_buffer.len();
4806 debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
4808
4809 let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
4810 let def = def_iter.as_mut().map(|d| d.next().unwrap());
4811
4812 let num_levels = rep
4814 .as_ref()
4815 .map(|r| r.num_levels)
4816 .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
4817 data_buffer.extend_from_slice(&num_levels.to_le_bytes());
4818
4819 if let Some(rep) = rep.as_ref() {
4821 let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
4822 Error::internal(format!(
4823 "Repetition buffer size ({} bytes) too large",
4824 rep.data.len()
4825 ))
4826 })?;
4827 data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
4828 }
4829 if let Some(def) = def.as_ref() {
4830 let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
4831 Error::internal(format!(
4832 "Definition buffer size ({} bytes) too large",
4833 def.data.len()
4834 ))
4835 })?;
4836 data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4837 }
4838
4839 if miniblock_chunk_size == MiniblockChunkSize::U32 {
4840 for &buffer_size in &chunk.buffer_sizes {
4841 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4842 }
4843 } else {
4844 for &buffer_size in &chunk.buffer_sizes {
4845 let buffer_size = u16::try_from(buffer_size).map_err(|_| {
4846 Error::internal(format!(
4847 "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
4848 buffer_size
4849 ))
4850 })?;
4851 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4852 }
4853 }
4854
4855 let add_padding = |data_buffer: &mut Vec<u8>| {
4857 let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4858 data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4859 };
4860 add_padding(&mut data_buffer);
4861
4862 if let Some(rep) = rep.as_ref() {
4864 data_buffer.extend_from_slice(&rep.data);
4865 add_padding(&mut data_buffer);
4866 }
4867 if let Some(def) = def.as_ref() {
4868 data_buffer.extend_from_slice(&def.data);
4869 add_padding(&mut data_buffer);
4870 }
4871 for (buffer_size, (buffer, buffer_offset)) in chunk
4872 .buffer_sizes
4873 .iter()
4874 .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4875 {
4876 let start = *buffer_offset;
4877 let end = start + *buffer_size as usize;
4878 *buffer_offset += *buffer_size as usize;
4879 data_buffer.extend_from_slice(&buffer[start..end]);
4880 add_padding(&mut data_buffer);
4881 }
4882
4883 let chunk_bytes = data_buffer.len() - start_pos;
4884 let max_chunk_size = match miniblock_chunk_size {
4885 MiniblockChunkSize::U16 => 32 * 1024,
4886 MiniblockChunkSize::U32 => 1_u64 << 31,
4887 };
4888 if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
4889 return Err(Error::internal(format!(
4890 "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
4891 chunk_bytes, max_chunk_size
4892 )));
4893 }
4894 if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
4895 return Err(Error::internal(format!(
4896 "Mini-block chunk size {} bytes is not aligned to {} bytes",
4897 chunk_bytes, MINIBLOCK_ALIGNMENT
4898 )));
4899 }
4900 if chunk.log_num_values > 15 {
4901 return Err(Error::internal(format!(
4902 "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
4903 chunk.log_num_values
4904 )));
4905 }
4906 let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4910 let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4911
4912 let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4913 if miniblock_chunk_size == MiniblockChunkSize::U32 {
4914 meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4915 } else {
4916 meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4917 }
4918 }
4919
4920 let data_buffer = LanceBuffer::from(data_buffer);
4921 let metadata_buffer = LanceBuffer::from(meta_buffer);
4922
4923 Ok(SerializedMiniBlockPage {
4924 num_buffers: miniblocks.data.len() as u64,
4925 data: data_buffer,
4926 metadata: metadata_buffer,
4927 })
4928 }
4929
4930 fn compress_levels(
4935 mut levels: RepDefSlicer<'_>,
4936 num_elements: u64,
4937 compression_strategy: &dyn CompressionStrategy,
4938 chunks: &[MiniBlockChunk],
4939 max_rep: u16,
4941 ) -> Result<CompressedLevels> {
4942 let mut rep_index = if max_rep > 0 {
4943 Vec::with_capacity(chunks.len())
4944 } else {
4945 vec![]
4946 };
4947 let num_levels = levels.num_levels() as u64;
4949 let levels_buf = levels.all_levels().clone();
4950
4951 let mut fixed_width_block = FixedWidthDataBlock {
4952 data: levels_buf,
4953 bits_per_value: 16,
4954 num_values: num_levels,
4955 block_info: BlockInfo::new(),
4956 };
4957 fixed_width_block.compute_stat();
4959
4960 let levels_block = DataBlock::FixedWidth(fixed_width_block);
4961 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4962 let (compressor, compressor_desc) =
4964 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4965 let mut level_chunks = Vec::with_capacity(chunks.len());
4967 let mut values_counter = 0;
4968 for (chunk_idx, chunk) in chunks.iter().enumerate() {
4969 let chunk_num_values = chunk.num_values(values_counter, num_elements);
4970 debug_assert!(chunk_num_values > 0);
4971 values_counter += chunk_num_values;
4972 let chunk_levels = if chunk_idx < chunks.len() - 1 {
4973 levels.slice_next(chunk_num_values as usize)
4974 } else {
4975 levels.slice_rest()
4976 };
4977 let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4978 if max_rep > 0 {
4979 let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4989 let rep_values = rep_values.as_ref();
4990
4991 let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4994 let num_leftovers = if chunk_idx < chunks.len() - 1 {
4995 rep_values
4996 .iter()
4997 .rev()
4998 .position(|v| *v == max_rep)
4999 .map(|pos| pos + 1)
5001 .unwrap_or(rep_values.len())
5002 } else {
5003 0
5005 };
5006
5007 if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
5008 let rep_len = rep_index.len();
5012 if rep_index[rep_len - 1] != 0 {
5013 rep_index[rep_len - 2] += 1;
5015 rep_index[rep_len - 1] = 0;
5016 }
5017 }
5018
5019 if chunk_idx == chunks.len() - 1 {
5020 num_rows += 1;
5022 }
5023 rep_index.push(num_rows as u64);
5024 rep_index.push(num_leftovers as u64);
5025 }
5026 let mut chunk_fixed_width = FixedWidthDataBlock {
5027 data: chunk_levels,
5028 bits_per_value: 16,
5029 num_values: num_chunk_levels,
5030 block_info: BlockInfo::new(),
5031 };
5032 chunk_fixed_width.compute_stat();
5033 let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
5034 let compressed_levels = compressor.compress(chunk_levels_block)?;
5035 let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
5036 Error::invalid_input_source(
5037 format!(
5038 "Mini-block cannot encode {} rep/def levels in one chunk. \
5039 This usually means a top-level row contains too much nested structure \
5040 for the current layout.",
5041 num_chunk_levels
5042 )
5043 .into(),
5044 )
5045 })?;
5046 level_chunks.push(CompressedLevelsChunk {
5047 data: compressed_levels,
5048 num_levels,
5049 });
5050 }
5051 debug_assert_eq!(levels.num_levels_remaining(), 0);
5052 let rep_index = if rep_index.is_empty() {
5053 None
5054 } else {
5055 Some(LanceBuffer::reinterpret_vec(rep_index))
5056 };
5057 Ok(CompressedLevels {
5058 data: level_chunks,
5059 compression: compressor_desc,
5060 rep_index,
5061 })
5062 }
5063
5064 fn encode_simple_all_null(
5065 column_idx: u32,
5066 num_rows: u64,
5067 row_number: u64,
5068 ) -> Result<EncodedPage> {
5069 let description =
5070 ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
5071 Ok(EncodedPage {
5072 column_idx,
5073 data: vec![],
5074 description: PageEncoding::Structural(description),
5075 num_rows,
5076 row_number,
5077 })
5078 }
5079
5080 fn encode_complex_all_null_vals(
5081 data: &Arc<[u16]>,
5082 compression_strategy: &dyn CompressionStrategy,
5083 ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
5084 let buffer = LanceBuffer::reinterpret_slice(data.clone());
5085 let mut fixed_width_block = FixedWidthDataBlock {
5086 data: buffer,
5087 bits_per_value: 16,
5088 num_values: data.len() as u64,
5089 block_info: BlockInfo::new(),
5090 };
5091 fixed_width_block.compute_stat();
5092
5093 let levels_block = DataBlock::FixedWidth(fixed_width_block);
5094 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
5095 let (compressor, encoding) =
5096 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
5097 let compressed_buffer = compressor.compress(levels_block)?;
5098 Ok((compressed_buffer, encoding))
5099 }
5100
5101 fn encode_complex_all_null(
5105 column_idx: u32,
5106 repdef: crate::repdef::SerializedRepDefs,
5107 row_number: u64,
5108 num_rows: u64,
5109 complex_null_encoding: ComplexNullEncoding,
5110 compression_strategy: &dyn CompressionStrategy,
5111 ) -> Result<EncodedPage> {
5112 if complex_null_encoding == ComplexNullEncoding::RawLevels {
5113 let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
5114 LanceBuffer::reinterpret_slice(rep.clone())
5115 } else {
5116 LanceBuffer::empty()
5117 };
5118
5119 let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
5120 LanceBuffer::reinterpret_slice(def.clone())
5121 } else {
5122 LanceBuffer::empty()
5123 };
5124
5125 let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
5126 return Ok(EncodedPage {
5127 column_idx,
5128 data: vec![rep_bytes, def_bytes],
5129 description: PageEncoding::Structural(description),
5130 num_rows,
5131 row_number,
5132 });
5133 }
5134
5135 let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
5136 repdef.repetition_levels.as_ref()
5137 {
5138 let num_values = rep.len() as u64;
5139 let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
5140 (buffer, Some(encoding), num_values)
5141 } else {
5142 (LanceBuffer::empty(), None, 0)
5143 };
5144
5145 let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
5146 repdef.definition_levels.as_ref()
5147 {
5148 let num_values = def.len() as u64;
5149 let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
5150 (buffer, Some(encoding), num_values)
5151 } else {
5152 (LanceBuffer::empty(), None, 0)
5153 };
5154
5155 let description = ProtobufUtils21::compressed_all_null_constant_layout(
5156 &repdef.def_meaning,
5157 rep_encoding,
5158 def_encoding,
5159 num_rep_values,
5160 num_def_values,
5161 );
5162 Ok(EncodedPage {
5163 column_idx,
5164 data: vec![rep_bytes, def_bytes],
5165 description: PageEncoding::Structural(description),
5166 num_rows,
5167 row_number,
5168 })
5169 }
5170
5171 fn leaf_validity(
5172 repdef: &crate::repdef::SerializedRepDefs,
5173 num_values: usize,
5174 ) -> Result<Option<BooleanBuffer>> {
5175 let rep = repdef
5176 .repetition_levels
5177 .as_ref()
5178 .map(|rep| rep.as_ref().to_vec());
5179 let def = repdef
5180 .definition_levels
5181 .as_ref()
5182 .map(|def| def.as_ref().to_vec());
5183 let mut unraveler = RepDefUnraveler::new(
5184 rep,
5185 def,
5186 repdef.def_meaning.clone().into(),
5187 num_values as u64,
5188 );
5189 if unraveler.is_all_valid() {
5190 return Ok(None);
5191 }
5192 let mut validity = BooleanBufferBuilder::new(num_values);
5193 unraveler.unravel_validity(&mut validity)?;
5194 Ok(Some(validity.finish()))
5195 }
5196
5197 fn is_constant_values(
5198 arrays: &[ArrayRef],
5199 scalar: &ArrayRef,
5200 validity: Option<&BooleanBuffer>,
5201 ) -> Result<bool> {
5202 debug_assert_eq!(scalar.len(), 1);
5203 debug_assert_eq!(scalar.null_count(), 0);
5204
5205 match scalar.data_type() {
5206 DataType::Boolean => {
5207 let mut global_idx = 0usize;
5208 let scalar_val = scalar.as_boolean().value(0);
5209 for arr in arrays {
5210 let bool_arr = arr.as_boolean();
5211 for i in 0..arr.len() {
5212 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
5213 global_idx += 1;
5214 if !is_valid {
5215 continue;
5216 }
5217 if bool_arr.value(i) != scalar_val {
5218 return Ok(false);
5219 }
5220 }
5221 }
5222 Ok(true)
5223 }
5224 DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
5225 DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
5226 DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
5227 DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
5228 data_type => {
5229 let mut global_idx = 0usize;
5230 let Some(byte_width) = data_type.byte_width_opt() else {
5231 return Ok(false);
5232 };
5233 let scalar_data = scalar.to_data();
5234 if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
5235 return Ok(false);
5236 }
5237 let scalar_bytes = scalar_data.buffers()[0].as_slice();
5238 if scalar_bytes.len() != byte_width {
5239 return Ok(false);
5240 }
5241
5242 for arr in arrays {
5243 let data = arr.to_data();
5244 if data.buffers().is_empty() {
5245 return Ok(false);
5246 }
5247 let buf = data.buffers()[0].as_slice();
5248 let base = data.offset();
5249 for i in 0..arr.len() {
5250 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
5251 global_idx += 1;
5252 if !is_valid {
5253 continue;
5254 }
5255 let start = (base + i) * byte_width;
5256 if buf[start..start + byte_width] != scalar_bytes[..] {
5257 return Ok(false);
5258 }
5259 }
5260 }
5261 Ok(true)
5262 }
5263 }
5264 }
5265
5266 fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
5267 arrays: &[ArrayRef],
5268 scalar: &ArrayRef,
5269 validity: Option<&BooleanBuffer>,
5270 ) -> Result<bool> {
5271 debug_assert_eq!(scalar.len(), 1);
5272 let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
5273 let mut global_idx = 0usize;
5274 for arr in arrays {
5275 let str_arr = arr.as_string::<O>();
5276 for i in 0..arr.len() {
5277 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
5278 global_idx += 1;
5279 if !is_valid {
5280 continue;
5281 }
5282 if str_arr.value(i).as_bytes() != scalar_val {
5283 return Ok(false);
5284 }
5285 }
5286 }
5287 Ok(true)
5288 }
5289
5290 fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
5291 arrays: &[ArrayRef],
5292 scalar: &ArrayRef,
5293 validity: Option<&BooleanBuffer>,
5294 ) -> Result<bool> {
5295 debug_assert_eq!(scalar.len(), 1);
5296 let scalar_val = scalar.as_binary::<O>().value(0);
5297 let mut global_idx = 0usize;
5298 for arr in arrays {
5299 let bin_arr = arr.as_binary::<O>();
5300 for i in 0..arr.len() {
5301 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
5302 global_idx += 1;
5303 if !is_valid {
5304 continue;
5305 }
5306 if bin_arr.value(i) != scalar_val {
5307 return Ok(false);
5308 }
5309 }
5310 }
5311 Ok(true)
5312 }
5313
5314 fn find_constant_scalar(
5315 arrays: &[ArrayRef],
5316 validity: Option<&BooleanBuffer>,
5317 ) -> Result<Option<ArrayRef>> {
5318 if arrays.is_empty() {
5319 return Ok(None);
5320 }
5321
5322 let global_scalar_idx = if let Some(validity) = validity {
5323 let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
5324 return Ok(None);
5325 };
5326 idx
5327 } else {
5328 0
5329 };
5330
5331 let mut idx_remaining = global_scalar_idx;
5332 let mut scalar_arr_idx = 0usize;
5333 while scalar_arr_idx < arrays.len() {
5334 let len = arrays[scalar_arr_idx].len();
5335 if idx_remaining < len {
5336 break;
5337 }
5338 idx_remaining -= len;
5339 scalar_arr_idx += 1;
5340 }
5341
5342 if scalar_arr_idx >= arrays.len() {
5343 return Ok(None);
5344 }
5345
5346 let scalar =
5347 lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
5348 if scalar.null_count() != 0 {
5349 return Ok(None);
5350 }
5351 if !Self::is_constant_values(arrays, &scalar, validity)? {
5352 return Ok(None);
5353 }
5354 Ok(Some(scalar))
5355 }
5356
5357 fn resolve_dict_values_compression_metadata(
5358 field_metadata: &HashMap<String, String>,
5359 env_compression: Option<String>,
5360 env_compression_level: Option<String>,
5361 ) -> HashMap<String, String> {
5362 let mut metadata = HashMap::new();
5363
5364 let compression = field_metadata
5365 .get(DICT_VALUES_COMPRESSION_META_KEY)
5366 .cloned()
5367 .or(env_compression)
5368 .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
5369 metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
5370
5371 if let Some(compression_level) = field_metadata
5372 .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
5373 .cloned()
5374 .or(env_compression_level)
5375 {
5376 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
5377 }
5378
5379 metadata
5380 }
5381
5382 fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
5383 let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
5388 dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
5389 &field.metadata,
5390 env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
5391 env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
5392 );
5393 Ok(dict_values_field)
5394 }
5395
5396 #[allow(clippy::too_many_arguments)]
5397 fn encode_miniblock(
5398 column_idx: u32,
5399 field: &Field,
5400 compression_strategy: &dyn CompressionStrategy,
5401 data: DataBlock,
5402 repdef: crate::repdef::SerializedRepDefs,
5403 row_number: u64,
5404 dictionary_data: Option<DataBlock>,
5405 num_rows: u64,
5406 miniblock_chunk_size: MiniblockChunkSize,
5407 ) -> Result<EncodedPage> {
5408 if let DataBlock::AllNull(_null_block) = data {
5409 unreachable!()
5412 }
5413
5414 let num_items = data.num_values();
5415
5416 let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
5417 let common_chunk_buffers =
5418 u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some());
5419 let support_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
5420 let compression_context =
5421 MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true);
5422 let (compressed_data, value_encoding) = compressor.compress(compression_context, data)?;
5423
5424 let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
5425
5426 let mut compressed_rep = repdef
5427 .rep_slicer()
5428 .map(|rep_slicer| {
5429 Self::compress_levels(
5430 rep_slicer,
5431 num_items,
5432 compression_strategy,
5433 &compressed_data.chunks,
5434 max_rep,
5435 )
5436 })
5437 .transpose()?;
5438
5439 let (rep_index, rep_index_depth) =
5440 match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
5441 Some(rep_index) => (Some(rep_index.clone()), 1),
5442 None => (None, 0),
5443 };
5444
5445 let mut compressed_def = repdef
5446 .def_slicer()
5447 .map(|def_slicer| {
5448 Self::compress_levels(
5449 def_slicer,
5450 num_items,
5451 compression_strategy,
5452 &compressed_data.chunks,
5453 0,
5454 )
5455 })
5456 .transpose()?;
5457
5458 let rep_data = compressed_rep
5464 .as_mut()
5465 .map(|cr| std::mem::take(&mut cr.data));
5466 let def_data = compressed_def
5467 .as_mut()
5468 .map(|cd| std::mem::take(&mut cd.data));
5469
5470 let serialized =
5471 Self::serialize_miniblocks(compressed_data, rep_data, def_data, miniblock_chunk_size)?;
5472 let has_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
5473
5474 let mut data = Vec::with_capacity(4);
5476 data.push(serialized.metadata);
5477 data.push(serialized.data);
5478
5479 if let Some(dictionary_data) = dictionary_data {
5480 let num_dictionary_items = dictionary_data.num_values();
5481 let dict_values_field = Self::build_dict_values_compressor_field(field)?;
5482
5483 let (compressor, dictionary_encoding) = compression_strategy
5484 .create_block_compressor(&dict_values_field, &dictionary_data)?;
5485 let dictionary_buffer = compressor.compress(dictionary_data)?;
5486
5487 data.push(dictionary_buffer);
5488 if let Some(rep_index) = rep_index {
5489 data.push(rep_index);
5490 }
5491
5492 let description = ProtobufUtils21::miniblock_layout(
5493 compressed_rep.map(|cr| cr.compression),
5494 compressed_def.map(|cd| cd.compression),
5495 value_encoding,
5496 rep_index_depth,
5497 serialized.num_buffers,
5498 Some((dictionary_encoding, num_dictionary_items)),
5499 &repdef.def_meaning,
5500 num_items,
5501 has_large_chunk,
5502 );
5503 Ok(EncodedPage {
5504 num_rows,
5505 column_idx,
5506 data,
5507 description: PageEncoding::Structural(description),
5508 row_number,
5509 })
5510 } else {
5511 let description = ProtobufUtils21::miniblock_layout(
5512 compressed_rep.map(|cr| cr.compression),
5513 compressed_def.map(|cd| cd.compression),
5514 value_encoding,
5515 rep_index_depth,
5516 serialized.num_buffers,
5517 None,
5518 &repdef.def_meaning,
5519 num_items,
5520 has_large_chunk,
5521 );
5522
5523 if let Some(rep_index) = rep_index {
5524 let view = rep_index.borrow_to_typed_slice::<u64>();
5525 let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
5526 debug_assert_eq!(total, num_rows);
5527
5528 data.push(rep_index);
5529 }
5530
5531 Ok(EncodedPage {
5532 num_rows,
5533 column_idx,
5534 data,
5535 description: PageEncoding::Structural(description),
5536 row_number,
5537 })
5538 }
5539 }
5540
5541 fn serialize_full_zip_fixed(
5543 fixed: FixedWidthDataBlock,
5544 mut repdef: ControlWordIterator,
5545 num_values: u64,
5546 ) -> Result<SerializedFullZip> {
5547 if !fixed.bits_per_value.is_multiple_of(8) {
5548 return Err(Error::invalid_input_source(
5549 format!(
5550 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
5551 fixed.bits_per_value
5552 )
5553 .into(),
5554 ));
5555 }
5556
5557 let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
5558 let mut zipped_data = Vec::with_capacity(len);
5559
5560 let max_rep_index_val = if repdef.has_repetition() {
5561 len as u64
5562 } else {
5563 0
5565 };
5566 let mut rep_index_builder =
5567 BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
5568
5569 let bytes_per_value = fixed.bits_per_value as usize / 8;
5570 let mut offset = 0;
5571
5572 if bytes_per_value == 0 {
5573 while let Some(control) = repdef.append_next(&mut zipped_data) {
5575 if control.is_new_row {
5576 debug_assert!(offset <= len);
5578 unsafe { rep_index_builder.append(offset as u64) };
5580 }
5581 offset = zipped_data.len();
5582 }
5583 } else {
5584 let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
5586 while let Some(control) = repdef.append_next(&mut zipped_data) {
5587 if control.is_new_row {
5588 debug_assert!(offset <= len);
5590 unsafe { rep_index_builder.append(offset as u64) };
5592 }
5593 if control.is_visible {
5594 let value = data_iter.next().unwrap();
5595 zipped_data.extend_from_slice(value);
5596 }
5597 offset = zipped_data.len();
5598 }
5599 }
5600
5601 debug_assert_eq!(zipped_data.len(), len);
5602 unsafe {
5605 rep_index_builder.append(zipped_data.len() as u64);
5606 }
5607
5608 let zipped_data = LanceBuffer::from(zipped_data);
5609 let rep_index = rep_index_builder.into_data();
5610 let rep_index = if rep_index.is_empty() {
5611 None
5612 } else {
5613 Some(LanceBuffer::from(rep_index))
5614 };
5615 Ok(SerializedFullZip {
5616 values: zipped_data,
5617 repetition_index: rep_index,
5618 })
5619 }
5620
5621 fn serialize_full_zip_variable(
5625 variable: VariableWidthBlock,
5626 mut repdef: ControlWordIterator,
5627 num_items: u64,
5628 ) -> Result<SerializedFullZip> {
5629 let bytes_per_offset = variable.bits_per_offset as usize / 8;
5630 if !variable.bits_per_offset.is_multiple_of(8) {
5631 return Err(Error::invalid_input_source(
5632 format!(
5633 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
5634 variable.bits_per_offset
5635 )
5636 .into(),
5637 ));
5638 }
5639 let len = variable.data.len()
5640 + repdef.bytes_per_word() * num_items as usize
5641 + bytes_per_offset * variable.num_values as usize;
5642 let mut buf = Vec::with_capacity(len);
5643
5644 let max_rep_index_val = len as u64;
5645 let mut rep_index_builder =
5646 BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
5647
5648 match bytes_per_offset {
5650 4 => {
5651 let offs = variable.offsets.borrow_to_typed_slice::<u32>();
5652 let mut rep_offset = 0;
5653 let mut windows_iter = offs.as_ref().windows(2);
5654 while let Some(control) = repdef.append_next(&mut buf) {
5655 if control.is_new_row {
5656 debug_assert!(rep_offset <= len);
5658 unsafe { rep_index_builder.append(rep_offset as u64) };
5660 }
5661 if control.is_visible {
5662 let window = windows_iter.next().unwrap();
5663 if control.is_valid_item {
5664 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
5665 buf.extend_from_slice(
5666 &variable.data[window[0] as usize..window[1] as usize],
5667 );
5668 }
5669 }
5670 rep_offset = buf.len();
5671 }
5672 }
5673 8 => {
5674 let offs = variable.offsets.borrow_to_typed_slice::<u64>();
5675 let mut rep_offset = 0;
5676 let mut windows_iter = offs.as_ref().windows(2);
5677 while let Some(control) = repdef.append_next(&mut buf) {
5678 if control.is_new_row {
5679 debug_assert!(rep_offset <= len);
5681 unsafe { rep_index_builder.append(rep_offset as u64) };
5683 }
5684 if control.is_visible {
5685 let window = windows_iter.next().unwrap();
5686 if control.is_valid_item {
5687 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
5688 buf.extend_from_slice(
5689 &variable.data[window[0] as usize..window[1] as usize],
5690 );
5691 }
5692 }
5693 rep_offset = buf.len();
5694 }
5695 }
5696 _ => {
5697 return Err(Error::invalid_input_source(
5698 format!(
5699 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
5700 variable.bits_per_offset
5701 )
5702 .into(),
5703 ));
5704 }
5705 }
5706
5707 debug_assert!(buf.len() <= len);
5710 unsafe {
5713 rep_index_builder.append(buf.len() as u64);
5714 }
5715
5716 let zipped_data = LanceBuffer::from(buf);
5717 let rep_index = rep_index_builder.into_data();
5718 debug_assert!(!rep_index.is_empty());
5719 let rep_index = Some(LanceBuffer::from(rep_index));
5720 Ok(SerializedFullZip {
5721 values: zipped_data,
5722 repetition_index: rep_index,
5723 })
5724 }
5725
5726 fn serialize_full_zip(
5729 compressed_data: PerValueDataBlock,
5730 repdef: ControlWordIterator,
5731 num_items: u64,
5732 ) -> Result<SerializedFullZip> {
5733 match compressed_data {
5734 PerValueDataBlock::Fixed(fixed) => {
5735 Self::serialize_full_zip_fixed(fixed, repdef, num_items)
5736 }
5737 PerValueDataBlock::Variable(var) => {
5738 Self::serialize_full_zip_variable(var, repdef, num_items)
5739 }
5740 }
5741 }
5742
5743 fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
5744 debug_assert_eq!(fixed.bits_per_value, 1);
5745 let num_values = fixed.num_values as usize;
5746 let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
5747 let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
5748 FixedWidthDataBlock {
5749 data: LanceBuffer::from(expanded),
5750 bits_per_value: 8,
5751 num_values: fixed.num_values,
5752 block_info: BlockInfo::new(),
5753 }
5754 }
5755
5756 fn encode_full_zip(
5757 column_idx: u32,
5758 field: &Field,
5759 compression_strategy: &dyn CompressionStrategy,
5760 data: DataBlock,
5761 repdef: crate::repdef::SerializedRepDefs,
5762 row_number: u64,
5763 num_lists: u64,
5764 ) -> Result<EncodedPage> {
5765 let max_rep = repdef
5766 .repetition_levels
5767 .as_ref()
5768 .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
5769 let max_def = repdef
5770 .definition_levels
5771 .as_ref()
5772 .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
5773
5774 let (num_items, num_visible_items) =
5778 if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
5779 (rep_levels.len() as u64, data.num_values())
5782 } else {
5783 (data.num_values(), data.num_values())
5785 };
5786
5787 let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
5788
5789 let repdef_iter = build_control_word_iterator(
5790 repdef.repetition_levels.as_deref(),
5791 max_rep,
5792 repdef.definition_levels.as_deref(),
5793 max_def,
5794 max_visible_def,
5795 num_items as usize,
5796 );
5797 let bits_rep = repdef_iter.bits_rep();
5798 let bits_def = repdef_iter.bits_def();
5799
5800 let data = match data {
5802 DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
5803 DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
5804 }
5805 other => other,
5806 };
5807
5808 let compressor = compression_strategy.create_per_value(field, &data)?;
5809 let (compressed_data, value_encoding) = compressor.compress(data)?;
5810
5811 let description = match &compressed_data {
5812 PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
5813 bits_rep,
5814 bits_def,
5815 fixed.bits_per_value as u32,
5816 value_encoding,
5817 &repdef.def_meaning,
5818 num_items as u32,
5819 num_visible_items as u32,
5820 ),
5821 PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
5822 bits_rep,
5823 bits_def,
5824 variable.bits_per_offset as u32,
5825 value_encoding,
5826 &repdef.def_meaning,
5827 num_items as u32,
5828 num_visible_items as u32,
5829 ),
5830 };
5831
5832 let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
5833
5834 let data = if let Some(repindex) = zipped.repetition_index {
5835 vec![zipped.values, repindex]
5836 } else {
5837 vec![zipped.values]
5838 };
5839
5840 Ok(EncodedPage {
5841 num_rows: num_lists,
5842 column_idx,
5843 data,
5844 description: PageEncoding::Structural(description),
5845 row_number,
5846 })
5847 }
5848
5849 fn should_dictionary_encode(
5850 data_block: &DataBlock,
5851 field: &Field,
5852 fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
5853 ) -> Option<DictEncodingBudget> {
5854 const DEFAULT_SAMPLE_SIZE: usize = 4096;
5855 const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5856
5857 match data_block {
5860 DataBlock::FixedWidth(fixed) => {
5861 if fixed.bits_per_value == 64
5862 && fixed_width_dictionary_encoding == FixedWidthDictionaryEncoding::Exclude64Bit
5863 {
5864 return None;
5865 }
5866 if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5867 return None;
5868 }
5869 if fixed.bits_per_value % 8 != 0 {
5870 return None;
5871 }
5872 }
5873 DataBlock::VariableWidth(var) => {
5874 if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5875 return None;
5876 }
5877 }
5878 _ => return None,
5879 }
5880
5881 let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5883 .ok()
5884 .and_then(|val| val.parse().ok())
5885 .unwrap_or(100);
5886 if data_block.num_values() < too_small {
5887 return None;
5888 }
5889
5890 let num_values = data_block.num_values();
5891
5892 let divisor: u64 = field
5895 .metadata
5896 .get(DICT_DIVISOR_META_KEY)
5897 .and_then(|val| val.parse().ok())
5898 .or_else(|| {
5899 env::var("LANCE_ENCODING_DICT_DIVISOR")
5900 .ok()
5901 .and_then(|val| val.parse().ok())
5902 })
5903 .unwrap_or(DEFAULT_DICT_DIVISOR);
5904
5905 let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5906 .ok()
5907 .and_then(|val| val.parse().ok())
5908 .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5909
5910 let threshold_cardinality = num_values
5911 .checked_div(divisor.max(1))
5912 .unwrap_or(0)
5913 .min(max_cardinality);
5914 if threshold_cardinality == 0 {
5915 return None;
5916 }
5917
5918 let threshold_ratio = field
5920 .metadata
5921 .get(DICT_SIZE_RATIO_META_KEY)
5922 .and_then(|val| val.parse::<f64>().ok())
5923 .or_else(|| {
5924 env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5925 .ok()
5926 .and_then(|val| val.parse().ok())
5927 })
5928 .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5929
5930 if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5931 panic!(
5932 "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5933 threshold_ratio
5934 );
5935 }
5936
5937 let data_size = data_block.data_size();
5938 if data_size == 0 {
5939 return None;
5940 }
5941
5942 let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5943 let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5944
5945 if let Some(sample_unique_ratio) =
5948 Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
5949 {
5950 if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
5951 return None;
5952 }
5953
5954 let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
5955 if projected_cardinality > threshold_cardinality {
5956 return None;
5957 }
5958 }
5959
5960 let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5961 Some(DictEncodingBudget {
5962 max_dict_entries,
5963 max_encoded_size,
5964 })
5965 }
5966
5967 fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
5975 use std::collections::HashSet;
5976
5977 const NUM_SAMPLE_BLOCKS: usize = 32;
5978 const MIN_RELIABLE_SAMPLES: usize = 1024;
5979
5980 let num_values = usize::try_from(data_block.num_values()).ok()?;
5981 if num_values == 0 {
5982 return Some(None);
5983 }
5984
5985 let sample_count = num_values.min(max_samples).max(1);
5986 if sample_count < MIN_RELIABLE_SAMPLES {
5987 return Some(None);
5988 }
5989
5990 let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
5991 let samples_per_block = (sample_count / block_count).max(1);
5992 let mut indices = Vec::with_capacity(sample_count);
5993 for block_idx in 0..block_count {
5994 let block_start = block_idx * num_values / block_count;
5995 let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
5996 let block_len = next_block_start.saturating_sub(block_start);
5997 let samples_in_block = samples_per_block.min(block_len);
5998 indices.extend((0..samples_in_block).map(|offset| block_start + offset));
5999 }
6000
6001 if indices.len() < MIN_RELIABLE_SAMPLES {
6002 return Some(None);
6003 }
6004
6005 let ratio = match data_block {
6006 DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
6007 64 => {
6008 let values = fixed.data.borrow_to_typed_slice::<u64>();
6009 let values = values.as_ref();
6010 let mut unique: HashSet<u64> =
6011 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
6012 for idx in indices.iter().copied() {
6013 unique.insert(values.get(idx).copied()?);
6014 }
6015 unique.len() as f64 / indices.len() as f64
6016 }
6017 128 => {
6018 let values = fixed.data.borrow_to_typed_slice::<u128>();
6019 let values = values.as_ref();
6020 let mut unique: HashSet<u128> =
6021 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
6022 for idx in indices.iter().copied() {
6023 unique.insert(values.get(idx).copied()?);
6024 }
6025 unique.len() as f64 / indices.len() as f64
6026 }
6027 _ => return Some(None),
6028 },
6029 DataBlock::VariableWidth(var) => {
6030 use xxhash_rust::xxh3::xxh3_64;
6031
6032 let mut unique: HashSet<u64> =
6034 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
6035 match var.bits_per_offset {
6036 32 => {
6037 let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
6038 let offsets: &[u32] = offsets_ref.as_ref();
6039 for i in indices.iter().copied() {
6040 let start = usize::try_from(*offsets.get(i)?).ok()?;
6041 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
6042 if start > end || end > var.data.len() {
6043 return None;
6044 }
6045 unique.insert(xxh3_64(&var.data[start..end]));
6046 }
6047 }
6048 64 => {
6049 let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
6050 let offsets: &[u64] = offsets_ref.as_ref();
6051 for i in indices.iter().copied() {
6052 let start = usize::try_from(*offsets.get(i)?).ok()?;
6053 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
6054 if start > end || end > var.data.len() {
6055 return None;
6056 }
6057 unique.insert(xxh3_64(&var.data[start..end]));
6058 }
6059 }
6060 _ => return Some(None),
6061 }
6062 unique.len() as f64 / indices.len() as f64
6063 }
6064 _ => return Some(None),
6065 };
6066
6067 Some(Some(ratio))
6068 }
6069
6070 fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
6071 let repetition_levels = repdef
6072 .repetition_levels
6073 .as_ref()
6074 .map(|levels| levels[range.clone()].to_vec());
6075 let definition_levels = repdef
6076 .definition_levels
6077 .as_ref()
6078 .map(|levels| levels[range].to_vec());
6079 SerializedRepDefs::new_with_fixed_size_list_levels(
6080 repetition_levels,
6081 definition_levels,
6082 repdef.def_meaning.clone(),
6083 repdef.has_fixed_size_list_levels(),
6084 )
6085 }
6086
6087 fn slice_arrays(
6088 arrays: &[ArrayRef],
6089 value_start: u64,
6090 num_values: u64,
6091 ) -> Result<Vec<ArrayRef>> {
6092 if num_values == 0 {
6093 return Ok(Vec::new());
6094 }
6095
6096 let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
6097 Error::invalid_input(format!("Value start {} is too large", value_start))
6098 })?;
6099 let mut values_remaining = usize::try_from(num_values).map_err(|_| {
6100 Error::invalid_input(format!("Value count {} is too large", num_values))
6101 })?;
6102 let mut sliced = Vec::new();
6103
6104 for array in arrays {
6105 if values_to_skip >= array.len() {
6106 values_to_skip -= array.len();
6107 continue;
6108 }
6109
6110 let offset = values_to_skip;
6111 let len = (array.len() - offset).min(values_remaining);
6112 sliced.push(array.slice(offset, len));
6113 values_remaining -= len;
6114 values_to_skip = 0;
6115
6116 if values_remaining == 0 {
6117 break;
6118 }
6119 }
6120
6121 if values_remaining != 0 {
6122 return Err(Error::internal(format!(
6123 "Page split requested {} values starting at {}, but the page did not contain enough values",
6124 num_values, value_start
6125 )));
6126 }
6127
6128 Ok(sliced)
6129 }
6130
6131 fn split_pages_for_miniblock_repdef_budget(
6132 arrays: Vec<ArrayRef>,
6133 repdef: SerializedRepDefs,
6134 budget: MiniBlockRepDefBudget,
6135 row_number: u64,
6136 num_rows: u64,
6137 ) -> Result<Vec<PrimitivePageData>> {
6138 if budget == MiniBlockRepDefBudget::WithinBudget {
6139 return Ok(vec![PrimitivePageData {
6140 arrays,
6141 structure: PrimitivePageStructure::Dense {
6142 repdef,
6143 single_row_miniblock_repdef_levels: None,
6144 },
6145 row_number,
6146 num_rows,
6147 }]);
6148 }
6149 if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget {
6150 return Ok(vec![PrimitivePageData {
6151 arrays,
6152 structure: PrimitivePageStructure::Dense {
6153 repdef,
6154 single_row_miniblock_repdef_levels: Some(num_levels),
6155 },
6156 row_number,
6157 num_rows,
6158 }]);
6159 }
6160
6161 let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else {
6162 unreachable!();
6163 };
6164
6165 let mut pages = Vec::with_capacity(splits.len());
6166 for split in splits {
6167 let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
6168 let repdef = Self::slice_repdef(&repdef, split.level_range);
6169 pages.push(PrimitivePageData {
6170 arrays,
6171 structure: PrimitivePageStructure::Dense {
6172 repdef,
6173 single_row_miniblock_repdef_levels: None,
6174 },
6175 row_number: row_number + split.row_start,
6176 num_rows: split.num_rows,
6177 });
6178 }
6179 Ok(pages)
6180 }
6181
6182 fn encode_dense_page(
6183 ctx: PrimitiveEncodeContext,
6184 page: PrimitivePageData,
6185 compression_strategy: Arc<dyn CompressionStrategy>,
6186 miniblock_chunk_size: MiniblockChunkSize,
6187 complex_null_encoding: ComplexNullEncoding,
6188 fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
6189 ) -> Result<EncodedPage> {
6190 let PrimitiveEncodeContext {
6191 column_idx,
6192 field,
6193 encoding_metadata,
6194 is_simple_validity,
6195 has_repdef_info,
6196 } = ctx;
6197 let PrimitivePageData {
6198 arrays,
6199 structure,
6200 row_number,
6201 num_rows,
6202 } = page;
6203 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
6204
6205 let (repdef, single_row_miniblock_repdef_levels) = match structure {
6206 PrimitivePageStructure::Dense {
6207 repdef,
6208 single_row_miniblock_repdef_levels,
6209 } => (repdef, single_row_miniblock_repdef_levels),
6210 PrimitivePageStructure::Sparse { .. } => {
6211 unreachable!("dense atom received sparse page")
6212 }
6213 };
6214
6215 if num_values == 0 {
6216 log::debug!(
6219 "Encoding column {} with {} items ({} rows) using complex-null layout",
6220 column_idx,
6221 num_values,
6222 num_rows
6223 );
6224 return Self::encode_complex_all_null(
6225 column_idx,
6226 repdef,
6227 row_number,
6228 num_rows,
6229 complex_null_encoding,
6230 compression_strategy.as_ref(),
6231 );
6232 }
6233
6234 let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
6235 let all_null = leaf_validity
6236 .as_ref()
6237 .map(|validity| validity.count_set_bits() == 0)
6238 .unwrap_or(false);
6239
6240 if all_null {
6241 return if is_simple_validity {
6242 log::debug!(
6243 "Encoding column {} with {} items ({} rows) using simple-null layout",
6244 column_idx,
6245 num_values,
6246 num_rows
6247 );
6248 Self::encode_simple_all_null(column_idx, num_values, row_number)
6249 } else {
6250 log::debug!(
6251 "Encoding column {} with {} items ({} rows) using complex-null layout",
6252 column_idx,
6253 num_values,
6254 num_rows
6255 );
6256 Self::encode_complex_all_null(
6257 column_idx,
6258 repdef,
6259 row_number,
6260 num_rows,
6261 complex_null_encoding,
6262 compression_strategy.as_ref(),
6263 )
6264 };
6265 }
6266
6267 if let DataType::Struct(fields) = &field.data_type()
6268 && fields.is_empty()
6269 {
6270 if has_repdef_info {
6271 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()));
6272 }
6273 return Self::encode_simple_all_null(column_idx, num_values, row_number);
6276 }
6277
6278 let data_block = DataBlock::from_arrays(&arrays, num_values);
6279
6280 if let Some(num_levels) = single_row_miniblock_repdef_levels {
6281 let requested_encoding = encoding_metadata
6282 .get(STRUCTURAL_ENCODING_META_KEY)
6283 .map(|requested| requested.to_lowercase());
6284 let fullzip_error = match &data_block {
6285 DataBlock::FixedWidth(fixed) if !fixed.bits_per_value.is_multiple_of(8) => {
6286 Some(format!(
6287 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
6288 fixed.bits_per_value
6289 ))
6290 }
6291 DataBlock::VariableWidth(variable)
6292 if !variable.bits_per_offset.is_multiple_of(8) =>
6293 {
6294 Some(format!(
6295 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
6296 variable.bits_per_offset
6297 ))
6298 }
6299 DataBlock::VariableWidth(variable)
6300 if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
6301 {
6302 Some(format!(
6303 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
6304 variable.bits_per_offset
6305 ))
6306 }
6307 DataBlock::Struct(struct_data_block)
6308 if !struct_data_block.has_variable_width_child() =>
6309 {
6310 Some(
6311 "Full-zip packed struct requires at least one variable-width child"
6312 .to_string(),
6313 )
6314 }
6315 DataBlock::Dictionary(_) => {
6316 Some("Full-zip does not encode dictionary data blocks directly".to_string())
6317 }
6318 DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
6319 Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
6320 Some(flat) => Some(format!(
6321 "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
6322 flat.bits_per_value
6323 )),
6324 None => Some(
6325 "Full-zip fixed-size-list capability requires a flat fixed-width child"
6326 .to_string(),
6327 ),
6328 },
6329 DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
6330 None
6331 }
6332 other => Some(format!(
6333 "Full-zip does not support value block type {}",
6334 other.name()
6335 )),
6336 };
6337 match requested_encoding.as_deref() {
6338 Some(STRUCTURAL_ENCODING_FULLZIP) => {
6339 if let Some(reason) = fullzip_error {
6340 return Err(Error::invalid_input_source(reason.into()));
6341 }
6342 return Self::encode_full_zip(
6343 column_idx,
6344 &field,
6345 compression_strategy.as_ref(),
6346 data_block,
6347 repdef,
6348 row_number,
6349 num_rows,
6350 );
6351 }
6352 Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
6353 if requested_encoding.is_none() && fullzip_error.is_none() {
6354 log::debug!(
6355 "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
6356 column_idx,
6357 num_values
6358 );
6359 return Self::encode_full_zip(
6360 column_idx,
6361 &field,
6362 compression_strategy.as_ref(),
6363 data_block,
6364 repdef,
6365 row_number,
6366 num_rows,
6367 );
6368 }
6369 return Err(Error::invalid_input_source(
6370 format!(
6371 "Mini-block cannot encode {} rep/def levels in one top-level row. \
6372 This usually means the row contains too much nested structure \
6373 for the current layout.",
6374 num_levels
6375 )
6376 .into(),
6377 ));
6378 }
6379 _ => {}
6380 }
6381 }
6382
6383 let requires_full_zip_packed_struct =
6384 if let DataBlock::Struct(ref struct_data_block) = data_block {
6385 struct_data_block.has_variable_width_child()
6386 } else {
6387 false
6388 };
6389
6390 if requires_full_zip_packed_struct {
6391 log::debug!(
6392 "Encoding column {} with {} items using full-zip packed struct layout",
6393 column_idx,
6394 num_values
6395 );
6396 return Self::encode_full_zip(
6397 column_idx,
6398 &field,
6399 compression_strategy.as_ref(),
6400 data_block,
6401 repdef,
6402 row_number,
6403 num_rows,
6404 );
6405 }
6406
6407 if let DataBlock::Dictionary(dict) = data_block {
6408 log::debug!(
6409 "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
6410 column_idx,
6411 num_values
6412 );
6413 let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
6414 indices_data_block.compute_stat();
6419 return Self::encode_miniblock(
6420 column_idx,
6421 &field,
6422 compression_strategy.as_ref(),
6423 indices_data_block,
6424 repdef,
6425 row_number,
6426 Some(dictionary_data_block),
6427 num_rows,
6428 miniblock_chunk_size,
6429 );
6430 }
6431
6432 let dict_result = Self::should_dictionary_encode(
6435 &data_block,
6436 &field,
6437 fixed_width_dictionary_encoding,
6438 )
6439 .and_then(|budget| {
6440 log::debug!(
6441 "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
6442 column_idx,
6443 num_values
6444 );
6445 dict::dictionary_encode(
6446 &data_block,
6447 budget.max_dict_entries,
6448 budget.max_encoded_size,
6449 )
6450 });
6451
6452 if let Some((indices_data_block, dictionary_data_block)) = dict_result {
6453 Self::encode_miniblock(
6454 column_idx,
6455 &field,
6456 compression_strategy.as_ref(),
6457 indices_data_block,
6458 repdef,
6459 row_number,
6460 Some(dictionary_data_block),
6461 num_rows,
6462 miniblock_chunk_size,
6463 )
6464 } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
6465 log::debug!(
6466 "Encoding column {} with {} items using mini-block layout",
6467 column_idx,
6468 num_values
6469 );
6470 Self::encode_miniblock(
6471 column_idx,
6472 &field,
6473 compression_strategy.as_ref(),
6474 data_block,
6475 repdef,
6476 row_number,
6477 None,
6478 num_rows,
6479 miniblock_chunk_size,
6480 )
6481 } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
6482 log::debug!(
6483 "Encoding column {} with {} items using full-zip layout",
6484 column_idx,
6485 num_values
6486 );
6487 Self::encode_full_zip(
6488 column_idx,
6489 &field,
6490 compression_strategy.as_ref(),
6491 data_block,
6492 repdef,
6493 row_number,
6494 num_rows,
6495 )
6496 } else {
6497 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()))
6498 }
6499 }
6500
6501 fn do_flush(
6503 &mut self,
6504 arrays: Vec<ArrayRef>,
6505 repdefs: Vec<RepDefBuilder>,
6506 row_number: u64,
6507 num_rows: u64,
6508 ) -> Result<Vec<EncodeTask>> {
6509 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
6510 let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
6511 let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
6512 let normalized = RepDefBuilder::normalize(repdefs);
6513 let plan_ctx = PrimitivePlanContext {
6514 column_idx: self.column_index,
6515 field: &self.field,
6516 encoding_metadata: &self.encoding_metadata,
6517 };
6518 let mut pages = None;
6519 for page_encoding in self.page_encodings.iter() {
6520 if let Some(planned) = page_encoding.behavior.try_plan_pages(
6521 &plan_ctx,
6522 &arrays,
6523 &normalized,
6524 row_number,
6525 num_rows,
6526 num_values,
6527 )? {
6528 pages = Some(planned);
6529 break;
6530 }
6531 }
6532 let pages = pages.ok_or_else(|| {
6533 Error::invalid_input_source(
6534 format!(
6535 "No primitive page planner supports field '{}'",
6536 self.field.name
6537 )
6538 .into(),
6539 )
6540 })?;
6541
6542 let mut tasks = Vec::with_capacity(pages.len());
6543 let ctx = PrimitiveEncodeContext {
6544 column_idx: self.column_index,
6545 field: self.field.clone(),
6546 encoding_metadata: self.encoding_metadata.clone(),
6547 is_simple_validity,
6548 has_repdef_info,
6549 };
6550 for page in pages {
6551 let ctx = ctx.clone();
6552 let page_encodings = self.page_encodings.clone();
6553 let task =
6554 spawn_cpu(move || Self::encode_page(page_encodings.as_ref(), &ctx, page)).boxed();
6555 tasks.push(task);
6556 }
6557 Ok(tasks)
6558 }
6559
6560 fn extract_validity_buf(
6561 array: Arc<dyn Array>,
6562 repdef: &mut RepDefBuilder,
6563 keep_original_array: bool,
6564 ) -> Result<Arc<dyn Array>> {
6565 if let Some(validity) = array.nulls() {
6566 if keep_original_array {
6567 repdef.add_validity_bitmap(validity.clone());
6568 } else {
6569 repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
6570 }
6571 let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
6572 Ok(make_array(data_no_nulls))
6573 } else {
6574 repdef.add_no_null(array.len());
6575 Ok(array)
6576 }
6577 }
6578
6579 fn extract_validity(
6580 mut array: Arc<dyn Array>,
6581 repdef: &mut RepDefBuilder,
6582 keep_original_array: bool,
6583 ) -> Result<Arc<dyn Array>> {
6584 match array.data_type() {
6585 DataType::Null => {
6586 repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
6587 Ok(array)
6588 }
6589 DataType::Dictionary(_, _) => {
6590 array = dict::normalize_dict_nulls(array)?;
6591 Self::extract_validity_buf(array, repdef, keep_original_array)
6592 }
6593 _ => Self::extract_validity_buf(array, repdef, keep_original_array),
6602 }
6603 }
6604}
6605
6606impl PrimitivePageEncodingBehavior for RejectSparsePrimitiveEncoding {
6607 fn validate_field(&self, field: &Field, metadata: &HashMap<String, String>) -> Result<()> {
6608 if metadata
6609 .get(STRUCTURAL_ENCODING_META_KEY)
6610 .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE))
6611 {
6612 return Err(Error::invalid_input_source(
6613 format!(
6614 "Field '{}' requests sparse structural encoding, which is not enabled by the selected file format",
6615 field.name
6616 )
6617 .into(),
6618 ));
6619 }
6620 Ok(())
6621 }
6622}
6623
6624fn plan_dense_primitive_pages(
6625 arrays: &[ArrayRef],
6626 normalized: &NormalizedStructuralPlan,
6627 row_number: u64,
6628 num_rows: u64,
6629 num_values: u64,
6630) -> Result<Vec<PrimitivePageData>> {
6631 let (repdef, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
6632 miniblock::max_repdef_levels_per_chunk,
6633 num_rows,
6634 num_values,
6635 )?;
6636 PrimitiveStructuralEncoder::split_pages_for_miniblock_repdef_budget(
6637 arrays.to_vec(),
6638 repdef,
6639 miniblock_repdef_budget,
6640 row_number,
6641 num_rows,
6642 )
6643}
6644
6645impl PrimitivePageEncodingBehavior for DenseU16PrimitiveEncoding {
6646 fn try_plan_pages(
6647 &self,
6648 _ctx: &PrimitivePlanContext<'_>,
6649 arrays: &[ArrayRef],
6650 normalized: &NormalizedStructuralPlan,
6651 row_number: u64,
6652 num_rows: u64,
6653 num_values: u64,
6654 ) -> Result<Option<Vec<PrimitivePageData>>> {
6655 Ok(Some(plan_dense_primitive_pages(
6656 arrays, normalized, row_number, num_rows, num_values,
6657 )?))
6658 }
6659
6660 fn try_encode_page(
6661 &self,
6662 ctx: &PrimitiveEncodeContext,
6663 page: PrimitivePageData,
6664 ) -> Result<PrimitiveEncodeAttempt> {
6665 if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
6666 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6667 }
6668 Ok(PrimitiveEncodeAttempt::Encoded(
6669 PrimitiveStructuralEncoder::encode_dense_page(
6670 ctx.clone(),
6671 page,
6672 self.compression.clone(),
6673 MiniblockChunkSize::U16,
6674 ComplexNullEncoding::RawLevels,
6675 FixedWidthDictionaryEncoding::Exclude64Bit,
6676 )?,
6677 ))
6678 }
6679}
6680
6681impl PrimitivePageEncodingBehavior for DenseU32PrimitiveEncoding {
6682 fn try_plan_pages(
6683 &self,
6684 _ctx: &PrimitivePlanContext<'_>,
6685 arrays: &[ArrayRef],
6686 normalized: &NormalizedStructuralPlan,
6687 row_number: u64,
6688 num_rows: u64,
6689 num_values: u64,
6690 ) -> Result<Option<Vec<PrimitivePageData>>> {
6691 Ok(Some(plan_dense_primitive_pages(
6692 arrays, normalized, row_number, num_rows, num_values,
6693 )?))
6694 }
6695
6696 fn try_encode_page(
6697 &self,
6698 ctx: &PrimitiveEncodeContext,
6699 page: PrimitivePageData,
6700 ) -> Result<PrimitiveEncodeAttempt> {
6701 if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
6702 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6703 }
6704 Ok(PrimitiveEncodeAttempt::Encoded(
6705 PrimitiveStructuralEncoder::encode_dense_page(
6706 ctx.clone(),
6707 page,
6708 self.compression.clone(),
6709 MiniblockChunkSize::U32,
6710 ComplexNullEncoding::CompressedLevels,
6711 FixedWidthDictionaryEncoding::Include64Bit,
6712 )?,
6713 ))
6714 }
6715}
6716
6717impl PrimitivePageEncodingBehavior for SparsePrimitiveEncoding {
6718 fn try_plan_pages(
6719 &self,
6720 ctx: &PrimitivePlanContext<'_>,
6721 arrays: &[ArrayRef],
6722 normalized: &NormalizedStructuralPlan,
6723 row_number: u64,
6724 num_rows: u64,
6725 num_values: u64,
6726 ) -> Result<Option<Vec<PrimitivePageData>>> {
6727 let requested_encoding = ctx.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY);
6728 let requests_sparse = requested_encoding
6729 .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE));
6730 if requests_sparse {
6731 let plan = sparse::writer::plan(normalized, num_values)?;
6732 if sparse::writer::uses_constant_layout(&plan, ctx.field) {
6733 return Ok(None);
6734 }
6735 return Ok(Some(vec![PrimitivePageData {
6736 arrays: arrays.to_vec(),
6737 structure: PrimitivePageStructure::Sparse {
6738 plan,
6739 prepared_values: None,
6740 },
6741 row_number,
6742 num_rows,
6743 }]));
6744 }
6745
6746 let (_, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
6747 miniblock::max_repdef_levels_per_chunk,
6748 num_rows,
6749 num_values,
6750 )?;
6751 let automatic_sparse = layout::select_automatic_sparse(
6752 requested_encoding.map(String::as_str),
6753 &miniblock_repdef_budget,
6754 || {
6755 let data = DataBlock::from_arrays(arrays, num_values);
6756 if !sparse::writer::supports_value_block(&data) {
6757 return Ok(None);
6758 }
6759 let prepared_values = match sparse::writer::prepare_values(
6760 ctx.field,
6761 self.compression.as_ref(),
6762 data,
6763 MiniblockChunkSize::U32,
6764 ) {
6765 Ok(prepared_values) => prepared_values,
6766 Err(error) => {
6767 debug!(
6768 "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}",
6769 ctx.column_idx, error
6770 );
6771 return Ok(None);
6772 }
6773 };
6774 let plan = sparse::writer::plan(normalized, num_values)?;
6775 if sparse::writer::uses_constant_layout(&plan, ctx.field) {
6776 return Ok(None);
6777 }
6778 Ok(Some((plan, prepared_values)))
6779 },
6780 )?;
6781 Ok(automatic_sparse.map(|(plan, prepared_values)| {
6782 vec![PrimitivePageData {
6783 arrays: arrays.to_vec(),
6784 structure: PrimitivePageStructure::Sparse {
6785 plan,
6786 prepared_values: Some(prepared_values),
6787 },
6788 row_number,
6789 num_rows,
6790 }]
6791 }))
6792 }
6793
6794 fn try_encode_page(
6795 &self,
6796 ctx: &PrimitiveEncodeContext,
6797 page: PrimitivePageData,
6798 ) -> Result<PrimitiveEncodeAttempt> {
6799 if !matches!(&page.structure, PrimitivePageStructure::Sparse { .. }) {
6800 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6801 }
6802 let PrimitivePageData {
6803 arrays,
6804 structure:
6805 PrimitivePageStructure::Sparse {
6806 plan,
6807 prepared_values,
6808 },
6809 row_number,
6810 num_rows,
6811 } = page
6812 else {
6813 unreachable!()
6814 };
6815 let num_values = arrays.iter().map(|array| array.len() as u64).sum();
6816 log::debug!(
6817 "Encoding column {} with {} visible items ({} rows) using sparse layout",
6818 ctx.column_idx,
6819 num_values,
6820 num_rows
6821 );
6822 Ok(PrimitiveEncodeAttempt::Encoded(
6823 sparse::writer::encode_page(
6824 ctx.column_idx,
6825 &ctx.field,
6826 self.compression.as_ref(),
6827 prepared_values.map_or_else(
6828 || {
6829 sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays(
6830 &arrays, num_values,
6831 ))
6832 },
6833 sparse::writer::SparseValueInput::Prepared,
6834 ),
6835 plan,
6836 row_number,
6837 num_rows,
6838 MiniblockChunkSize::U32,
6839 )?,
6840 ))
6841 }
6842}
6843
6844impl PrimitivePageEncodingBehavior for ConstantPrimitiveEncoding {
6845 fn try_encode_page(
6846 &self,
6847 ctx: &PrimitiveEncodeContext,
6848 page: PrimitivePageData,
6849 ) -> Result<PrimitiveEncodeAttempt> {
6850 let PrimitivePageStructure::Dense { repdef, .. } = &page.structure else {
6851 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6852 };
6853 let num_values: u64 = page.arrays.iter().map(|array| array.len() as u64).sum();
6854 if num_values == 0 {
6855 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6856 }
6857 let leaf_validity = PrimitiveStructuralEncoder::leaf_validity(repdef, num_values as usize)?;
6858 if leaf_validity
6859 .as_ref()
6860 .is_some_and(|validity| validity.count_set_bits() == 0)
6861 || matches!(ctx.field.data_type(), DataType::Struct(fields) if fields.is_empty())
6862 {
6863 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6864 }
6865 let Some(scalar) =
6866 PrimitiveStructuralEncoder::find_constant_scalar(&page.arrays, leaf_validity.as_ref())?
6867 else {
6868 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
6869 };
6870 let PrimitivePageData {
6871 structure: PrimitivePageStructure::Dense { repdef, .. },
6872 row_number,
6873 num_rows,
6874 ..
6875 } = page
6876 else {
6877 unreachable!()
6878 };
6879 log::debug!(
6880 "Encoding column {} with {} items ({} rows) using constant layout",
6881 ctx.column_idx,
6882 num_values,
6883 num_rows
6884 );
6885 Ok(PrimitiveEncodeAttempt::Encoded(
6886 constant::encode_constant_page(ctx.column_idx, scalar, repdef, row_number, num_rows)?,
6887 ))
6888 }
6889}
6890
6891impl FieldEncoder for PrimitiveStructuralEncoder {
6892 fn maybe_encode(
6894 &mut self,
6895 array: ArrayRef,
6896 _external_buffers: &mut OutOfLineBuffers,
6897 mut repdef: RepDefBuilder,
6898 row_number: u64,
6899 num_rows: u64,
6900 ) -> Result<Vec<EncodeTask>> {
6901 let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
6902 self.accumulated_repdefs.push(repdef);
6903
6904 if let Some((arrays, row_number, num_rows)) =
6905 self.accumulation_queue.insert(array, row_number, num_rows)
6906 {
6907 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
6908 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
6909 } else {
6910 Ok(vec![])
6911 }
6912 }
6913
6914 fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
6916 if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
6917 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
6918 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
6919 } else {
6920 Ok(vec![])
6921 }
6922 }
6923
6924 fn num_columns(&self) -> u32 {
6925 1
6926 }
6927
6928 fn finish(
6929 &mut self,
6930 _external_buffers: &mut OutOfLineBuffers,
6931 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
6932 std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
6933 }
6934}
6935
6936#[cfg(test)]
6937#[allow(clippy::single_range_in_vec_init)]
6938mod tests {
6939 use super::{
6940 ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
6941 FixedWidthDataBlock, FixedWidthDictionaryEncoding, FullZipCacheableState,
6942 FullZipDecodeDetails, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler,
6943 LazyLevels, LevelCodec, LevelCursor, LevelPlan, MiniBlockChunk, MiniBlockChunkIndex,
6944 MiniBlockCompressed, MiniblockChunkSize, PerValueDecompressor, PreambleAction,
6945 RunEndsBuilder, RunPosition, RunStorage, StructuralPageScheduler, VariableFullZipDecoder,
6946 dense_levels_from_block, validate_complex_all_null_levels,
6947 };
6948 use crate::buffer::LanceBuffer;
6949 use crate::compression::{BlockCompressor, DefaultDecompressionStrategy};
6950 use crate::constants::{
6951 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
6952 DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
6953 STRUCTURAL_ENCODING_MINIBLOCK,
6954 };
6955 use crate::data::BlockInfo;
6956 use crate::decoder::{PageEncoding, StructuralFieldDecoder};
6957 use crate::encodings::logical::primitive::{
6958 ChunkDrainInstructions, PrimitiveStructuralEncoder, StructuralPrimitiveFieldDecoder,
6959 };
6960 use crate::encodings::physical::rle::{RleDecompressor, RleEncoder, RleRuns, RunLengthWidth};
6961 use crate::format::ProtobufUtils21;
6962 use crate::format::pb21;
6963 use crate::format::pb21::compressive_encoding::Compression;
6964 use crate::repdef::build_control_word_iterator;
6965 use crate::testing::TestEncoding;
6966 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6967 use arrow_array::{ArrayRef, Int8Array, StringArray};
6968 use arrow_buffer::ScalarBuffer;
6969 use arrow_schema::{DataType, Field as ArrowField};
6970 use std::collections::HashMap;
6971 use std::{collections::VecDeque, sync::Arc};
6972
6973 #[test]
6974 fn test_is_narrow() {
6975 let int8_array = Int8Array::from(vec![1, 2, 3]);
6976 let array_ref: ArrayRef = Arc::new(int8_array);
6977 let block = DataBlock::from_array(array_ref);
6978
6979 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
6980
6981 let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
6982 let block = DataBlock::from_array(string_array);
6983 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
6984
6985 let string_array = StringArray::from(vec![
6986 Some("hello world".repeat(100)),
6987 Some("world".to_string()),
6988 ]);
6989 let block = DataBlock::from_array(string_array);
6990 assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
6991 }
6992
6993 #[test]
6994 fn test_primitive_decoder_empty_page_queue_returns_error() {
6995 let field = Arc::new(ArrowField::new("vector", DataType::Float32, true));
6996 let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
6997
6998 let err = decoder.drain(1).unwrap_err();
6999 assert!(
7000 matches!(&err, lance_core::Error::Internal { .. }),
7001 "expected internal error, got: {err:?}"
7002 );
7003 let message = err.to_string();
7004 for expected in [
7005 "Primitive decoder missing page decoder",
7006 "field 'vector'",
7007 "data_type=Float32",
7008 "requested_rows=1",
7009 "remaining_rows=1",
7010 "rows_drained_in_current=0",
7011 "queued_pages=0",
7012 ] {
7013 assert!(
7014 message.contains(expected),
7015 "expected error to contain {expected:?}, got: {message}"
7016 );
7017 }
7018 }
7019
7020 #[test]
7021 fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
7022 let fixed = FixedWidthDataBlock {
7023 data: LanceBuffer::from(vec![0_u8]),
7024 bits_per_value: 1,
7025 num_values: 8,
7026 block_info: BlockInfo::new(),
7027 };
7028 let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
7029
7030 let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
7031 else {
7032 panic!("expected full-zip to reject 1-bit fixed-width values");
7033 };
7034 assert!(
7035 err.to_string().contains("byte aligned"),
7036 "unexpected error: {err}"
7037 );
7038 }
7039
7040 #[test]
7041 fn test_map_range() {
7042 let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
7045 let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
7046 let max_visible_def = 0;
7047 let total_items = 8;
7048 let max_rep = 1;
7049
7050 let check = |range, expected_item_range, expected_level_range| {
7051 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7052 range,
7053 rep.as_ref(),
7054 def.as_ref(),
7055 max_rep,
7056 max_visible_def,
7057 total_items,
7058 PreambleAction::Absent,
7059 );
7060 assert_eq!(item_range, expected_item_range);
7061 assert_eq!(level_range, expected_level_range);
7062 };
7063
7064 check(0..1, 0..3, 0..3);
7065 check(1..2, 3..5, 3..5);
7066 check(2..3, 5..5, 5..6);
7067 check(3..4, 5..8, 6..9);
7068 check(0..2, 0..5, 0..5);
7069 check(1..3, 3..5, 3..6);
7070 check(2..4, 5..8, 5..9);
7071 check(0..3, 0..5, 0..6);
7072 check(1..4, 3..8, 3..9);
7073 check(0..4, 0..8, 0..9);
7074
7075 let rep = Some(vec![1, 1, 0, 1]);
7078 let def = Some(vec![1, 0, 0, 0]);
7079 let max_visible_def = 0;
7080 let total_items = 3;
7081
7082 let check = |range, expected_item_range, expected_level_range| {
7083 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7084 range,
7085 rep.as_ref(),
7086 def.as_ref(),
7087 max_rep,
7088 max_visible_def,
7089 total_items,
7090 PreambleAction::Absent,
7091 );
7092 assert_eq!(item_range, expected_item_range);
7093 assert_eq!(level_range, expected_level_range);
7094 };
7095
7096 check(0..1, 0..0, 0..1);
7097 check(1..2, 0..2, 1..3);
7098 check(2..3, 2..3, 3..4);
7099 check(0..2, 0..2, 0..3);
7100 check(1..3, 0..3, 1..4);
7101 check(0..3, 0..3, 0..4);
7102
7103 let rep = Some(vec![1, 1, 0, 1]);
7106 let def = Some(vec![0, 0, 0, 1]);
7107 let max_visible_def = 0;
7108 let total_items = 3;
7109
7110 let check = |range, expected_item_range, expected_level_range| {
7111 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7112 range,
7113 rep.as_ref(),
7114 def.as_ref(),
7115 max_rep,
7116 max_visible_def,
7117 total_items,
7118 PreambleAction::Absent,
7119 );
7120 assert_eq!(item_range, expected_item_range);
7121 assert_eq!(level_range, expected_level_range);
7122 };
7123
7124 check(0..1, 0..1, 0..1);
7125 check(1..2, 1..3, 1..3);
7126 check(2..3, 3..3, 3..4);
7127 check(0..2, 0..3, 0..3);
7128 check(1..3, 1..3, 1..4);
7129 check(0..3, 0..3, 0..4);
7130
7131 let rep = Some(vec![1, 0, 1, 0, 1, 0]);
7134 let def: Option<&[u16]> = None;
7135 let max_visible_def = 0;
7136 let total_items = 6;
7137
7138 let check = |range, expected_item_range, expected_level_range| {
7139 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7140 range,
7141 rep.as_ref(),
7142 def.as_ref(),
7143 max_rep,
7144 max_visible_def,
7145 total_items,
7146 PreambleAction::Absent,
7147 );
7148 assert_eq!(item_range, expected_item_range);
7149 assert_eq!(level_range, expected_level_range);
7150 };
7151
7152 check(0..1, 0..2, 0..2);
7153 check(1..2, 2..4, 2..4);
7154 check(2..3, 4..6, 4..6);
7155 check(0..2, 0..4, 0..4);
7156 check(1..3, 2..6, 2..6);
7157 check(0..3, 0..6, 0..6);
7158
7159 let rep: Option<&[u16]> = None;
7162 let def = Some(vec![0, 0, 1, 0]);
7163 let max_visible_def = 1;
7164 let total_items = 4;
7165
7166 let check = |range, expected_item_range, expected_level_range| {
7167 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7168 range,
7169 rep.as_ref(),
7170 def.as_ref(),
7171 max_rep,
7172 max_visible_def,
7173 total_items,
7174 PreambleAction::Absent,
7175 );
7176 assert_eq!(item_range, expected_item_range);
7177 assert_eq!(level_range, expected_level_range);
7178 };
7179
7180 check(0..1, 0..1, 0..1);
7181 check(1..2, 1..2, 1..2);
7182 check(2..3, 2..3, 2..3);
7183 check(0..2, 0..2, 0..2);
7184 check(1..3, 1..3, 1..3);
7185 check(0..3, 0..3, 0..3);
7186
7187 let rep = Some(vec![0, 1, 0, 1]);
7192 let def = Some(vec![0, 0, 0, 1]);
7193 let max_visible_def = 0;
7194 let total_items = 3;
7195
7196 let check = |range, expected_item_range, expected_level_range| {
7197 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7198 range,
7199 rep.as_ref(),
7200 def.as_ref(),
7201 max_rep,
7202 max_visible_def,
7203 total_items,
7204 PreambleAction::Take,
7205 );
7206 assert_eq!(item_range, expected_item_range);
7207 assert_eq!(level_range, expected_level_range);
7208 };
7209
7210 check(0..1, 0..3, 0..3);
7212 check(0..2, 0..3, 0..4);
7213
7214 let check = |range, expected_item_range, expected_level_range| {
7215 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7216 range,
7217 rep.as_ref(),
7218 def.as_ref(),
7219 max_rep,
7220 max_visible_def,
7221 total_items,
7222 PreambleAction::Skip,
7223 );
7224 assert_eq!(item_range, expected_item_range);
7225 assert_eq!(level_range, expected_level_range);
7226 };
7227
7228 check(0..1, 1..3, 1..3);
7229 check(1..2, 3..3, 3..4);
7230 check(0..2, 1..3, 1..4);
7231
7232 let rep = Some(vec![0, 1, 1, 0]);
7237 let def = Some(vec![0, 1, 0, 0]);
7238 let max_visible_def = 0;
7239 let total_items = 4;
7240
7241 let check = |range, expected_item_range, expected_level_range| {
7242 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7243 range,
7244 rep.as_ref(),
7245 def.as_ref(),
7246 max_rep,
7247 max_visible_def,
7248 total_items,
7249 PreambleAction::Take,
7250 );
7251 assert_eq!(item_range, expected_item_range);
7252 assert_eq!(level_range, expected_level_range);
7253 };
7254
7255 check(0..1, 0..1, 0..2);
7257 check(0..2, 0..3, 0..4);
7258
7259 let check = |range, expected_item_range, expected_level_range| {
7260 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7261 range,
7262 rep.as_ref(),
7263 def.as_ref(),
7264 max_rep,
7265 max_visible_def,
7266 total_items,
7267 PreambleAction::Skip,
7268 );
7269 assert_eq!(item_range, expected_item_range);
7270 assert_eq!(level_range, expected_level_range);
7271 };
7272
7273 check(0..1, 1..1, 1..2);
7275 check(1..2, 1..3, 2..4);
7276 check(0..2, 1..3, 1..4);
7277
7278 let rep = Some(vec![0, 1, 0, 1]);
7281 let def: Option<Vec<u16>> = None;
7282 let max_visible_def = 0;
7283 let total_items = 4;
7284
7285 let check = |range, expected_item_range, expected_level_range| {
7286 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7287 range,
7288 rep.as_ref(),
7289 def.as_ref(),
7290 max_rep,
7291 max_visible_def,
7292 total_items,
7293 PreambleAction::Take,
7294 );
7295 assert_eq!(item_range, expected_item_range);
7296 assert_eq!(level_range, expected_level_range);
7297 };
7298
7299 check(0..1, 0..3, 0..3);
7301 check(0..2, 0..4, 0..4);
7302
7303 let check = |range, expected_item_range, expected_level_range| {
7304 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7305 range,
7306 rep.as_ref(),
7307 def.as_ref(),
7308 max_rep,
7309 max_visible_def,
7310 total_items,
7311 PreambleAction::Skip,
7312 );
7313 assert_eq!(item_range, expected_item_range);
7314 assert_eq!(level_range, expected_level_range);
7315 };
7316
7317 check(0..1, 1..3, 1..3);
7318 check(1..2, 3..4, 3..4);
7319 check(0..2, 1..4, 1..4);
7320
7321 let rep = Some(vec![2, 1, 2, 0, 1, 2]);
7325 let def = Some(vec![0, 1, 2, 0, 0, 0]);
7326 let max_rep = 2;
7327 let max_visible_def = 0;
7328 let total_items = 4;
7329
7330 let check = |range, expected_item_range, expected_level_range| {
7331 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7332 range,
7333 rep.as_ref(),
7334 def.as_ref(),
7335 max_rep,
7336 max_visible_def,
7337 total_items,
7338 PreambleAction::Absent,
7339 );
7340 assert_eq!(item_range, expected_item_range);
7341 assert_eq!(level_range, expected_level_range);
7342 };
7343
7344 check(0..3, 0..4, 0..6);
7345 check(0..1, 0..1, 0..2);
7346 check(1..2, 1..3, 2..5);
7347 check(2..3, 3..4, 5..6);
7348
7349 let rep = Some(vec![0, 0, 1, 0, 1, 1]);
7351 let def = Some(vec![0, 1, 0, 0, 0, 0]);
7352 let max_rep = 1;
7353 let max_visible_def = 0;
7354 let total_items = 5;
7355
7356 let check = |range, expected_item_range, expected_level_range| {
7357 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7358 range,
7359 rep.as_ref(),
7360 def.as_ref(),
7361 max_rep,
7362 max_visible_def,
7363 total_items,
7364 PreambleAction::Take,
7365 );
7366 assert_eq!(item_range, expected_item_range);
7367 assert_eq!(level_range, expected_level_range);
7368 };
7369
7370 check(0..0, 0..1, 0..2);
7371 check(0..1, 0..3, 0..4);
7372 check(0..2, 0..4, 0..5);
7373
7374 let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
7377 let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
7378 let max_rep = 1;
7379 let max_visible_def = 0;
7380 let total_items = 5;
7381
7382 let check = |range, expected_item_range, expected_level_range| {
7383 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
7384 range,
7385 rep.as_ref(),
7386 def.as_ref(),
7387 max_rep,
7388 max_visible_def,
7389 total_items,
7390 PreambleAction::Skip,
7391 );
7392 assert_eq!(item_range, expected_item_range);
7393 assert_eq!(level_range, expected_level_range);
7394 };
7395
7396 check(2..3, 2..4, 5..7);
7397 }
7398
7399 #[test]
7400 fn test_slice_batch_data_and_rebase_offsets_u32() {
7401 let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
7402 let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
7403
7404 let (sliced_data, normalized_offsets) =
7405 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
7406 .unwrap();
7407
7408 assert_eq!(sliced_data.as_ref(), b"6789ab");
7409 let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
7410 assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
7411 }
7412
7413 #[test]
7414 fn test_slice_batch_data_and_rebase_offsets_u64() {
7415 let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
7416 let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
7417
7418 let (sliced_data, normalized_offsets) =
7419 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
7420 .unwrap();
7421
7422 assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
7423 let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
7424 assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
7425 }
7426
7427 #[test]
7428 fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
7429 let data = LanceBuffer::copy_slice(b"abcd");
7430 let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
7431
7432 let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
7433 .expect_err("offset end before start should error");
7434 assert!(err.to_string().contains("less than base"));
7435 }
7436
7437 #[test]
7438 fn test_schedule_instructions() {
7439 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
7441 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
7442 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
7443
7444 let check = |user_ranges, expected_instructions| {
7445 let instructions = ChunkInstructions::schedule_instructions(&chunk_index, user_ranges);
7446 assert_eq!(instructions, expected_instructions);
7447 };
7448
7449 let expected_take_all = vec![
7451 ChunkInstructions {
7452 chunk_idx: 0,
7453 preamble: PreambleAction::Absent,
7454 rows_to_skip: 0,
7455 rows_to_take: 6,
7456 take_trailer: true,
7457 },
7458 ChunkInstructions {
7459 chunk_idx: 1,
7460 preamble: PreambleAction::Take,
7461 rows_to_skip: 0,
7462 rows_to_take: 2,
7463 take_trailer: false,
7464 },
7465 ChunkInstructions {
7466 chunk_idx: 2,
7467 preamble: PreambleAction::Absent,
7468 rows_to_skip: 0,
7469 rows_to_take: 5,
7470 take_trailer: true,
7471 },
7472 ChunkInstructions {
7473 chunk_idx: 3,
7474 preamble: PreambleAction::Take,
7475 rows_to_skip: 0,
7476 rows_to_take: 1,
7477 take_trailer: false,
7478 },
7479 ];
7480
7481 check(&[0..14], expected_take_all.clone());
7483
7484 check(
7486 &[
7487 0..1,
7488 1..2,
7489 2..3,
7490 3..4,
7491 4..5,
7492 5..6,
7493 6..7,
7494 7..8,
7495 8..9,
7496 9..10,
7497 10..11,
7498 11..12,
7499 12..13,
7500 13..14,
7501 ],
7502 expected_take_all,
7503 );
7504
7505 check(
7509 &[0..1, 3..4],
7510 vec![
7511 ChunkInstructions {
7512 chunk_idx: 0,
7513 preamble: PreambleAction::Absent,
7514 rows_to_skip: 0,
7515 rows_to_take: 1,
7516 take_trailer: false,
7517 },
7518 ChunkInstructions {
7519 chunk_idx: 0,
7520 preamble: PreambleAction::Absent,
7521 rows_to_skip: 3,
7522 rows_to_take: 1,
7523 take_trailer: false,
7524 },
7525 ],
7526 );
7527
7528 check(
7530 &[5..6],
7531 vec![
7532 ChunkInstructions {
7533 chunk_idx: 0,
7534 preamble: PreambleAction::Absent,
7535 rows_to_skip: 5,
7536 rows_to_take: 1,
7537 take_trailer: true,
7538 },
7539 ChunkInstructions {
7540 chunk_idx: 1,
7541 preamble: PreambleAction::Take,
7542 rows_to_skip: 0,
7543 rows_to_take: 0,
7544 take_trailer: false,
7545 },
7546 ],
7547 );
7548
7549 check(
7551 &[7..10],
7552 vec![
7553 ChunkInstructions {
7554 chunk_idx: 1,
7555 preamble: PreambleAction::Skip,
7556 rows_to_skip: 1,
7557 rows_to_take: 1,
7558 take_trailer: false,
7559 },
7560 ChunkInstructions {
7561 chunk_idx: 2,
7562 preamble: PreambleAction::Absent,
7563 rows_to_skip: 0,
7564 rows_to_take: 2,
7565 take_trailer: false,
7566 },
7567 ],
7568 );
7569 }
7570
7571 #[test]
7572 fn test_drain_instructions() {
7573 fn drain_from_instructions(
7574 instructions: &mut VecDeque<ChunkInstructions>,
7575 mut rows_desired: u64,
7576 need_preamble: &mut bool,
7577 skip_in_chunk: &mut u64,
7578 ) -> Vec<ChunkDrainInstructions> {
7579 let mut drain_instructions = Vec::with_capacity(instructions.len());
7581 while rows_desired > 0 || *need_preamble {
7582 let (next_instructions, consumed_chunk) = instructions
7583 .front()
7584 .unwrap()
7585 .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
7586 if consumed_chunk {
7587 instructions.pop_front();
7588 }
7589 drain_instructions.push(next_instructions);
7590 }
7591 drain_instructions
7592 }
7593
7594 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
7596 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
7597 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
7598 let user_ranges = vec![1..7, 10..14];
7599
7600 let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
7602
7603 let mut to_drain = VecDeque::from(scheduled.clone());
7604
7605 let mut need_preamble = false;
7608 let mut skip_in_chunk = 0;
7609
7610 let next_batch =
7611 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
7612
7613 assert!(!need_preamble);
7614 assert_eq!(skip_in_chunk, 4);
7615 assert_eq!(
7616 next_batch,
7617 vec![ChunkDrainInstructions {
7618 chunk_instructions: scheduled[0].clone(),
7619 rows_to_take: 4,
7620 rows_to_skip: 0,
7621 preamble_action: PreambleAction::Absent,
7622 }]
7623 );
7624
7625 let next_batch =
7626 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
7627
7628 assert!(!need_preamble);
7629 assert_eq!(skip_in_chunk, 2);
7630
7631 assert_eq!(
7632 next_batch,
7633 vec![
7634 ChunkDrainInstructions {
7635 chunk_instructions: scheduled[0].clone(),
7636 rows_to_take: 1,
7637 rows_to_skip: 4,
7638 preamble_action: PreambleAction::Absent,
7639 },
7640 ChunkDrainInstructions {
7641 chunk_instructions: scheduled[1].clone(),
7642 rows_to_take: 1,
7643 rows_to_skip: 0,
7644 preamble_action: PreambleAction::Take,
7645 },
7646 ChunkDrainInstructions {
7647 chunk_instructions: scheduled[2].clone(),
7648 rows_to_take: 2,
7649 rows_to_skip: 0,
7650 preamble_action: PreambleAction::Absent,
7651 }
7652 ]
7653 );
7654
7655 let next_batch =
7656 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
7657
7658 assert!(!need_preamble);
7659 assert_eq!(skip_in_chunk, 0);
7660
7661 assert_eq!(
7662 next_batch,
7663 vec![
7664 ChunkDrainInstructions {
7665 chunk_instructions: scheduled[2].clone(),
7666 rows_to_take: 1,
7667 rows_to_skip: 2,
7668 preamble_action: PreambleAction::Absent,
7669 },
7670 ChunkDrainInstructions {
7671 chunk_instructions: scheduled[3].clone(),
7672 rows_to_take: 1,
7673 rows_to_skip: 0,
7674 preamble_action: PreambleAction::Take,
7675 },
7676 ]
7677 );
7678
7679 let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
7681 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
7682 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
7683 let user_ranges = vec![0..28];
7684
7685 let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
7687
7688 let mut to_drain = VecDeque::from(scheduled.clone());
7689
7690 let mut need_preamble = false;
7693 let mut skip_in_chunk = 0;
7694
7695 let next_batch =
7696 drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
7697
7698 assert_eq!(
7699 next_batch,
7700 vec![
7701 ChunkDrainInstructions {
7702 chunk_instructions: scheduled[0].clone(),
7703 rows_to_take: 6,
7704 rows_to_skip: 0,
7705 preamble_action: PreambleAction::Absent,
7706 },
7707 ChunkDrainInstructions {
7708 chunk_instructions: scheduled[1].clone(),
7709 rows_to_take: 1,
7710 rows_to_skip: 0,
7711 preamble_action: PreambleAction::Take,
7712 },
7713 ]
7714 );
7715
7716 assert!(!need_preamble);
7717 assert_eq!(skip_in_chunk, 1);
7718
7719 let next_batch =
7722 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
7723
7724 assert_eq!(
7725 next_batch,
7726 vec![
7727 ChunkDrainInstructions {
7728 chunk_instructions: scheduled[1].clone(),
7729 rows_to_take: 2,
7730 rows_to_skip: 1,
7731 preamble_action: PreambleAction::Skip,
7732 },
7733 ChunkDrainInstructions {
7734 chunk_instructions: scheduled[2].clone(),
7735 rows_to_take: 0,
7736 rows_to_skip: 0,
7737 preamble_action: PreambleAction::Take,
7738 },
7739 ]
7740 );
7741
7742 assert!(!need_preamble);
7743 assert_eq!(skip_in_chunk, 0);
7744 }
7745
7746 use super::chunk_index::{PrefixSums, RowMapping};
7747 use super::{MINIBLOCK_ALIGNMENT, Words, build_chunk_index};
7748 use bytes::Bytes;
7749 use lance_core::cache::{Context, DeepSizeOf};
7750 use rstest::rstest;
7751
7752 fn words_from(entries: &[(u32, u32)]) -> (Words, u64) {
7755 let mut raw = Vec::with_capacity(entries.len() * 2);
7756 let mut total = 0u64;
7757 for &(log, num_bytes) in entries {
7758 assert!(num_bytes > 0 && num_bytes % MINIBLOCK_ALIGNMENT as u32 == 0);
7759 let divided = num_bytes / MINIBLOCK_ALIGNMENT as u32 - 1;
7760 let word = (divided << 4) | log;
7761 assert!(word <= u16::MAX as u32, "test word {word} exceeds u16");
7762 raw.extend_from_slice(&(word as u16).to_le_bytes());
7763 total += num_bytes as u64;
7764 }
7765 (Words::from_bytes(Bytes::from(raw), false).unwrap(), total)
7766 }
7767
7768 fn rep_bytes_from(values: &[u64]) -> Vec<u8> {
7769 values.iter().flat_map(|v| v.to_le_bytes()).collect()
7770 }
7771
7772 #[rstest]
7773 #[case::uniform_partial_last(&[(3, 16), (3, 24), (0, 8)], 19, "uniform_flat", 8, 3)]
7776 #[case::single_chunk(&[(0, 24)], 5, "uniform_flat", 5, 5)]
7778 #[case::exact_multiple(&[(3, 16), (3, 16)], 16, "uniform_flat", 8, 8)]
7780 #[case::non_uniform(&[(4, 16), (2, 16), (0, 8)], 21, "flat", 16, 1)]
7782 fn test_flat_detection(
7783 #[case] entries: &[(u32, u32)],
7784 #[case] items_in_page: u64,
7785 #[case] expected_kind: &str,
7786 #[case] expected_first_items: u64,
7787 #[case] expected_last_items: u64,
7788 ) {
7789 let base = 100u64;
7790 let (words, data_buf_size) = words_from(entries);
7791 let index = build_chunk_index(&words, items_in_page, base, data_buf_size, None, 0);
7792
7793 assert_eq!(index.row_mapping_debug(), expected_kind);
7794 assert_eq!(index.num_chunks(), entries.len());
7795 assert_eq!(index.items_in_chunk(0), expected_first_items);
7796 assert_eq!(index.items_in_chunk(entries.len() - 1), expected_last_items);
7797
7798 let mut expected_start = base;
7800 for (i, &(_, num_bytes)) in entries.iter().enumerate() {
7801 let range = index.byte_range(i);
7802 assert_eq!(range.start, expected_start);
7803 assert_eq!(range.end - range.start, num_bytes as u64);
7804 expected_start = range.end;
7805 }
7806 assert_eq!(expected_start, base + data_buf_size);
7807
7808 let total_items: u64 = (0..index.num_chunks())
7810 .map(|i| index.items_in_chunk(i))
7811 .sum();
7812 assert_eq!(total_items, items_in_page);
7813 }
7814
7815 #[test]
7816 fn test_nested_detection_and_axes() {
7817 let rep = rep_bytes_from(&[5, 0, 4, 0, 3, 0]);
7819
7820 let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (0, 8)]);
7822 let index = build_chunk_index(&words, 10, 0, data_buf_size, Some(&rep), 1);
7823 assert_eq!(index.row_mapping_debug(), "nested");
7824 assert_eq!(index.num_chunks(), 3);
7825 assert_eq!(index.first_row(0), 0);
7827 assert_eq!(index.rows_in_chunk(0), 5);
7828 assert_eq!(index.first_row(1), 5);
7829 assert_eq!(index.rows_in_chunk(1), 4);
7830 assert_eq!(index.first_row(2), 9);
7831 assert_eq!(index.rows_in_chunk(2), 3);
7832 assert_eq!(index.items_in_chunk(0), 4);
7834 assert_eq!(index.items_in_chunk(1), 4);
7835 assert_eq!(index.items_in_chunk(2), 2);
7836
7837 let (words_nu, dbs_nu) = words_from(&[(3, 8), (1, 8), (0, 8)]);
7839 let index_nu = build_chunk_index(&words_nu, 15, 0, dbs_nu, Some(&rep), 1);
7840 assert_eq!(index_nu.row_mapping_debug(), "nested");
7841 assert_eq!(index_nu.items_in_chunk(0), 8);
7842 assert_eq!(index_nu.items_in_chunk(1), 2);
7843 assert_eq!(index_nu.items_in_chunk(2), 5);
7844 assert_eq!(index_nu.rows_in_chunk(0), 5);
7846 }
7847
7848 #[test]
7849 fn test_uniform_flat_matches_prefix_sum_flat() {
7850 let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (2, 8), (0, 8)]);
7852 let uniform = build_chunk_index(&words, 15, 0, data_buf_size, None, 0);
7853 assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
7854
7855 let byte_starts = PrefixSums::from_deltas([8u64, 8, 8, 8].into_iter(), 4, 32);
7857 let value_starts = PrefixSums::from_deltas([4u64, 4, 4, 3].into_iter(), 4, 15);
7858 let flat = MiniBlockChunkIndex::new(0, byte_starts, RowMapping::Flat { value_starts });
7859 assert_eq!(flat.row_mapping_debug(), "flat");
7860
7861 for i in 0..4 {
7863 assert_eq!(uniform.byte_range(i), flat.byte_range(i));
7864 assert_eq!(uniform.items_in_chunk(i), flat.items_in_chunk(i));
7865 }
7866
7867 let range_sets: Vec<Vec<std::ops::Range<u64>>> = vec![
7869 vec![0..15],
7870 vec![0..1],
7871 vec![7..8],
7872 vec![14..15],
7873 vec![3..10],
7874 vec![0..2, 5..6, 12..15],
7875 ];
7876 for ranges in &range_sets {
7877 let from_uniform = ChunkInstructions::schedule_instructions(&uniform, ranges);
7878 let from_flat = ChunkInstructions::schedule_instructions(&flat, ranges);
7879 assert_eq!(from_uniform, from_flat, "mismatch for ranges {ranges:?}");
7880 }
7881
7882 let full = ChunkInstructions::schedule_instructions(&uniform, &[0..15]);
7884 assert_eq!(full.len(), 4);
7885 for (i, inst) in full.iter().enumerate() {
7886 assert_eq!(inst.chunk_idx, i);
7887 assert_eq!(inst.preamble, PreambleAction::Absent);
7888 assert_eq!(inst.rows_to_skip, 0);
7889 assert!(!inst.take_trailer);
7890 }
7891 assert_eq!(full.iter().map(|i| i.rows_to_take).sum::<u64>(), 15);
7892 }
7893
7894 #[test]
7895 fn test_deep_size_per_variant_below_legacy() {
7896 const LEGACY_PER_CHUNK: usize = 48;
7899 let num_chunks = 3;
7900 let heap = |index: &MiniBlockChunkIndex| index.deep_size_of_children(&mut Context::new());
7901
7902 let (uniform_words, uniform_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
7903 let uniform = build_chunk_index(&uniform_words, 10, 0, uniform_dbs, None, 0);
7904 assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
7905 assert!(heap(&uniform) < LEGACY_PER_CHUNK * num_chunks);
7906
7907 let (flat_words, flat_dbs) = words_from(&[(3, 8), (1, 8), (0, 8)]);
7908 let flat = build_chunk_index(&flat_words, 11, 0, flat_dbs, None, 0);
7909 assert_eq!(flat.row_mapping_debug(), "flat");
7910 assert!(heap(&flat) < LEGACY_PER_CHUNK * num_chunks);
7911 assert!(heap(&flat) > heap(&uniform));
7913
7914 let rep = rep_bytes_from(&[4, 0, 3, 0, 3, 0]);
7915 let (nested_words, nested_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
7916 let nested = build_chunk_index(&nested_words, 10, 0, nested_dbs, Some(&rep), 1);
7917 assert_eq!(nested.row_mapping_debug(), "nested");
7918 assert!(heap(&nested) < LEGACY_PER_CHUNK * num_chunks);
7919 }
7920
7921 #[tokio::test]
7922 async fn test_fullzip_initialize_is_lazy() {
7923 use futures::{FutureExt, future::BoxFuture};
7924 use std::ops::Range;
7925 use std::sync::Mutex;
7926
7927 #[derive(Debug, Clone)]
7928 struct RecordingScheduler {
7929 data: bytes::Bytes,
7930 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
7931 }
7932
7933 impl RecordingScheduler {
7934 fn new(data: bytes::Bytes) -> Self {
7935 Self {
7936 data,
7937 requests: Arc::new(Mutex::new(Vec::new())),
7938 }
7939 }
7940
7941 fn requests(&self) -> Vec<Vec<Range<u64>>> {
7942 self.requests.lock().unwrap().clone()
7943 }
7944 }
7945
7946 impl crate::EncodingsIo for RecordingScheduler {
7947 fn submit_request(
7948 &self,
7949 ranges: Vec<Range<u64>>,
7950 _priority: u64,
7951 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
7952 self.requests.lock().unwrap().push(ranges.clone());
7953 let data = ranges
7954 .into_iter()
7955 .map(|range| self.data.slice(range.start as usize..range.end as usize))
7956 .collect::<Vec<_>>();
7957 std::future::ready(Ok(data)).boxed()
7958 }
7959 }
7960
7961 #[derive(Debug)]
7962 struct TestFixedDecompressor;
7963
7964 impl FixedPerValueDecompressor for TestFixedDecompressor {
7965 fn decompress(
7966 &self,
7967 _data: FixedWidthDataBlock,
7968 _num_rows: u64,
7969 ) -> crate::Result<DataBlock> {
7970 unimplemented!("Test decompressor")
7971 }
7972
7973 fn bits_per_value(&self) -> u64 {
7974 32
7975 }
7976 }
7977
7978 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
7979 0;
7980 16 * 1024
7981 ])));
7982 let mut scheduler = FullZipScheduler {
7983 data_buf_position: 0,
7984 data_buf_size: 4096,
7985 rep_index: Some(FullZipRepIndexDetails {
7986 buf_position: 1000,
7987 bytes_per_value: 4,
7988 }),
7989 priority: 0,
7990 rows_in_page: 100,
7991 bits_per_offset: 32,
7992 details: Arc::new(FullZipDecodeDetails {
7993 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
7994 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
7995 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
7996 max_rep: 0,
7997 max_visible_def: 0,
7998 }),
7999 cached_state: None,
8000 enable_cache: false,
8001 };
8002
8003 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
8004 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
8005
8006 assert!(
8007 cached_data
8008 .as_arc_any()
8009 .downcast_ref::<super::NoCachedPageData>()
8010 .is_some(),
8011 "FullZip initialize should not eagerly load repetition index data"
8012 );
8013 assert!(scheduler.cached_state.is_none());
8014 assert!(
8015 io.requests().is_empty(),
8016 "FullZip initialize should not issue any I/O"
8017 );
8018 }
8019
8020 #[tokio::test]
8021 async fn test_fullzip_read_source_slices_prefetched_page() {
8022 let page_start = 200_u64;
8023 let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
8024 let source = FullZipReadSource::PrefetchedPage {
8025 base_offset: page_start,
8026 data: page_data,
8027 };
8028 let ranges = vec![
8029 page_start..(page_start + 3),
8030 (page_start + 4)..(page_start + 8),
8031 ];
8032 let mut data = source.fetch(&ranges, 0).await.unwrap();
8033 assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
8034 assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
8035 }
8036
8037 #[tokio::test]
8038 async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
8039 use futures::{FutureExt, future::BoxFuture};
8040 use std::ops::Range;
8041 use std::sync::Mutex;
8042
8043 #[derive(Debug, Clone)]
8044 struct RecordingScheduler {
8045 data: bytes::Bytes,
8046 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
8047 }
8048
8049 impl RecordingScheduler {
8050 fn new(data: bytes::Bytes) -> Self {
8051 Self {
8052 data,
8053 requests: Arc::new(Mutex::new(Vec::new())),
8054 }
8055 }
8056
8057 fn requests(&self) -> Vec<Vec<Range<u64>>> {
8058 self.requests.lock().unwrap().clone()
8059 }
8060 }
8061
8062 impl crate::EncodingsIo for RecordingScheduler {
8063 fn submit_request(
8064 &self,
8065 ranges: Vec<Range<u64>>,
8066 _priority: u64,
8067 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
8068 self.requests.lock().unwrap().push(ranges.clone());
8069 let data = ranges
8070 .into_iter()
8071 .map(|range| self.data.slice(range.start as usize..range.end as usize))
8072 .collect::<Vec<_>>();
8073 std::future::ready(Ok(data)).boxed()
8074 }
8075 }
8076
8077 #[derive(Debug)]
8078 struct TestFixedDecompressor;
8079
8080 impl FixedPerValueDecompressor for TestFixedDecompressor {
8081 fn decompress(
8082 &self,
8083 _data: FixedWidthDataBlock,
8084 _num_rows: u64,
8085 ) -> crate::Result<DataBlock> {
8086 unimplemented!("Test decompressor")
8087 }
8088
8089 fn bits_per_value(&self) -> u64 {
8090 32
8091 }
8092 }
8093
8094 let rows_in_page = 100_u64;
8095 let bytes_per_value = 4_u64;
8096 let rep_start = 1000_u64;
8097 let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
8098 let mut data = vec![0_u8; 16 * 1024];
8099 data[rep_start as usize..rep_start as usize + rep_size].fill(7);
8100 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
8101
8102 let mut scheduler = FullZipScheduler {
8103 data_buf_position: 0,
8104 data_buf_size: 4096,
8105 rep_index: Some(FullZipRepIndexDetails {
8106 buf_position: rep_start,
8107 bytes_per_value,
8108 }),
8109 priority: 0,
8110 rows_in_page,
8111 bits_per_offset: 32,
8112 details: Arc::new(FullZipDecodeDetails {
8113 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
8114 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
8115 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
8116 max_rep: 0,
8117 max_visible_def: 0,
8118 }),
8119 cached_state: None,
8120 enable_cache: true,
8121 };
8122
8123 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
8124 let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
8125 assert!(
8126 cached_data
8127 .as_arc_any()
8128 .downcast_ref::<FullZipCacheableState>()
8129 .is_some()
8130 );
8131 assert!(scheduler.cached_state.is_some());
8132 assert_eq!(
8133 io.requests(),
8134 vec![vec![
8135 rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
8136 ]]
8137 );
8138 }
8139
8140 #[tokio::test]
8141 async fn test_fullzip_full_page_bypasses_rep_index_io() {
8142 use futures::{FutureExt, future::BoxFuture};
8143 use std::ops::Range;
8144 use std::sync::Mutex;
8145
8146 #[derive(Debug, Clone)]
8147 struct RecordingScheduler {
8148 data: bytes::Bytes,
8149 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
8150 }
8151
8152 impl RecordingScheduler {
8153 fn new(data: bytes::Bytes) -> Self {
8154 Self {
8155 data,
8156 requests: Arc::new(Mutex::new(Vec::new())),
8157 }
8158 }
8159
8160 fn requests(&self) -> Vec<Vec<Range<u64>>> {
8161 self.requests.lock().unwrap().clone()
8162 }
8163 }
8164
8165 impl crate::EncodingsIo for RecordingScheduler {
8166 fn submit_request(
8167 &self,
8168 ranges: Vec<Range<u64>>,
8169 _priority: u64,
8170 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
8171 self.requests.lock().unwrap().push(ranges.clone());
8172 let data = ranges
8173 .into_iter()
8174 .map(|range| self.data.slice(range.start as usize..range.end as usize))
8175 .collect::<Vec<_>>();
8176 std::future::ready(Ok(data)).boxed()
8177 }
8178 }
8179
8180 #[derive(Debug)]
8181 struct TestFixedDecompressor;
8182
8183 impl FixedPerValueDecompressor for TestFixedDecompressor {
8184 fn decompress(
8185 &self,
8186 _data: FixedWidthDataBlock,
8187 _num_rows: u64,
8188 ) -> crate::Result<DataBlock> {
8189 unimplemented!("Test decompressor")
8190 }
8191
8192 fn bits_per_value(&self) -> u64 {
8193 32
8194 }
8195 }
8196
8197 let rows_in_page = 100_u64;
8198 let data_start = 256_u64;
8199 let data_size = 500_u64;
8200 let rep_start = 4096_u64;
8201 let bytes_per_value = 4_u64;
8202
8203 let mut bytes = vec![0_u8; 16 * 1024];
8204 for i in 0..=rows_in_page {
8205 let offset = (i * 5) as u32;
8206 let pos = rep_start as usize + (i * bytes_per_value) as usize;
8207 bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
8208 }
8209 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
8210
8211 let scheduler = FullZipScheduler {
8212 data_buf_position: data_start,
8213 data_buf_size: data_size,
8214 rep_index: Some(FullZipRepIndexDetails {
8215 buf_position: rep_start,
8216 bytes_per_value,
8217 }),
8218 priority: 0,
8219 rows_in_page,
8220 bits_per_offset: 32,
8221 details: Arc::new(FullZipDecodeDetails {
8222 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
8223 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
8224 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
8225 max_rep: 0,
8226 max_visible_def: 0,
8227 }),
8228 cached_state: None,
8229 enable_cache: false,
8230 };
8231
8232 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
8233 let tasks = scheduler
8234 .schedule_ranges_rep(
8235 &[0..rows_in_page],
8236 &io_dyn,
8237 FullZipRepIndexDetails {
8238 buf_position: rep_start,
8239 bytes_per_value,
8240 },
8241 )
8242 .unwrap();
8243
8244 let requests = io.requests();
8245 assert_eq!(requests.len(), 1);
8246 assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
8247
8248 let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
8249 let requests_after_await = io.requests();
8250 assert_eq!(
8251 requests_after_await.len(),
8252 1,
8253 "full page path should not issue rep-index I/O"
8254 );
8255 }
8256
8257 #[tokio::test]
8259 async fn test_fuzz_issue_4492_empty_rep_values() {
8260 use lance_datagen::{RowCount, Seed, array, gen_batch};
8261
8262 let seed = 1823859942947654717u64;
8263 let num_rows = 2741usize;
8264
8265 let batch_gen = gen_batch().with_seed(Seed::from(seed));
8267 let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
8268 let list_generator = array::rand_list_any(base_generator, false);
8269
8270 let batch = batch_gen
8271 .anon_col(list_generator)
8272 .into_batch_rows(RowCount::from(num_rows as u64))
8273 .unwrap();
8274
8275 let list_array = batch.column(0).clone();
8276
8277 let mut metadata = HashMap::new();
8279 metadata.insert(
8280 STRUCTURAL_ENCODING_META_KEY.to_string(),
8281 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
8282 );
8283
8284 let test_cases = TestCases::default()
8285 .with_structural_encodings()
8286 .with_batch_size(100)
8287 .with_range(0..num_rows.min(500) as u64)
8288 .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
8289
8290 check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
8291 }
8292
8293 async fn test_minichunk_size_helper(
8294 string_data: Vec<Option<String>>,
8295 minichunk_size: u64,
8296 encodings: &[TestEncoding],
8297 ) {
8298 use crate::constants::MINICHUNK_SIZE_META_KEY;
8299 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
8300 use arrow_array::{ArrayRef, StringArray};
8301 use std::sync::Arc;
8302
8303 let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
8304
8305 let mut metadata = HashMap::new();
8306 metadata.insert(
8307 MINICHUNK_SIZE_META_KEY.to_string(),
8308 minichunk_size.to_string(),
8309 );
8310 metadata.insert(
8311 STRUCTURAL_ENCODING_META_KEY.to_string(),
8312 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
8313 );
8314
8315 let test_cases = TestCases::default()
8316 .with_encodings(encodings.iter().copied())
8317 .with_batch_size(1000);
8318
8319 check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
8320 }
8321
8322 #[tokio::test]
8323 async fn test_minichunk_size_roundtrip() {
8324 let mut string_data = Vec::new();
8326 for i in 0..100 {
8327 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
8328 }
8329 test_minichunk_size_helper(
8331 string_data,
8332 64,
8333 &[
8334 TestEncoding::StructuralU16,
8335 TestEncoding::StructuralU32,
8336 TestEncoding::StructuralSparse,
8337 ],
8338 )
8339 .await;
8340 }
8341
8342 #[tokio::test]
8343 async fn test_minichunk_size_128kb_v2_2() {
8344 let mut string_data = Vec::new();
8346 for i in 0..10000 {
8348 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
8349 }
8350 test_minichunk_size_helper(
8351 string_data,
8352 128 * 1024,
8353 &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
8354 )
8355 .await;
8356 }
8357
8358 #[tokio::test]
8359 async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
8360 let mut string_data = Vec::new();
8361 for i in 0..10000 {
8363 string_data.push(Some(format!("t_{}", i)));
8364 }
8365 test_minichunk_size_helper(
8366 string_data,
8367 128 * 1024,
8368 &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
8369 )
8370 .await;
8371 }
8372
8373 #[tokio::test]
8374 async fn test_large_dictionary_general_compression() {
8375 use arrow_array::{ArrayRef, StringArray};
8376 use std::collections::HashMap;
8377 use std::sync::Arc;
8378
8379 let unique_values: Vec<String> = (0..100)
8382 .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
8383 .collect();
8384
8385 let repeated_strings: Vec<_> = unique_values
8387 .iter()
8388 .cycle()
8389 .take(100_000)
8390 .map(|s| Some(s.as_str()))
8391 .collect();
8392
8393 let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
8394
8395 let test_cases = TestCases::default()
8397 .with_u32_structural_encodings()
8398 .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
8399 assert_eq!(cols.len(), 1);
8400 let col = &cols[0];
8401
8402 if let Some(PageEncoding::Structural(page_layout)) =
8404 &col.final_pages.first().map(|p| &p.description)
8405 && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
8406 &page_layout.layout
8407 && let Some(dictionary_encoding) = &mini_block.dictionary
8408 {
8409 match dictionary_encoding.compression.as_ref() {
8410 Some(Compression::General(general)) => {
8411 let compression = general.compression.as_ref().unwrap();
8413 assert!(
8414 compression.scheme()
8415 == pb21::CompressionScheme::CompressionAlgorithmLz4
8416 || compression.scheme()
8417 == pb21::CompressionScheme::CompressionAlgorithmZstd,
8418 "Expected LZ4 or Zstd compression for large dictionary"
8419 );
8420 }
8421 _ => panic!("Expected General compression for large dictionary"),
8422 }
8423 }
8424 }));
8425
8426 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
8427 }
8428
8429 fn dictionary_encoding_from_page(
8430 page: &crate::encoder::EncodedPage,
8431 ) -> &crate::format::pb21::CompressiveEncoding {
8432 let PageEncoding::Structural(layout) = &page.description else {
8433 panic!("Expected structural page encoding");
8434 };
8435 let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
8436 else {
8437 panic!("Expected mini-block layout");
8438 };
8439 layout
8440 .dictionary
8441 .as_ref()
8442 .unwrap_or_else(|| panic!("Expected dictionary encoding"))
8443 }
8444
8445 async fn encode_variable_dict_page(
8446 metadata: HashMap<String, String>,
8447 ) -> crate::encoder::EncodedPage {
8448 use arrow_array::types::Int32Type;
8449 use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
8450
8451 let values = Arc::new(StringArray::from(
8452 (0..128)
8453 .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
8454 .collect::<Vec<_>>(),
8455 )) as ArrayRef;
8456 let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
8457 let dict_array =
8458 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
8459
8460 let field = arrow_schema::Field::new(
8461 "dict_col",
8462 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8463 false,
8464 )
8465 .with_metadata(metadata);
8466
8467 encode_first_page(field, dict_array, TestEncoding::StructuralU32).await
8468 }
8469
8470 async fn encode_auto_fixed_dict_page(
8471 metadata: HashMap<String, String>,
8472 ) -> crate::encoder::EncodedPage {
8473 use arrow_array::{ArrayRef, Decimal128Array};
8474
8475 let values = (0..20_000)
8477 .map(|i| match i % 3 {
8478 0 => 10_i128,
8479 1 => 20_i128,
8480 _ => 30_i128,
8481 })
8482 .collect::<Vec<_>>();
8483 let decimal = Decimal128Array::from_iter_values(values)
8484 .with_precision_and_scale(38, 0)
8485 .unwrap();
8486 let decimal = Arc::new(decimal) as ArrayRef;
8487
8488 let mut field_metadata = metadata;
8489 field_metadata.insert(
8491 "lance-encoding:dict-size-ratio".to_string(),
8492 "0.99".to_string(),
8493 );
8494 let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
8495 .with_metadata(field_metadata);
8496
8497 encode_first_page(field, decimal, TestEncoding::StructuralU32).await
8498 }
8499
8500 #[tokio::test]
8501 async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
8502 let page = encode_variable_dict_page(HashMap::new()).await;
8503 let dictionary_encoding = dictionary_encoding_from_page(&page);
8504 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
8505 panic!("Expected General compression for dictionary values");
8506 };
8507 let compression = general.compression.as_ref().unwrap();
8508 assert_eq!(
8509 compression.scheme(),
8510 pb21::CompressionScheme::CompressionAlgorithmLz4
8511 );
8512 }
8513
8514 #[tokio::test]
8515 async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
8516 let page = encode_auto_fixed_dict_page(HashMap::new()).await;
8517 let dictionary_encoding = dictionary_encoding_from_page(&page);
8518 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
8519 panic!("Expected General compression for dictionary values");
8520 };
8521 let compression = general.compression.as_ref().unwrap();
8522 assert_eq!(
8523 compression.scheme(),
8524 pb21::CompressionScheme::CompressionAlgorithmLz4
8525 );
8526 }
8527
8528 #[tokio::test]
8529 async fn test_dict_values_general_compression_zstd() {
8530 let mut metadata = HashMap::new();
8531 metadata.insert(
8532 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
8533 "zstd".to_string(),
8534 );
8535 let page = encode_variable_dict_page(metadata).await;
8536 let dictionary_encoding = dictionary_encoding_from_page(&page);
8537 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
8538 panic!("Expected General compression for dictionary values");
8539 };
8540 let compression = general.compression.as_ref().unwrap();
8541 assert_eq!(
8542 compression.scheme(),
8543 pb21::CompressionScheme::CompressionAlgorithmZstd
8544 );
8545 }
8546
8547 #[tokio::test]
8548 async fn test_dict_values_general_compression_none() {
8549 let mut metadata = HashMap::new();
8550 metadata.insert(
8551 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
8552 "none".to_string(),
8553 );
8554 let page = encode_variable_dict_page(metadata).await;
8555 let dictionary_encoding = dictionary_encoding_from_page(&page);
8556 assert!(
8557 !matches!(
8558 dictionary_encoding.compression.as_ref(),
8559 Some(Compression::General(_))
8560 ),
8561 "Expected dictionary values to avoid General compression"
8562 );
8563 }
8564
8565 #[test]
8566 fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
8567 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
8568 &HashMap::new(),
8569 None,
8570 None,
8571 );
8572 assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
8573 assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
8574 }
8575
8576 #[test]
8577 fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
8578 let field_metadata = HashMap::from([
8579 (
8580 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
8581 "none".to_string(),
8582 ),
8583 (
8584 DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
8585 "7".to_string(),
8586 ),
8587 ]);
8588 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
8589 &field_metadata,
8590 Some("zstd".to_string()),
8591 Some("3".to_string()),
8592 );
8593 assert_eq!(
8594 metadata.get(COMPRESSION_META_KEY),
8595 Some(&"none".to_string()),
8596 );
8597 assert_eq!(
8598 metadata.get(COMPRESSION_LEVEL_META_KEY),
8599 Some(&"7".to_string()),
8600 );
8601 }
8602
8603 #[test]
8604 fn test_resolve_dict_values_compression_metadata_env_fallback() {
8605 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
8606 &HashMap::new(),
8607 Some("zstd".to_string()),
8608 Some("9".to_string()),
8609 );
8610 assert_eq!(
8611 metadata.get(COMPRESSION_META_KEY),
8612 Some(&"zstd".to_string()),
8613 );
8614 assert_eq!(
8615 metadata.get(COMPRESSION_LEVEL_META_KEY),
8616 Some(&"9".to_string()),
8617 );
8618 }
8619
8620 #[tokio::test]
8621 async fn test_dictionary_encode_int64() {
8622 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
8623 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
8624 use arrow_array::{ArrayRef, Int64Array};
8625 use std::collections::HashMap;
8626 use std::sync::Arc;
8627
8628 let values = (0..1000)
8630 .map(|i| match i % 3 {
8631 0 => 10i64,
8632 1 => 20i64,
8633 _ => 30i64,
8634 })
8635 .collect::<Vec<_>>();
8636 let array = Arc::new(Int64Array::from(values)) as ArrayRef;
8637
8638 let mut metadata = HashMap::new();
8639 metadata.insert(
8640 STRUCTURAL_ENCODING_META_KEY.to_string(),
8641 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
8642 );
8643 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
8644
8645 let test_cases = TestCases::default()
8646 .with_u32_structural_encodings()
8647 .with_batch_size(1000)
8648 .with_range(0..1000)
8649 .with_indices(vec![0, 1, 10, 999])
8650 .with_expected_encoding("dictionary");
8651
8652 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
8653 }
8654
8655 #[tokio::test]
8656 async fn test_dictionary_encode_float64() {
8657 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
8658 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
8659 use arrow_array::{ArrayRef, Float64Array};
8660 use std::collections::HashMap;
8661 use std::sync::Arc;
8662
8663 let values = (0..1000)
8665 .map(|i| match i % 3 {
8666 0 => 0.1f64,
8667 1 => 0.2f64,
8668 _ => 0.3f64,
8669 })
8670 .collect::<Vec<_>>();
8671 let array = Arc::new(Float64Array::from(values)) as ArrayRef;
8672
8673 let mut metadata = HashMap::new();
8674 metadata.insert(
8675 STRUCTURAL_ENCODING_META_KEY.to_string(),
8676 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
8677 );
8678 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
8679
8680 let test_cases = TestCases::default()
8681 .with_u32_structural_encodings()
8682 .with_batch_size(1000)
8683 .with_range(0..1000)
8684 .with_indices(vec![0, 1, 10, 999])
8685 .with_expected_encoding("dictionary");
8686
8687 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
8688 }
8689
8690 #[test]
8691 fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
8692 let rows = 10_000;
8693 let unique_values = 2_000;
8694
8695 let dictionary_encoding =
8696 ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
8697 let layout = pb21::MiniBlockLayout {
8698 rep_compression: None,
8699 def_compression: None,
8700 value_compression: Some(ProtobufUtils21::flat(64, None)),
8701 dictionary: Some(dictionary_encoding),
8702 num_dictionary_items: unique_values,
8703 layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
8704 num_buffers: 1,
8705 repetition_index_depth: 0,
8706 num_items: rows,
8707 has_large_chunk: false,
8708 };
8709
8710 let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
8711 let scheduler = super::MiniBlockScheduler::try_new(
8712 &buffer_offsets_and_sizes,
8713 0,
8714 rows,
8715 &layout,
8716 &DefaultDecompressionStrategy::default(),
8717 )
8718 .unwrap();
8719
8720 let dictionary = scheduler.dictionary.unwrap();
8721 assert_eq!(dictionary.num_dictionary_items, unique_values);
8722 assert_eq!(
8723 dictionary.dictionary_data_alignment,
8724 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
8725 );
8726 }
8727
8728 fn create_test_fixed_data_block(
8730 num_values: u64,
8731 cardinality: u64,
8732 bits_per_value: u64,
8733 ) -> DataBlock {
8734 assert!(cardinality > 0);
8735 assert!(cardinality <= num_values);
8736 let block_info = BlockInfo::default();
8737
8738 assert_eq!(bits_per_value % 8, 0);
8739 let data = match bits_per_value {
8740 32 => {
8741 let values = (0..num_values)
8742 .map(|i| (i % cardinality) as u32)
8743 .collect::<Vec<_>>();
8744 crate::buffer::LanceBuffer::reinterpret_vec(values)
8745 }
8746 64 => {
8747 let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
8748 crate::buffer::LanceBuffer::reinterpret_vec(values)
8749 }
8750 128 => {
8751 let values = (0..num_values)
8752 .map(|i| (i % cardinality) as u128)
8753 .collect::<Vec<_>>();
8754 crate::buffer::LanceBuffer::reinterpret_vec(values)
8755 }
8756 _ => unreachable!(),
8757 };
8758 DataBlock::FixedWidth(FixedWidthDataBlock {
8759 bits_per_value,
8760 data,
8761 num_values,
8762 block_info,
8763 })
8764 }
8765
8766 fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
8768 use arrow_array::StringArray;
8769
8770 assert!(cardinality <= num_values && cardinality > 0);
8771
8772 let mut values = Vec::with_capacity(num_values as usize);
8773 for i in 0..num_values {
8774 values.push(format!("value_{:016}", i % cardinality));
8775 }
8776
8777 let array = StringArray::from(values);
8778 DataBlock::from_array(Arc::new(array) as ArrayRef)
8779 }
8780
8781 fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
8782 use arrow_array::StringArray;
8783
8784 assert!(cardinality <= num_values && cardinality > 0);
8785
8786 let mut values = Vec::with_capacity(num_values as usize);
8787 for i in 0..num_values {
8788 let value_idx = i * cardinality / num_values;
8789 values.push(format!("value_{:016}", value_idx));
8790 }
8791
8792 Arc::new(StringArray::from(values)) as ArrayRef
8793 }
8794
8795 fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
8796 DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
8797 }
8798
8799 #[test]
8800 fn test_should_dictionary_encode() {
8801 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8802 use lance_core::datatypes::Field as LanceField;
8803
8804 let block = create_test_variable_width_block(1000, 10);
8806
8807 let mut metadata = HashMap::new();
8808 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
8809 let arrow_field =
8810 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
8811 let field = LanceField::try_from(&arrow_field).unwrap();
8812
8813 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
8814 &block,
8815 &field,
8816 FixedWidthDictionaryEncoding::Exclude64Bit,
8817 );
8818
8819 assert!(
8820 result.is_some(),
8821 "Should use dictionary encode based on size"
8822 );
8823 }
8824
8825 #[test]
8826 fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
8827 let sample_count: usize = 4096;
8828 let num_values: u64 = 200_000;
8829 let cardinality: u64 = 8_000;
8830 let run_length = num_values / cardinality;
8831 let stride = num_values as usize / sample_count;
8832 assert!(
8833 stride > run_length as usize,
8834 "test must construct the stride > run_length case"
8835 );
8836
8837 let block = create_sorted_variable_width_block(num_values, cardinality);
8838 let sample_unique_ratio =
8839 PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
8840
8841 assert!(
8842 sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
8843 "sorted low-cardinality data must not be classified as near-unique"
8844 );
8845 }
8846
8847 #[test]
8848 fn test_should_dictionary_encode_sorted_low_cardinality() {
8849 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8850 use lance_core::datatypes::Field as LanceField;
8851
8852 let block = create_sorted_variable_width_block(200_000, 8_000);
8853
8854 let mut metadata = HashMap::new();
8855 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
8856 let arrow_field =
8857 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
8858 let field = LanceField::try_from(&arrow_field).unwrap();
8859
8860 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
8861 &block,
8862 &field,
8863 FixedWidthDictionaryEncoding::Include64Bit,
8864 );
8865
8866 assert!(
8867 result.is_some(),
8868 "sorted low-cardinality data should reach dictionary encoding"
8869 );
8870 }
8871
8872 #[test]
8873 fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
8874 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8875 use lance_core::datatypes::Field as LanceField;
8876
8877 let num_values = 200_002;
8878 let cardinality = 100_001;
8879 let block = create_sorted_variable_width_block(num_values, cardinality);
8880
8881 let mut metadata = HashMap::new();
8882 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
8883 let arrow_field =
8884 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
8885 let field = LanceField::try_from(&arrow_field).unwrap();
8886
8887 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
8888 &block,
8889 &field,
8890 FixedWidthDictionaryEncoding::Include64Bit,
8891 );
8892
8893 assert!(
8894 result.is_none(),
8895 "sorted high-cardinality short runs should not trigger a full dictionary probe"
8896 );
8897 }
8898
8899 #[tokio::test]
8900 async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
8901 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8902
8903 let mut metadata = HashMap::new();
8904 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
8905 let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
8906 let array = create_sorted_string_array(200_000, 8_000);
8907
8908 let page = encode_first_page(field, array, TestEncoding::StructuralU32).await;
8909 let _ = dictionary_encoding_from_page(&page);
8910 }
8911
8912 #[test]
8913 fn test_should_not_dictionary_encode_unsupported_bits() {
8914 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8915 use lance_core::datatypes::Field as LanceField;
8916
8917 let block = create_test_fixed_data_block(1000, 1000, 32);
8918
8919 let mut metadata = HashMap::new();
8920 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
8921 let arrow_field =
8922 arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
8923 let field = LanceField::try_from(&arrow_field).unwrap();
8924
8925 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
8926 &block,
8927 &field,
8928 FixedWidthDictionaryEncoding::Exclude64Bit,
8929 );
8930
8931 assert!(
8932 result.is_none(),
8933 "Should not use dictionary encode for unsupported bit width"
8934 );
8935 }
8936
8937 #[test]
8938 fn test_should_not_dictionary_encode_near_unique_sample() {
8939 use crate::constants::DICT_SIZE_RATIO_META_KEY;
8940 use lance_core::datatypes::Field as LanceField;
8941
8942 let num_values = 5000;
8943 let block = create_test_variable_width_block(num_values, num_values);
8944
8945 let mut metadata = HashMap::new();
8946 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
8947 let arrow_field =
8948 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
8949 let field = LanceField::try_from(&arrow_field).unwrap();
8950
8951 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
8952 &block,
8953 &field,
8954 FixedWidthDictionaryEncoding::Exclude64Bit,
8955 );
8956
8957 assert!(
8958 result.is_none(),
8959 "Should not probe dictionary encoding for near-unique data"
8960 );
8961 }
8962
8963 #[test]
8964 fn test_v2_1_miniblock_serializes_log_num_values_15() {
8965 let miniblocks = MiniBlockCompressed {
8966 data: vec![LanceBuffer::from(vec![1_u8; 16])],
8967 chunks: vec![
8968 MiniBlockChunk {
8969 buffer_sizes: vec![8],
8970 log_num_values: 15,
8971 },
8972 MiniBlockChunk {
8973 buffer_sizes: vec![8],
8974 log_num_values: 0,
8975 },
8976 ],
8977 num_values: 32_769,
8978 };
8979
8980 let serialized = PrimitiveStructuralEncoder::serialize_miniblocks(
8981 miniblocks,
8982 None,
8983 None,
8984 MiniblockChunkSize::U16,
8985 )
8986 .unwrap();
8987
8988 let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
8989 assert_eq!(chunk_metadata.len(), 2);
8990 assert_eq!(
8991 chunk_metadata[0] & 0x0F,
8992 15,
8993 "V2.1 metadata should use all 4 bits for log_num_values"
8994 );
8995 }
8996
8997 async fn encode_first_page(
8998 field: arrow_schema::Field,
8999 array: ArrayRef,
9000 version: TestEncoding,
9001 ) -> crate::encoder::EncodedPage {
9002 use crate::repdef::RepDefBuilder;
9003 use crate::{
9004 encoder::{
9005 ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
9006 },
9007 testing::{create_test_field_encoder, test_encoding_strategy},
9008 };
9009
9010 let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
9011 let encoding_strategy = test_encoding_strategy(version);
9012 let mut column_index_seq = ColumnIndexSequence::default();
9013 let encoding_options = EncodingOptions {
9014 cache_bytes_per_column: 1,
9015 max_page_bytes: 32 * 1024 * 1024,
9016 keep_original_array: true,
9017 buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
9018 };
9019
9020 let mut encoder = create_test_field_encoder(
9021 encoding_strategy.as_ref(),
9022 &lance_field,
9023 &mut column_index_seq,
9024 &encoding_options,
9025 )
9026 .unwrap();
9027
9028 let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
9029 let repdef = RepDefBuilder::default();
9030 let num_rows = array.len() as u64;
9031 let mut pages = Vec::new();
9032 for task in encoder
9033 .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
9034 .unwrap()
9035 {
9036 pages.push(task.await.unwrap());
9037 }
9038 for task in encoder.flush(&mut external_buffers).unwrap() {
9039 pages.push(task.await.unwrap());
9040 }
9041 pages.into_iter().next().unwrap()
9042 }
9043
9044 #[tokio::test]
9045 async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
9046 use crate::format::pb21::page_layout::Layout;
9047
9048 let val = vec![0xABu8; 33];
9049 let arr: ArrayRef = Arc::new(
9050 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
9051 std::iter::repeat_n(Some(val.as_slice()), 256),
9052 33,
9053 )
9054 .unwrap(),
9055 );
9056 let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
9057 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9058
9059 let PageEncoding::Structural(layout) = &page.description else {
9060 panic!("Expected structural encoding");
9061 };
9062 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
9063 panic!("Expected constant layout in slot 2");
9064 };
9065 assert!(layout.inline_value.is_none());
9066 assert_eq!(page.data.len(), 1);
9067
9068 let test_cases = TestCases::default()
9069 .with_encoding(TestEncoding::StructuralU32)
9070 .with_page_sizes(vec![4096]);
9071 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9072 }
9073
9074 #[tokio::test]
9075 async fn test_constant_layout_out_of_line_utf8_v2_2() {
9076 use crate::format::pb21::page_layout::Layout;
9077
9078 let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
9079 std::iter::repeat_n("hello", 512),
9080 ));
9081 let field = arrow_schema::Field::new("c", DataType::Utf8, true);
9082 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9083
9084 let PageEncoding::Structural(layout) = &page.description else {
9085 panic!("Expected structural encoding");
9086 };
9087 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
9088 panic!("Expected constant layout in slot 2");
9089 };
9090 assert!(layout.inline_value.is_none());
9091 assert_eq!(page.data.len(), 1);
9092
9093 let test_cases = TestCases::default()
9094 .with_encoding(TestEncoding::StructuralU32)
9095 .with_page_sizes(vec![4096]);
9096 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9097 }
9098
9099 #[tokio::test]
9100 async fn test_constant_layout_nullable_item_v2_2() {
9101 use crate::format::pb21::page_layout::Layout;
9102
9103 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
9104 Some(7),
9105 None,
9106 Some(7),
9107 None,
9108 Some(7),
9109 ]));
9110 let field = arrow_schema::Field::new("c", DataType::Int32, true);
9111 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9112
9113 let PageEncoding::Structural(layout) = &page.description else {
9114 panic!("Expected structural encoding");
9115 };
9116 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
9117 panic!("Expected constant layout in slot 2");
9118 };
9119 assert!(layout.inline_value.is_some());
9120 assert_eq!(page.data.len(), 2);
9121
9122 let test_cases = TestCases::default()
9123 .with_encoding(TestEncoding::StructuralU32)
9124 .with_page_sizes(vec![4096]);
9125 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9126 }
9127
9128 #[tokio::test]
9129 async fn test_constant_layout_list_repdef_v2_2() {
9130 use crate::format::pb21::page_layout::Layout;
9131 use arrow_array::builder::{Int32Builder, ListBuilder};
9132
9133 let mut builder = ListBuilder::new(Int32Builder::new());
9134 builder.values().append_value(7);
9135 builder.values().append_null();
9136 builder.values().append_value(7);
9137 builder.append(true);
9138
9139 builder.append(true);
9140
9141 builder.values().append_value(7);
9142 builder.append(true);
9143
9144 builder.append_null();
9145
9146 let arr: ArrayRef = Arc::new(builder.finish());
9147 let field = arrow_schema::Field::new(
9148 "c",
9149 DataType::List(Arc::new(arrow_schema::Field::new(
9150 "item",
9151 DataType::Int32,
9152 true,
9153 ))),
9154 true,
9155 );
9156 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9157
9158 let PageEncoding::Structural(layout) = &page.description else {
9159 panic!("Expected structural encoding");
9160 };
9161 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
9162 panic!("Expected constant layout in slot 2");
9163 };
9164 assert!(layout.inline_value.is_some());
9165 assert_eq!(page.data.len(), 2);
9166
9167 let test_cases = TestCases::default()
9168 .with_encoding(TestEncoding::StructuralU32)
9169 .with_page_sizes(vec![4096]);
9170 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9171 }
9172
9173 #[tokio::test]
9174 async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
9175 use crate::format::pb21::page_layout::Layout;
9176 use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
9177
9178 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
9179 for _ in 0..64 {
9180 builder.values().append_value(1);
9181 builder.values().append_null();
9182 builder.values().append_value(3);
9183 builder.append(true);
9184 }
9185 let arr: ArrayRef = Arc::new(builder.finish());
9186 let field = arrow_schema::Field::new(
9187 "c",
9188 DataType::FixedSizeList(
9189 Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
9190 3,
9191 ),
9192 true,
9193 );
9194 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9195
9196 if let PageEncoding::Structural(layout) = &page.description {
9197 assert!(
9198 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
9199 "FixedSizeList should not use constant layout yet"
9200 );
9201 }
9202
9203 let test_cases = TestCases::default()
9204 .with_encoding(TestEncoding::StructuralU32)
9205 .with_page_sizes(vec![4096]);
9206 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9207 }
9208
9209 #[tokio::test]
9210 async fn test_constant_layout_not_written_before_v2_2() {
9211 use crate::format::pb21::page_layout::Layout;
9212
9213 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
9214 let field = arrow_schema::Field::new("c", DataType::Int32, true);
9215 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU16).await;
9216
9217 let PageEncoding::Structural(layout) = &page.description else {
9218 return;
9219 };
9220 assert!(
9221 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
9222 "Should not emit constant layout before v2.2"
9223 );
9224
9225 let test_cases = TestCases::default()
9226 .with_encoding(TestEncoding::StructuralU16)
9227 .with_page_sizes(vec![4096]);
9228 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9229 }
9230
9231 #[tokio::test]
9232 async fn test_all_null_constant_layout_still_works_v2_2() {
9233 use crate::format::pb21::page_layout::Layout;
9234
9235 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
9236 let field = arrow_schema::Field::new("c", DataType::Int32, true);
9237 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
9238
9239 let PageEncoding::Structural(layout) = &page.description else {
9240 panic!("Expected structural encoding");
9241 };
9242 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
9243 panic!("Expected layout in slot 2");
9244 };
9245 assert!(layout.inline_value.is_none());
9246 assert_eq!(page.data.len(), 0);
9247
9248 let test_cases = TestCases::default()
9249 .with_encoding(TestEncoding::StructuralU32)
9250 .with_page_sizes(vec![4096]);
9251 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
9252 }
9253
9254 #[test]
9255 fn test_encode_decode_complex_all_null_vals_roundtrip() {
9256 use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy};
9257
9258 let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
9259
9260 let compression_strategy = crate::testing::test_compression_strategy(
9261 TestEncoding::StructuralU16,
9262 crate::compression_config::CompressionParams::default(),
9263 );
9264 let decompression_strategy = DefaultDecompressionStrategy::default();
9265
9266 let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
9267 &values,
9268 compression_strategy.as_ref(),
9269 )
9270 .unwrap();
9271
9272 let decompressor = decompression_strategy
9273 .create_block_decompressor(&encoding)
9274 .unwrap();
9275 let decompressed = decompressor
9276 .decompress(compressed_buf, values.len() as u64)
9277 .unwrap();
9278 let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
9279 assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
9280 assert_eq!(decompressed_fixed_width.bits_per_value, 16);
9281 let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
9282 assert_eq!(rep_result.as_ref(), values.as_ref());
9283 }
9284
9285 #[tokio::test]
9286 async fn test_complex_all_null_compression_gated_by_version() {
9287 use crate::format::pb21::page_layout::Layout;
9288 use arrow_array::ListArray;
9289
9290 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
9291 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
9292 );
9293 let arr: ArrayRef = Arc::new(list_array);
9294 let field = arrow_schema::Field::new(
9295 "c",
9296 DataType::List(Arc::new(arrow_schema::Field::new(
9297 "item",
9298 DataType::Int32,
9299 true,
9300 ))),
9301 true,
9302 );
9303
9304 let page_v21 =
9305 encode_first_page(field.clone(), arr.clone(), TestEncoding::StructuralU16).await;
9306 let PageEncoding::Structural(layout_v21) = &page_v21.description else {
9307 panic!("Expected structural encoding");
9308 };
9309 let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
9310 panic!("Expected constant layout");
9311 };
9312 assert!(layout_v21.rep_compression.is_none());
9313 assert!(layout_v21.def_compression.is_none());
9314 assert_eq!(layout_v21.num_rep_values, 0);
9315 assert_eq!(layout_v21.num_def_values, 0);
9316
9317 let page_v22 = encode_first_page(field, arr, TestEncoding::StructuralU32).await;
9318 let PageEncoding::Structural(layout_v22) = &page_v22.description else {
9319 panic!("Expected structural encoding");
9320 };
9321 let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
9322 panic!("Expected constant layout");
9323 };
9324 assert!(layout_v22.def_compression.is_some());
9325 assert!(layout_v22.num_def_values > 0);
9326 }
9327
9328 #[tokio::test]
9329 async fn test_complex_all_null_round_trip() {
9330 use arrow_array::ListArray;
9331
9332 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
9333 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
9334 );
9335
9336 let test_cases = TestCases::default().with_u32_structural_encodings();
9337 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
9338 .await;
9339 }
9340
9341 #[tokio::test]
9342 async fn test_complex_all_null_constant_def_round_trip() {
9343 use arrow_array::ListArray;
9344
9345 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
9348 (0..5000).map(|_| None::<Vec<Option<i32>>>),
9349 );
9350
9351 let test_cases = TestCases::default().with_u32_structural_encodings();
9352 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
9353 .await;
9354 }
9355
9356 fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer {
9357 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
9358 data: LanceBuffer::reinterpret_slice(Arc::from(levels)),
9359 bits_per_value: 16,
9360 num_values: levels.len() as u64,
9361 block_info: BlockInfo::new(),
9362 });
9363 BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block)
9364 .unwrap()
9365 }
9366
9367 fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns {
9368 let frame = encoded_u16_frame(levels, run_length_width);
9369 RleDecompressor::with_run_length_width(16, run_length_width)
9370 .decode_u16_runs(frame, levels.len() as u64)
9371 .unwrap()
9372 }
9373
9374 fn physical_levels(levels: &[u16]) -> LazyLevels {
9375 LazyLevels::Runs(Arc::new(RunStorage::Physical(
9376 encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(),
9377 )))
9378 }
9379
9380 fn coalesced_levels(levels: &[u16]) -> LazyLevels {
9381 let mut values = Vec::new();
9382 let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len());
9383 for (index, &value) in levels.iter().enumerate() {
9384 if values.last() == Some(&value) {
9385 ends.set_last(index + 1).unwrap();
9386 } else {
9387 values.push(value);
9388 ends.push(index + 1).unwrap();
9389 }
9390 }
9391 LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
9392 values: values.into_boxed_slice(),
9393 ends: ends.finish(),
9394 }))
9395 }
9396
9397 #[test]
9398 fn lazy_levels_runs_match_dense() {
9399 let expanded: Vec<u16> = vec![3, 3, 1, 3, 3, 3, 0, 0];
9401 let coalesced = coalesced_levels(&expanded);
9402 let physical = physical_levels(&expanded);
9403 let dense = LazyLevels::Dense(ScalarBuffer::<u16>::from(expanded.clone()));
9404 let n = expanded.len();
9405
9406 assert_eq!(coalesced.len(), n);
9407 assert_eq!(physical.len(), n);
9408 assert_eq!(dense.len(), n);
9409
9410 let max_rep = 3u16;
9412 let row_starts: Vec<usize> = (0..n).filter(|&i| expanded[i] == max_rep).collect();
9413 for target in 0..=row_starts.len() as u64 {
9414 let want = row_starts.get(target as usize).copied().unwrap_or(n);
9415 for runs in [&coalesced, &physical] {
9416 let mut cursor = LevelCursor::default();
9417 assert_eq!(
9418 runs.seek_row_start(&mut cursor, target, max_rep).unwrap(),
9419 want,
9420 "seek_row_start({target})"
9421 );
9422 }
9423 let mut c_dense = LevelCursor::default();
9424 assert_eq!(
9425 dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(),
9426 want
9427 );
9428 }
9429
9430 for start in 0..=n {
9433 for end in start..=n {
9434 for max in [0u16, 1, 2, 3] {
9435 let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64;
9436 for runs in [&coalesced, &physical] {
9437 let mut cursor = RunPosition::default();
9438 assert_eq!(
9439 runs.count_le_cursor(&mut cursor, start..end, max).0,
9440 want,
9441 "count_le_cursor({start}..{end}, {max})"
9442 );
9443 }
9444 let mut d_cur = RunPosition::default();
9445 assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want);
9446 }
9447 for runs in [&coalesced, &physical] {
9448 let mut got = Vec::new();
9449 runs.extend_into(start..end, RunPosition::default(), &mut got);
9450 assert_eq!(
9451 got,
9452 expanded[start..end].to_vec(),
9453 "extend_into({start}..{end})"
9454 );
9455 }
9456 let mut got_dense = Vec::new();
9457 dense.extend_into(start..end, RunPosition::default(), &mut got_dense);
9458 assert_eq!(got_dense, expanded[start..end].to_vec());
9459 }
9460 }
9461 }
9462
9463 #[test]
9464 fn physical_run_hints_support_deferred_materialization() {
9465 let expanded: Vec<u16> = vec![3, 3, 1, 1, 2, 2, 0, 0];
9466 let physical = physical_levels(&expanded);
9467 let LazyLevels::Runs(runs) = &physical else {
9468 panic!("expected physical runs");
9469 };
9470 let mut first_hint = RunPosition::default();
9471 runs.seek(&mut first_hint, 2);
9472 let mut second_hint = RunPosition::default();
9473 runs.seek(&mut second_hint, 6);
9474
9475 let mut second = Vec::new();
9476 physical.extend_into(6..8, second_hint, &mut second);
9477 let mut first = Vec::new();
9478 physical.extend_into(2..4, first_hint, &mut first);
9479 assert_eq!(second, expanded[6..8]);
9480 assert_eq!(first, expanded[2..4]);
9481 }
9482
9483 mod complex_all_null_drain_parity {
9488 use std::ops::Range;
9489
9490 use arrow_buffer::ScalarBuffer;
9491 use proptest::prelude::*;
9492
9493 use super::super::{LazyLevels, LevelCursor, RunPosition};
9494 use super::{coalesced_levels, physical_levels};
9495 use crate::Result;
9496
9497 #[derive(Debug, Clone)]
9498 struct DrainInput {
9499 max_rep: u16,
9500 max_visible: u16,
9501 rep: Option<Vec<u16>>,
9502 def: Option<Vec<u16>>,
9503 ranges: Vec<Range<u64>>,
9504 }
9505
9506 fn dense_levels(levels: &[u16]) -> LazyLevels {
9507 LazyLevels::Dense(ScalarBuffer::from(levels.to_vec()))
9508 }
9509
9510 fn rle_levels(levels: &[u16]) -> LazyLevels {
9511 coalesced_levels(levels)
9512 }
9513
9514 fn seek(
9515 rep: Option<&LazyLevels>,
9516 cursor: &mut LevelCursor,
9517 row: u64,
9518 max_rep: u16,
9519 ) -> Result<usize> {
9520 match rep {
9521 Some(rep) => rep.seek_row_start(cursor, row, max_rep),
9522 None => {
9523 cursor.row = row;
9524 cursor.level = row as usize;
9525 Ok(row as usize)
9526 }
9527 }
9528 }
9529
9530 fn simulate_drain(
9533 rep: Option<&LazyLevels>,
9534 def: Option<&LazyLevels>,
9535 max_rep: u16,
9536 max_visible: u16,
9537 ranges: &[Range<u64>],
9538 ) -> Result<(Vec<Range<usize>>, u64)> {
9539 let mut rep_cursor = LevelCursor::default();
9540 let mut def_run_cursor = RunPosition::default();
9541 let mut slices: Vec<Range<usize>> = Vec::new();
9542 let mut visible = 0u64;
9543 for range in ranges {
9544 let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?;
9545 let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?;
9546 visible += match def {
9547 Some(def) => {
9548 def.count_le_cursor(
9549 &mut def_run_cursor,
9550 level_start..level_end,
9551 max_visible,
9552 )
9553 .0
9554 }
9555 None => (level_end - level_start) as u64,
9556 };
9557 match slices.last_mut() {
9558 Some(last) if last.end == level_start => last.end = level_end,
9559 _ => slices.push(level_start..level_end),
9560 }
9561 }
9562 Ok((slices, visible))
9563 }
9564
9565 fn reference_drain(
9567 rep: Option<&[u16]>,
9568 def: Option<&[u16]>,
9569 max_rep: u16,
9570 max_visible: u16,
9571 ranges: &[Range<u64>],
9572 ) -> (Vec<Range<usize>>, u64) {
9573 let total_levels = rep
9574 .map(|r| r.len())
9575 .or_else(|| def.map(|d| d.len()))
9576 .unwrap_or(0);
9577 let row_starts: Vec<usize> = match rep {
9579 Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(),
9580 None => (0..total_levels).collect(),
9581 };
9582 let level_of_row = |row: u64| {
9583 row_starts
9584 .get(row as usize)
9585 .copied()
9586 .unwrap_or(total_levels)
9587 };
9588
9589 let mut slices: Vec<Range<usize>> = Vec::new();
9590 let mut visible = 0u64;
9591 for range in ranges {
9592 let ls = level_of_row(range.start);
9593 let le = level_of_row(range.end);
9594 visible += match def {
9595 Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64,
9596 None => (le - ls) as u64,
9597 };
9598 match slices.last_mut() {
9599 Some(last) if last.end == ls => last.end = le,
9600 _ => slices.push(ls..le),
9601 }
9602 }
9603 (slices, visible)
9604 }
9605
9606 fn ranges_strategy(num_rows: u64) -> BoxedStrategy<Vec<Range<u64>>> {
9607 if num_rows == 0 {
9608 return Just(Vec::new()).boxed();
9609 }
9610 proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8)
9613 .prop_map(move |pairs| {
9614 let mut ranges = Vec::new();
9615 let mut pos = 0u64;
9616 for (gap, len) in pairs {
9617 pos = pos.saturating_add(gap);
9618 if pos >= num_rows {
9619 break;
9620 }
9621 let end = (pos + len).min(num_rows);
9622 ranges.push(pos..end);
9623 pos = end;
9624 }
9625 ranges
9626 })
9627 .boxed()
9628 }
9629
9630 fn drain_input() -> impl Strategy<Value = DrainInput> {
9631 (
9632 1u16..=3,
9633 0u16..=3,
9634 any::<bool>(),
9635 any::<bool>(),
9636 1usize..=48,
9637 )
9638 .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| {
9639 let has_def = has_def || !has_rep;
9642 let rep = if has_rep {
9643 proptest::collection::vec(0u16..=max_rep, len)
9644 .prop_map(move |mut v| {
9645 v[0] = max_rep;
9647 Some(v)
9648 })
9649 .boxed()
9650 } else {
9651 Just(None).boxed()
9652 };
9653 let def = if has_def {
9654 proptest::collection::vec(0u16..=(max_visible + 2), len)
9655 .prop_map(Some)
9656 .boxed()
9657 } else {
9658 Just(None).boxed()
9659 };
9660 (Just(max_rep), Just(max_visible), rep, def)
9661 })
9662 .prop_flat_map(|(max_rep, max_visible, rep, def)| {
9663 let num_rows = match &rep {
9664 Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64,
9665 None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0),
9666 };
9667 ranges_strategy(num_rows).prop_map(move |ranges| DrainInput {
9668 max_rep,
9669 max_visible,
9670 rep: rep.clone(),
9671 def: def.clone(),
9672 ranges,
9673 })
9674 })
9675 }
9676
9677 proptest! {
9678 #![proptest_config(ProptestConfig::with_cases(512))]
9679
9680 #[test]
9681 fn drain_matches_reference(input in drain_input()) {
9682 let DrainInput { max_rep, max_visible, rep, def, ranges } = input;
9683
9684 let reference =
9685 reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges);
9686
9687 let rep_dense = rep.as_deref().map(dense_levels);
9688 let def_dense = def.as_deref().map(dense_levels);
9689 let got_dense =
9690 simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges)
9691 .unwrap();
9692 prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference");
9693
9694 let rep_rle = rep.as_deref().map(rle_levels);
9695 let def_rle = def.as_deref().map(rle_levels);
9696 let got_rle =
9697 simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges)
9698 .unwrap();
9699 prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference");
9700
9701 let rep_physical = rep.as_deref().map(physical_levels);
9702 let def_physical = def.as_deref().map(physical_levels);
9703 let got_physical =
9704 simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges)
9705 .unwrap();
9706 prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference");
9707 }
9708 }
9709 }
9710
9711 #[test]
9712 fn lazy_levels_runs_are_compact() {
9713 let single_run = |n: usize| {
9714 let mut ends = RunEndsBuilder::with_capacity(n, 1);
9715 ends.push(n).unwrap();
9716 LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
9717 values: vec![1u16].into_boxed_slice(),
9718 ends: ends.finish(),
9719 }))
9720 };
9721 assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size());
9723 assert!(single_run(10_000_000).deep_size() < 100);
9724 assert_eq!(single_run(10_000_000).len(), 10_000_000);
9725 assert_eq!(
9727 LazyLevels::Dense(ScalarBuffer::<u16>::from(vec![1u16; 1000])).deep_size(),
9728 2000
9729 );
9730 }
9731
9732 #[test]
9733 fn lazy_levels_selects_smallest_representation() {
9734 let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8);
9735 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
9736
9737 let equal_size: Vec<u16> = std::iter::repeat_n(0, 256)
9738 .chain(std::iter::repeat_n(1, 100))
9739 .chain(std::iter::repeat_n(2, 100))
9740 .collect();
9741 let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8);
9742 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
9743
9744 let moderate_runs: Vec<u16> = (0..250)
9745 .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
9746 .collect();
9747 let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8);
9748 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
9749
9750 let split_constant = vec![7u16; 5000];
9751 let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8);
9752 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
9753
9754 let high_density: Vec<u16> = (0..70_000).map(|index| (index % 2) as u16).collect();
9755 let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32);
9756 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
9757 }
9758
9759 #[test]
9760 fn physical_runs_detach_from_large_encoded_frame() {
9761 let levels: Vec<u16> = (0..250)
9762 .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
9763 .collect();
9764 let frame = encoded_u16_frame(&levels, RunLengthWidth::U8);
9765 let frame_offset = 4096;
9766 let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000];
9767 allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref());
9768 let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len());
9769 let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8)
9770 .decode_u16_runs(frame, levels.len() as u64)
9771 .unwrap();
9772
9773 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
9774 let cached = LazyLevels::from_rle_runs(runs).unwrap();
9775 assert!(
9776 matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_)))
9777 );
9778 assert_eq!(cached.len(), levels.len());
9779 assert!(cached.deep_size() < 4096);
9780 }
9781
9782 #[test]
9783 fn complex_all_null_levels_reject_invalid_values_and_lengths() {
9784 let invalid_levels = vec![0u16, 3];
9785 for levels in [
9786 LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())),
9787 physical_levels(&invalid_levels),
9788 coalesced_levels(&invalid_levels),
9789 ] {
9790 let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err();
9791 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9792 assert!(error.to_string().contains("Invalid definition level 3"));
9793 }
9794
9795 let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2])));
9796 let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16])));
9797 let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err();
9798 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9799 assert!(
9800 error
9801 .to_string()
9802 .contains("repetition has 2, definition has 1")
9803 );
9804 }
9805
9806 #[test]
9807 fn block_levels_reject_malformed_payload_size() {
9808 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
9809 data: LanceBuffer::from(vec![0]),
9810 bits_per_value: 16,
9811 num_values: 1,
9812 block_info: BlockInfo::new(),
9813 });
9814 let error = dense_levels_from_block(block, 1, "definition").unwrap_err();
9815 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9816 assert!(
9817 error
9818 .to_string()
9819 .contains("expected 2 bytes for 1 values, got 1")
9820 );
9821 }
9822
9823 #[test]
9824 fn complex_all_null_level_codec_validates_rle_metadata() {
9825 let encoding = pb21::CompressiveEncoding {
9826 compression: Some(Compression::Rle(Box::new(pb21::Rle {
9827 values: None,
9828 run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))),
9829 }))),
9830 };
9831
9832 let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default())
9833 .unwrap_err();
9834 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9835 assert!(
9836 error
9837 .to_string()
9838 .contains("RLE compression missing values encoding")
9839 );
9840 }
9841
9842 #[tokio::test]
9844 async fn test_sparse_boolean_list_roundtrip() {
9845 use arrow_array::builder::{BooleanBuilder, ListBuilder};
9846
9847 let mut list_builder = ListBuilder::new(BooleanBuilder::new());
9848 for i in 0..1000i32 {
9849 if i % 64 == 0 {
9850 list_builder.values().append_value(i % 128 == 0);
9852 list_builder.append(true);
9853 } else {
9854 list_builder.append(false);
9855 }
9856 }
9857 let list_array = Arc::new(list_builder.finish());
9858
9859 let test_cases = TestCases::default().with_structural_encodings();
9860 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
9861 }
9862
9863 fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
9864 use crate::compression::VariablePerValueDecompressor;
9865 use crate::encodings::physical::binary::VariableDecoder;
9866 use crate::repdef::{ControlWordParser, DefinitionInterpretation};
9867 use std::sync::Arc;
9868 Arc::new(super::FullZipDecodeDetails {
9869 value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
9870 VariableDecoder::default(),
9871 )
9872 as Arc<dyn VariablePerValueDecompressor>),
9873 def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
9874 ctrl_word_parser: ControlWordParser::new(0, 0),
9875 max_rep: 0,
9876 max_visible_def: 0,
9877 })
9878 }
9879
9880 fn decode_variable_full_zip(
9881 buf: Vec<u8>,
9882 bits_per_offset: u8,
9883 ) -> lance_core::Result<super::VariableFullZipDecoder> {
9884 use std::collections::VecDeque;
9885 let mut data = VecDeque::new();
9886 data.push_back(crate::buffer::LanceBuffer::from(buf));
9887 super::VariableFullZipDecoder::new(
9888 truncated_tail_details(),
9889 data,
9890 1,
9891 bits_per_offset,
9892 bits_per_offset,
9893 )
9894 }
9895
9896 #[test]
9898 fn variable_full_zip_wellformed_length_prefix() {
9899 assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok());
9900 assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok());
9901 }
9902
9903 #[test]
9912 fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
9913 use lance_core::Error;
9914
9915 for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
9916 let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
9917 .expect_err("a truncated length prefix must not decode");
9918 assert!(
9919 matches!(err, Error::CorruptFile { .. }),
9920 "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
9921 bits,
9922 buf_len,
9923 err
9924 );
9925 let msg = err.to_string();
9926 assert!(
9927 msg.contains("truncated length prefix"),
9928 "error should say what is wrong, got: {msg}"
9929 );
9930 }
9931 }
9932}