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