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::{
28 Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, new_null_array, types::UInt64Type,
29};
30use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
31use arrow_schema::{DataType, Field as ArrowField};
32use bytes::Bytes;
33use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
34use itertools::Itertools;
35use lance_arrow::DataTypeExt;
36use lance_arrow::deepcopy::deep_copy_nulls;
37use lance_core::{
38 cache::{CacheKey, CacheKeySchema, Context, DeepSizeOf, KeyBuilder, LanceCache},
39 error::{Error, LanceOptionExt},
40 utils::bit::pad_bytes,
41};
42use log::{debug, trace};
43
44use crate::encodings::logical::primitive::miniblock::MiniBlockChunk;
45use crate::encodings::physical::rle::{RleDecompressor, RleRuns};
46use crate::utils::bytepack::ByteUnpacker;
47use crate::{
48 compression::{
49 BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
50 compress_required_block, create_rle_decompressor,
51 },
52 data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
53 decoder::{IO_REQUEST_BATCH_BYTES, RequestBatch},
54 utils::bytepack::BytepackedIntegerEncoder,
55};
56use crate::{
57 compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
58 encodings::logical::primitive::fullzip::PerValueDataBlock,
59};
60use crate::{
61 encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext},
62 statistics::{ComputeStat, GetStat, Stat},
63};
64use crate::{
65 repdef::{
66 CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
67 MiniBlockRepDefBudget, NormalizedStructuralPlan, RepDefSlicer, SerializedRepDefs,
68 build_control_word_iterator, max_visible_level,
69 },
70 utils::accumulation::AccumulationQueue,
71};
72use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
73
74use crate::constants::{
75 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
76 DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
77 DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
78 DICT_VALUES_COMPRESSION_META_KEY,
79};
80use crate::{
81 EncodingsIo,
82 buffer::LanceBuffer,
83 data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
84 decoder::{
85 ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
86 MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
87 StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
88 StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
89 },
90 encoder::{
91 EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
92 },
93 repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
94};
95
96pub mod blob;
97mod chunk_index;
98pub mod constant;
99pub mod dict;
100pub mod fullzip;
101mod layout;
102pub mod miniblock;
103pub(crate) mod sparse;
104
105use chunk_index::{ItemCounts, MiniBlockChunkIndex, PrefixSums, RowMapping, parse_nested_rep};
106
107const FILL_BYTE: u8 = 0xFE;
108const DEFAULT_DICT_DIVISOR: u64 = 2;
109const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
110const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
111const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
112
113const MAX_LEGACY_RLE_LEVELS: u64 = ((u16::MAX as u64 - std::mem::size_of::<u64>() as u64)
118 / (std::mem::size_of::<u16>() as u64 + std::mem::size_of::<u8>() as u64))
119 * u8::MAX as u64;
120
121struct PageLoadTask {
122 decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
123 num_rows: u64,
124}
125
126#[derive(Debug)]
131struct PageInitialization {
132 buffer_ranges: Vec<Option<Range<u64>>>,
133}
134
135impl PageInitialization {
136 fn new(buffer_ranges: impl IntoIterator<Item = Option<Range<u64>>>) -> Self {
137 Self {
138 buffer_ranges: buffer_ranges.into_iter().collect(),
139 }
140 }
141
142 fn required_ranges(&self) -> impl Iterator<Item = &Range<u64>> {
143 self.buffer_ranges.iter().flatten()
144 }
145
146 fn num_required_ranges(&self) -> usize {
147 self.required_ranges().count()
148 }
149
150 fn with_buffers(self, buffers: Vec<Bytes>) -> Result<PageInitializationBuffers> {
151 let expected = self.num_required_ranges();
152 if buffers.len() != expected {
153 return Err(Error::invalid_input_source(
154 format!(
155 "Page initialization received {} buffers for {expected} declared ranges",
156 buffers.len()
157 )
158 .into(),
159 ));
160 }
161
162 let mut buffers = buffers.into_iter();
163 let mut buffer_slots = Vec::with_capacity(self.buffer_ranges.len());
164 for range in self.buffer_ranges {
165 if range.is_some() {
166 let buffer = buffers.next().ok_or_else(|| {
167 Error::internal("Page initialization buffer was not assigned")
168 })?;
169 buffer_slots.push(Some(buffer));
170 } else {
171 buffer_slots.push(None);
172 }
173 }
174 debug_assert!(
175 buffers.next().is_none(),
176 "validated page initialization buffer count must be exhausted"
177 );
178 Ok(PageInitializationBuffers { buffer_slots })
179 }
180}
181
182#[derive(Debug)]
183struct PageInitializationBuffers {
184 buffer_slots: Vec<Option<Bytes>>,
185}
186
187impl PageInitializationBuffers {
188 fn try_into_array<const N: usize>(self, layout: &str) -> Result<[Option<Bytes>; N]> {
189 let actual = self.buffer_slots.len();
190 self.buffer_slots.try_into().map_err(|_| {
191 Error::internal(format!(
192 "{layout} page initialization declared {actual} buffer slots, expected {N}"
193 ))
194 })
195 }
196
197 fn into_required(self, layout: &str) -> Result<Vec<Bytes>> {
198 self.buffer_slots
199 .into_iter()
200 .enumerate()
201 .map(|(index, buffer)| {
202 buffer.ok_or_else(|| {
203 Error::internal(format!(
204 "{layout} page initialization buffer slot {index} is absent"
205 ))
206 })
207 })
208 .collect()
209 }
210}
211
212trait StructuralPageScheduler: std::fmt::Debug + Send {
215 fn needs_initialization(&self) -> bool {
217 true
218 }
219
220 fn init_layout(&self) -> Result<PageInitialization>;
226
227 fn init_from_buffers<'a>(
231 &'a mut self,
232 buffers: PageInitializationBuffers,
233 io: &Arc<dyn EncodingsIo>,
234 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
235
236 fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()>;
238
239 fn cache_view(&self) -> PageCacheView {
245 PageCacheView::Native
246 }
247
248 fn schedule_ranges(
257 &self,
258 ranges: &[Range<u64>],
259 io: &Arc<dyn EncodingsIo>,
260 ) -> Result<Vec<PageLoadTask>>;
261}
262
263#[derive(Debug)]
265struct ChunkMeta {
266 num_values: u64,
267 chunk_size_bytes: u64,
268 offset_bytes: u64,
269}
270
271#[derive(Debug, Clone)]
273struct DecodedMiniBlockChunk {
274 rep: Option<ScalarBuffer<u16>>,
275 def: Option<ScalarBuffer<u16>>,
276 values: DataBlock,
277}
278
279#[derive(Debug)]
287struct DecodeMiniBlockTask {
288 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
289 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
290 value_decompressor: Arc<dyn MiniBlockDecompressor>,
291 dictionary_data: Option<Arc<DataBlock>>,
292 def_meaning: Arc<[DefinitionInterpretation]>,
293 num_buffers: u64,
294 max_visible_level: u16,
295 instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
296 has_large_chunk: bool,
297}
298
299impl DecodeMiniBlockTask {
300 fn decoded_size_bytes(&self) -> Option<u64> {
301 if self.rep_decompressor.is_some() || self.def_decompressor.is_some() {
302 return None;
303 }
304 let num_values = self
305 .instructions
306 .iter()
307 .try_fold(0_u64, |total, (instruction, _)| {
308 total.checked_add(instruction.rows_to_take)
309 })?;
310 self.value_decompressor.decoded_size_bytes(num_values)
311 }
312
313 fn resolve_num_levels(
314 rep_decompressor: Option<&dyn BlockDecompressor>,
315 rep_levels: Option<&LanceBuffer>,
316 def_decompressor: Option<&dyn BlockDecompressor>,
317 def_levels: Option<&LanceBuffer>,
318 declared_num_levels: u16,
319 ) -> Result<u64> {
320 let rep_num_levels = match (rep_decompressor, rep_levels) {
321 (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?,
322 (None, None) => None,
323 _ => {
324 return Err(Error::invalid_input_source(
325 "miniblock repetition codec and payload presence disagree".into(),
326 ));
327 }
328 };
329 let def_num_levels = match (def_decompressor, def_levels) {
330 (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?,
331 (None, None) => None,
332 _ => {
333 return Err(Error::invalid_input_source(
334 "miniblock definition codec and payload presence disagree".into(),
335 ));
336 }
337 };
338
339 let inferred_num_levels = match (rep_num_levels, def_num_levels) {
340 (Some(rep_num_levels), Some(def_num_levels)) if rep_num_levels != def_num_levels => {
341 return Err(Error::invalid_input_source(
342 format!(
343 "miniblock structural streams disagree on the level count: repetition inferred {rep_num_levels}, definition inferred {def_num_levels}"
344 )
345 .into(),
346 ));
347 }
348 (Some(num_levels), _) | (_, Some(num_levels)) => num_levels,
349 (None, None) => return Ok(u64::from(declared_num_levels)),
350 };
351
352 let declared_num_levels = u64::from(declared_num_levels);
353 if inferred_num_levels == declared_num_levels {
354 return Ok(inferred_num_levels);
355 }
356 let u16_modulus = u64::from(u16::MAX) + 1;
357 if inferred_num_levels <= declared_num_levels
358 || inferred_num_levels % u16_modulus != declared_num_levels
359 {
360 return Err(Error::invalid_input_source(
361 format!(
362 "miniblock payload proves {inferred_num_levels} levels but the header declared {declared_num_levels}; the counts are not congruent modulo 65536"
363 )
364 .into(),
365 ));
366 }
367 if inferred_num_levels > MAX_LEGACY_RLE_LEVELS {
368 return Err(Error::invalid_input_source(
369 format!(
370 "miniblock payload proves {inferred_num_levels} levels, exceeding the legacy RLE payload bound of {MAX_LEGACY_RLE_LEVELS}"
371 )
372 .into(),
373 ));
374 }
375 Ok(inferred_num_levels)
376 }
377
378 fn decode_levels(
379 decompressor: &dyn BlockDecompressor,
380 levels: LanceBuffer,
381 expected_num_levels: u64,
382 ) -> Result<ScalarBuffer<u16>> {
383 let levels = decompressor.decompress(Some(levels), expected_num_levels)?;
384 let levels = levels.as_fixed_width().ok_or_else(|| {
385 Error::invalid_input_source(
386 "miniblock levels did not decode to fixed-width data".into(),
387 )
388 })?;
389 if levels.num_values != expected_num_levels {
390 return Err(Error::invalid_input_source(
391 format!(
392 "miniblock levels decoded {} values, expected {expected_num_levels}",
393 levels.num_values
394 )
395 .into(),
396 ));
397 }
398 if levels.bits_per_value != 16 {
399 return Err(Error::invalid_input_source(
400 format!(
401 "miniblock levels decoded {} bits per value, expected 16",
402 levels.bits_per_value
403 )
404 .into(),
405 ));
406 }
407 let expected_bytes = expected_num_levels.checked_mul(2).ok_or_else(|| {
408 Error::invalid_input_source(
409 format!("miniblock level byte count overflowed for {expected_num_levels} levels")
410 .into(),
411 )
412 })?;
413 if levels.data.len() as u64 != expected_bytes {
414 return Err(Error::invalid_input_source(
415 format!(
416 "miniblock levels decoded {} bytes, expected {expected_bytes}",
417 levels.data.len()
418 )
419 .into(),
420 ));
421 }
422 Ok(levels.data.borrow_to_typed_slice::<u16>())
423 }
424
425 fn extend_levels(
432 range: Range<u64>,
433 levels: &mut Option<LevelBuffer>,
434 level_buf: &Option<impl AsRef<[u16]>>,
435 dest_offset: usize,
436 ) {
437 if let Some(level_buf) = level_buf {
438 if levels.is_none() {
439 let mut new_levels_vec =
442 LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
443 new_levels_vec.extend(iter::repeat_n(0, dest_offset));
444 *levels = Some(new_levels_vec);
445 }
446 levels.as_mut().unwrap().extend(
447 level_buf.as_ref()[range.start as usize..range.end as usize]
448 .iter()
449 .copied(),
450 );
451 } else if let Some(levels) = levels {
452 let num_values = (range.end - range.start) as usize;
453 levels.extend(iter::repeat_n(0, num_values));
456 }
457 }
458
459 fn map_range(
496 range: Range<u64>,
497 rep: Option<&impl AsRef<[u16]>>,
498 def: Option<&impl AsRef<[u16]>>,
499 max_rep: u16,
500 max_visible_def: u16,
501 total_items: u64,
504 preamble_action: PreambleAction,
505 ) -> (Range<u64>, Range<u64>) {
506 if let Some(rep) = rep {
507 let mut rep = rep.as_ref();
508 let mut items_in_preamble = 0_u64;
511 let first_row_start = match preamble_action {
512 PreambleAction::Skip | PreambleAction::Take => {
513 let first_row_start = if let Some(def) = def.as_ref() {
514 let mut first_row_start = None;
515 for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
516 if *rep == max_rep {
517 first_row_start = Some(idx as u64);
518 break;
519 }
520 if *def <= max_visible_def {
521 items_in_preamble += 1;
522 }
523 }
524 first_row_start
525 } else {
526 let first_row_start =
527 rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
528 items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
529 first_row_start
530 };
531 if first_row_start.is_none() {
534 assert!(preamble_action == PreambleAction::Take);
535 return (0..total_items, 0..rep.len() as u64);
536 }
537 let first_row_start = first_row_start.unwrap();
538 rep = &rep[first_row_start as usize..];
539 first_row_start
540 }
541 PreambleAction::Absent => {
542 debug_assert!(rep[0] == max_rep);
543 0
544 }
545 };
546
547 if range.start == range.end {
549 debug_assert!(preamble_action == PreambleAction::Take);
550 debug_assert!(items_in_preamble <= total_items);
551 return (0..items_in_preamble, 0..first_row_start);
552 }
553 assert!(range.start < range.end);
554
555 let mut rows_seen = 0;
556 let mut new_start = 0;
557 let mut new_levels_start = 0;
558
559 if let Some(def) = def {
560 let def = &def.as_ref()[first_row_start as usize..];
561
562 let mut lead_invis_seen = 0;
564
565 if range.start > 0 {
566 if def[0] > max_visible_def {
567 lead_invis_seen += 1;
568 }
569 for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
570 if *rep == max_rep {
571 rows_seen += 1;
572 if rows_seen == range.start {
573 new_start = idx as u64 + 1 - lead_invis_seen;
574 new_levels_start = idx as u64 + 1;
575 break;
576 }
577 }
578 if *def > max_visible_def {
579 lead_invis_seen += 1;
580 }
581 }
582 }
583
584 rows_seen += 1;
585
586 let mut new_end = u64::MAX;
587 let mut new_levels_end = rep.len() as u64;
588 let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
589 let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
590 for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
591 .iter()
592 .zip(&def[(new_levels_start + 1) as usize..])
593 .enumerate()
594 {
595 if *rep == max_rep {
596 rows_seen += 1;
597 if rows_seen == range.end + 1 {
598 new_end = idx as u64 + new_start + 1 - tail_invis_seen;
599 new_levels_end = idx as u64 + new_levels_start + 1;
600 break;
601 }
602 }
603 if *def > max_visible_def {
604 tail_invis_seen += 1;
605 }
606 }
607
608 if new_end == u64::MAX {
609 new_levels_end = rep.len() as u64;
610 let total_invis_seen = lead_invis_seen + tail_invis_seen;
611 new_end = rep.len() as u64 - total_invis_seen;
612 }
613
614 assert_ne!(new_end, u64::MAX);
615
616 if preamble_action == PreambleAction::Skip {
618 new_start += items_in_preamble;
619 new_end += items_in_preamble;
620 new_levels_start += first_row_start;
621 new_levels_end += first_row_start;
622 } else if preamble_action == PreambleAction::Take {
623 debug_assert_eq!(new_start, 0);
624 debug_assert_eq!(new_levels_start, 0);
625 new_end += items_in_preamble;
626 new_levels_end += first_row_start;
627 }
628
629 debug_assert!(new_end <= total_items);
630 (new_start..new_end, new_levels_start..new_levels_end)
631 } else {
632 if range.start > 0 {
638 for (idx, rep) in rep.iter().skip(1).enumerate() {
639 if *rep == max_rep {
640 rows_seen += 1;
641 if rows_seen == range.start {
642 new_start = idx as u64 + 1;
643 break;
644 }
645 }
646 }
647 }
648 let mut new_end = rep.len() as u64;
649 if range.end < total_items {
651 for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
652 if *rep == max_rep {
653 rows_seen += 1;
654 if rows_seen == range.end {
655 new_end = idx as u64 + new_start + 1;
656 break;
657 }
658 }
659 }
660 }
661
662 if preamble_action == PreambleAction::Skip {
664 new_start += first_row_start;
665 new_end += first_row_start;
666 } else if preamble_action == PreambleAction::Take {
667 debug_assert_eq!(new_start, 0);
668 new_end += first_row_start;
669 }
670
671 debug_assert!(new_end <= total_items);
672 (new_start..new_end, new_start..new_end)
673 }
674 } else {
675 (range.clone(), range)
678 }
679 }
680
681 fn read_buffer_sizes<const LARGE: bool>(
683 buf: &[u8],
684 offset: &mut usize,
685 num_buffers: u64,
686 ) -> Vec<u32> {
687 let read_size = if LARGE { 4 } else { 2 };
688 (0..num_buffers)
689 .map(|_| {
690 let bytes = &buf[*offset..*offset + read_size];
691 let size = if LARGE {
692 u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
693 } else {
694 u16::from_le_bytes([bytes[0], bytes[1]]) as u32
696 };
697 *offset += read_size;
698 size
699 })
700 .collect()
701 }
702
703 fn decode_miniblock_chunk(
705 &self,
706 buf: &LanceBuffer,
707 items_in_chunk: u64,
708 ) -> Result<DecodedMiniBlockChunk> {
709 let mut offset = 0;
710 let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
711 offset += 2;
712
713 let rep_size = if self.rep_decompressor.is_some() {
714 let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
715 offset += 2;
716 Some(rep_size)
717 } else {
718 None
719 };
720 let def_size = if self.def_decompressor.is_some() {
721 let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
722 offset += 2;
723 Some(def_size)
724 } else {
725 None
726 };
727
728 let buffer_sizes = if self.has_large_chunk {
729 Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
730 } else {
731 Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
732 };
733
734 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
735
736 let rep = rep_size.map(|rep_size| {
737 let rep = buf.slice_with_length(offset, rep_size as usize);
738 offset += rep_size as usize;
739 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
740 rep
741 });
742
743 let def = def_size.map(|def_size| {
744 let def = buf.slice_with_length(offset, def_size as usize);
745 offset += def_size as usize;
746 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
747 def
748 });
749
750 let num_levels = Self::resolve_num_levels(
751 self.rep_decompressor.as_deref(),
752 rep.as_ref(),
753 self.def_decompressor.as_deref(),
754 def.as_ref(),
755 num_levels,
756 )?;
757
758 let buffers = buffer_sizes
759 .into_iter()
760 .map(|buf_size| {
761 let buf = buf.slice_with_length(offset, buf_size as usize);
762 offset += buf_size as usize;
763 offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
764 buf
765 })
766 .collect::<Vec<_>>();
767
768 let values = self
769 .value_decompressor
770 .decompress(buffers, items_in_chunk)?;
771
772 let rep = rep
773 .map(|rep| {
774 Self::decode_levels(
775 self.rep_decompressor.as_ref().unwrap().as_ref(),
776 rep,
777 num_levels,
778 )
779 })
780 .transpose()?;
781 let def = def
782 .map(|def| {
783 Self::decode_levels(
784 self.def_decompressor.as_ref().unwrap().as_ref(),
785 def,
786 num_levels,
787 )
788 })
789 .transpose()?;
790
791 if let (Some(rep), Some(def)) = (&rep, &def)
792 && rep.len() != def.len()
793 {
794 return Err(Error::invalid_input_source(
795 format!(
796 "miniblock structural streams decoded different level counts: repetition {}, definition {}",
797 rep.len(),
798 def.len()
799 )
800 .into(),
801 ));
802 }
803
804 Ok(DecodedMiniBlockChunk { rep, def, values })
805 }
806}
807
808impl DecodePageTask for DecodeMiniBlockTask {
809 fn decode(self: Box<Self>) -> Result<DecodedPage> {
810 let mut repbuf: Option<LevelBuffer> = None;
812 let mut defbuf: Option<LevelBuffer> = None;
813
814 let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
815
816 let estimated_size_bytes = self.decoded_size_bytes().unwrap_or_else(|| {
817 self.instructions
819 .iter()
820 .map(|(_, chunk)| chunk.data.len() as u64)
821 .sum::<u64>()
822 * 2
823 });
824 let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes);
825
826 let mut level_offset = 0;
828
829 let needs_caching: Vec<bool> = self
831 .instructions
832 .windows(2)
833 .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
834 .chain(std::iter::once(false)) .collect();
836
837 let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
839
840 for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
842 let should_cache_this_chunk = needs_caching[idx];
843
844 let decoded_chunk = match &chunk_cache {
845 Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
846 cached_chunk.clone()
848 }
849 _ => {
850 let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
852
853 if should_cache_this_chunk {
855 chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
856 }
857 decoded
858 }
859 };
860
861 let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
862
863 let row_range_start =
865 instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
866 let row_range_end = row_range_start + instructions.rows_to_take;
867
868 let (item_range, level_range) = Self::map_range(
870 row_range_start..row_range_end,
871 rep.as_ref(),
872 def.as_ref(),
873 max_rep,
874 self.max_visible_level,
875 chunk.items_in_chunk,
876 instructions.preamble_action,
877 );
878 if item_range.end - item_range.start > chunk.items_in_chunk {
879 return Err(lance_core::Error::internal(format!(
880 "Item range {:?} is greater than chunk items in chunk {:?}",
881 item_range, chunk.items_in_chunk
882 )));
883 }
884
885 Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
887 Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
888 level_offset += (level_range.end - level_range.start) as usize;
889 data_builder.append(&values, item_range)?;
890 }
891
892 let mut data = data_builder.finish();
893
894 let unraveler =
895 RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
896
897 if let Some(dictionary) = &self.dictionary_data {
898 let DataBlock::FixedWidth(indices) = data else {
900 return Err(lance_core::Error::internal(format!(
901 "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
902 data
903 )));
904 };
905 data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
906 indices,
907 dictionary.as_ref().clone(),
908 ));
909 }
910
911 Ok(DecodedPage {
912 data,
913 repdef: unraveler,
914 })
915 }
916}
917
918#[derive(Debug)]
921struct LoadedChunk {
922 data: LanceBuffer,
923 items_in_chunk: u64,
924 byte_range: Range<u64>,
925 chunk_idx: usize,
926}
927
928impl Clone for LoadedChunk {
929 fn clone(&self) -> Self {
930 Self {
931 data: self.data.clone(),
933 items_in_chunk: self.items_in_chunk,
934 byte_range: self.byte_range.clone(),
935 chunk_idx: self.chunk_idx,
936 }
937 }
938}
939
940#[derive(Debug)]
943struct MiniBlockDecoder {
944 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
945 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
946 value_decompressor: Arc<dyn MiniBlockDecompressor>,
947 def_meaning: Arc<[DefinitionInterpretation]>,
948 loaded_chunks: VecDeque<LoadedChunk>,
949 instructions: VecDeque<ChunkInstructions>,
950 offset_in_current_chunk: u64,
951 num_rows: u64,
952 num_buffers: u64,
953 dictionary: Option<Arc<DataBlock>>,
954 has_large_chunk: bool,
955}
956
957impl StructuralPageDecoder for MiniBlockDecoder {
960 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
961 let mut items_desired = num_rows;
962 let mut need_preamble = false;
963 let mut skip_in_chunk = self.offset_in_current_chunk;
964 let mut drain_instructions = Vec::new();
965 while items_desired > 0 || need_preamble {
966 let (instructions, consumed) = self
967 .instructions
968 .front()
969 .unwrap()
970 .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
971
972 while self.loaded_chunks.front().unwrap().chunk_idx
973 != instructions.chunk_instructions.chunk_idx
974 {
975 self.loaded_chunks.pop_front();
976 }
977 drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
978 if consumed {
979 self.instructions.pop_front();
980 }
981 }
982 self.offset_in_current_chunk = skip_in_chunk;
985
986 let max_visible_level = max_visible_level(&self.def_meaning);
987
988 Ok(Box::new(DecodeMiniBlockTask {
989 instructions: drain_instructions,
990 def_decompressor: self.def_decompressor.clone(),
991 rep_decompressor: self.rep_decompressor.clone(),
992 value_decompressor: self.value_decompressor.clone(),
993 dictionary_data: self.dictionary.clone(),
994 def_meaning: self.def_meaning.clone(),
995 num_buffers: self.num_buffers,
996 max_visible_level,
997 has_large_chunk: self.has_large_chunk,
998 }))
999 }
1000
1001 fn num_rows(&self) -> u64 {
1002 self.num_rows
1003 }
1004}
1005
1006#[derive(Debug, Clone)]
1010pub(crate) enum LevelCodec {
1011 Uncompressed,
1013 Rle(Arc<RleDecompressor>),
1015 Block(Arc<dyn BlockDecompressor>),
1018}
1019
1020impl LevelCodec {
1021 fn try_new(
1022 encoding: Option<&CompressiveEncoding>,
1023 decompression_strategy: &dyn DecompressionStrategy,
1024 ) -> Result<Self> {
1025 match encoding {
1026 None => Ok(Self::Uncompressed),
1027 Some(encoding) => match encoding.compression.as_ref() {
1028 Some(Compression::Rle(rle)) => Ok(Self::Rle(Arc::new(create_rle_decompressor(
1029 rle,
1030 decompression_strategy,
1031 )?))),
1032 _ => Ok(Self::Block(Arc::from(
1033 decompression_strategy.create_block_decompressor(encoding)?,
1034 ))),
1035 },
1036 }
1037 }
1038}
1039
1040#[derive(Debug)]
1041enum RunEnds {
1042 U16(Box<[u16]>),
1043 U32(Box<[u32]>),
1044 U64(Box<[u64]>),
1045}
1046
1047impl RunEnds {
1048 fn width_for(num_values: usize) -> usize {
1049 if u16::try_from(num_values).is_ok() {
1050 std::mem::size_of::<u16>()
1051 } else if u32::try_from(num_values).is_ok() {
1052 std::mem::size_of::<u32>()
1053 } else {
1054 std::mem::size_of::<u64>()
1055 }
1056 }
1057
1058 fn len(&self) -> usize {
1059 match self {
1060 Self::U16(ends) => ends.len(),
1061 Self::U32(ends) => ends.len(),
1062 Self::U64(ends) => ends.len(),
1063 }
1064 }
1065
1066 fn get(&self, run: usize) -> usize {
1067 match self {
1068 Self::U16(ends) => ends[run] as usize,
1069 Self::U32(ends) => ends[run] as usize,
1070 Self::U64(ends) => ends[run] as usize,
1071 }
1072 }
1073
1074 fn partition_point(&self, logical_index: usize) -> usize {
1075 match self {
1076 Self::U16(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1077 Self::U32(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1078 Self::U64(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1079 }
1080 }
1081
1082 fn deep_size(&self) -> usize {
1083 match self {
1084 Self::U16(ends) => std::mem::size_of_val(ends.as_ref()),
1085 Self::U32(ends) => std::mem::size_of_val(ends.as_ref()),
1086 Self::U64(ends) => std::mem::size_of_val(ends.as_ref()),
1087 }
1088 }
1089}
1090
1091enum RunEndsBuilder {
1092 U16(Vec<u16>),
1093 U32(Vec<u32>),
1094 U64(Vec<u64>),
1095}
1096
1097impl RunEndsBuilder {
1098 fn with_capacity(num_values: usize, capacity: usize) -> Self {
1099 if u16::try_from(num_values).is_ok() {
1100 Self::U16(Vec::with_capacity(capacity))
1101 } else if u32::try_from(num_values).is_ok() {
1102 Self::U32(Vec::with_capacity(capacity))
1103 } else {
1104 Self::U64(Vec::with_capacity(capacity))
1105 }
1106 }
1107
1108 fn push(&mut self, end: usize) -> Result<()> {
1109 match self {
1110 Self::U16(ends) => ends.push(
1111 u16::try_from(end)
1112 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?,
1113 ),
1114 Self::U32(ends) => ends.push(
1115 u32::try_from(end)
1116 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?,
1117 ),
1118 Self::U64(ends) => ends.push(end as u64),
1119 }
1120 Ok(())
1121 }
1122
1123 fn set_last(&mut self, end: usize) -> Result<()> {
1124 match self {
1125 Self::U16(ends) => {
1126 let last = ends.last_mut().ok_or_else(|| {
1127 Error::internal("Cannot extend an empty coalesced run buffer")
1128 })?;
1129 *last = u16::try_from(end)
1130 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?;
1131 }
1132 Self::U32(ends) => {
1133 let last = ends.last_mut().ok_or_else(|| {
1134 Error::internal("Cannot extend an empty coalesced run buffer")
1135 })?;
1136 *last = u32::try_from(end)
1137 .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?;
1138 }
1139 Self::U64(ends) => {
1140 let last = ends.last_mut().ok_or_else(|| {
1141 Error::internal("Cannot extend an empty coalesced run buffer")
1142 })?;
1143 *last = end as u64;
1144 }
1145 }
1146 Ok(())
1147 }
1148
1149 fn finish(self) -> RunEnds {
1150 match self {
1151 Self::U16(ends) => RunEnds::U16(ends.into_boxed_slice()),
1152 Self::U32(ends) => RunEnds::U32(ends.into_boxed_slice()),
1153 Self::U64(ends) => RunEnds::U64(ends.into_boxed_slice()),
1154 }
1155 }
1156}
1157
1158#[derive(Debug)]
1159enum RunStorage {
1160 Physical(RleRuns),
1161 Coalesced { values: Box<[u16]>, ends: RunEnds },
1162}
1163
1164impl RunStorage {
1165 fn len(&self) -> usize {
1166 match self {
1167 Self::Physical(runs) => runs.num_values(),
1168 Self::Coalesced { ends, .. } => ends.get(ends.len() - 1),
1169 }
1170 }
1171
1172 fn num_runs(&self) -> usize {
1173 match self {
1174 Self::Physical(runs) => runs.num_runs(),
1175 Self::Coalesced { values, .. } => values.len(),
1176 }
1177 }
1178
1179 fn value(&self, run: usize) -> u16 {
1180 match self {
1181 Self::Physical(runs) => runs.value(run),
1182 Self::Coalesced { values, .. } => values[run],
1183 }
1184 }
1185
1186 fn first_value_above(&self, max: u16) -> Option<(usize, u16)> {
1187 (0..self.num_runs()).find_map(|run| {
1188 let value = self.value(run);
1189 (value > max).then_some((run, value))
1190 })
1191 }
1192
1193 fn seek(&self, position: &mut RunPosition, logical_index: usize) {
1194 if logical_index >= self.len() {
1195 *position = RunPosition {
1196 run: self.num_runs(),
1197 start: self.len(),
1198 end: self.len(),
1199 };
1200 return;
1201 }
1202
1203 match self {
1204 Self::Physical(runs) => {
1205 if position.run >= runs.num_runs()
1206 || position.end == 0
1207 || logical_index < position.start
1208 {
1209 *position = RunPosition {
1210 run: 0,
1211 start: 0,
1212 end: runs.length(0),
1213 };
1214 }
1215 while position.end <= logical_index {
1216 self.advance(position);
1217 }
1218 }
1219 Self::Coalesced { ends, .. } => {
1220 if logical_index < position.start || logical_index >= position.end {
1221 let run = ends.partition_point(logical_index);
1222 *position = RunPosition {
1223 run,
1224 start: if run == 0 { 0 } else { ends.get(run - 1) },
1225 end: ends.get(run),
1226 };
1227 }
1228 }
1229 }
1230 }
1231
1232 fn advance(&self, position: &mut RunPosition) {
1233 let next_run = position.run + 1;
1234 if next_run >= self.num_runs() {
1235 *position = RunPosition {
1236 run: self.num_runs(),
1237 start: self.len(),
1238 end: self.len(),
1239 };
1240 return;
1241 }
1242
1243 let start = position.end;
1244 position.run = next_run;
1245 position.start = start;
1246 position.end = match self {
1247 Self::Physical(runs) => start + runs.length(next_run),
1248 Self::Coalesced { ends, .. } => ends.get(next_run),
1249 };
1250 }
1251
1252 fn deep_size(&self) -> usize {
1253 match self {
1254 Self::Physical(runs) => runs.deep_size(),
1255 Self::Coalesced { values, ends } => {
1256 std::mem::size_of_val(values.as_ref()) + ends.deep_size()
1257 }
1258 }
1259 }
1260}
1261
1262#[derive(Debug, Clone)]
1268enum LazyLevels {
1269 Dense(ScalarBuffer<u16>),
1270 Runs(Arc<RunStorage>),
1271}
1272
1273#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1274enum LevelPlan {
1275 Physical,
1276 Coalesced,
1277 Dense,
1278}
1279
1280#[derive(Debug, Default, Clone, Copy)]
1281struct RunPosition {
1282 run: usize,
1283 start: usize,
1284 end: usize,
1285}
1286
1287#[derive(Debug, Default, Clone, Copy)]
1293struct LevelCursor {
1294 level: usize,
1296 row: u64,
1298 run: RunPosition,
1300}
1301
1302impl LazyLevels {
1303 fn from_rle_runs(runs: RleRuns) -> Result<Self> {
1304 let plan = Self::select_plan(&runs);
1305 match plan {
1306 LevelPlan::Physical => Ok(Self::Runs(Arc::new(RunStorage::Physical(
1307 runs.into_owned(),
1308 )))),
1309 LevelPlan::Coalesced => Self::build_coalesced(runs),
1310 LevelPlan::Dense => Self::build_dense(runs),
1311 }
1312 }
1313
1314 fn select_plan(runs: &RleRuns) -> LevelPlan {
1317 if runs.num_values() == 0 {
1318 return LevelPlan::Dense;
1319 }
1320
1321 let run_storage_size = std::mem::size_of::<RunStorage>() as u128;
1322 let physical_size = run_storage_size + runs.owned_size() as u128;
1323 let coalesced_size = run_storage_size
1324 + (runs.coalesced_runs() as u128)
1325 * (std::mem::size_of::<u16>() + RunEnds::width_for(runs.num_values())) as u128;
1326 let dense_size = (runs.num_values() as u128) * std::mem::size_of::<u16>() as u128;
1327 [
1328 (physical_size, runs.num_runs(), 0usize, LevelPlan::Physical),
1329 (
1330 coalesced_size,
1331 runs.coalesced_runs(),
1332 1usize,
1333 LevelPlan::Coalesced,
1334 ),
1335 (dense_size, runs.num_values(), 2usize, LevelPlan::Dense),
1336 ]
1337 .into_iter()
1338 .min_by_key(|(size, traversal, priority, _)| (*size, *traversal, *priority))
1339 .map(|(_, _, _, plan)| plan)
1340 .unwrap_or(LevelPlan::Dense)
1341 }
1342
1343 fn build_coalesced(runs: RleRuns) -> Result<Self> {
1344 let mut values = Vec::with_capacity(runs.coalesced_runs());
1345 let mut ends = RunEndsBuilder::with_capacity(runs.num_values(), runs.coalesced_runs());
1346 let mut logical_end = 0usize;
1347 for (value, length) in runs.iter() {
1348 logical_end = logical_end
1349 .checked_add(length)
1350 .ok_or_else(|| Error::internal("Validated RLE run length sum overflowed usize"))?;
1351 if values.last().copied() == Some(value) {
1352 ends.set_last(logical_end)?;
1353 } else {
1354 values.push(value);
1355 ends.push(logical_end)?;
1356 }
1357 }
1358 Ok(Self::Runs(Arc::new(RunStorage::Coalesced {
1359 values: values.into_boxed_slice(),
1360 ends: ends.finish(),
1361 })))
1362 }
1363
1364 fn build_dense(runs: RleRuns) -> Result<Self> {
1365 let mut values = Vec::new();
1366 values.try_reserve_exact(runs.num_values()).map_err(|_| {
1367 Error::internal(format!(
1368 "Cannot allocate {} dense repetition/definition levels",
1369 runs.num_values()
1370 ))
1371 })?;
1372 for (value, length) in runs.iter() {
1373 values.resize(values.len() + length, value);
1374 }
1375 Ok(Self::Dense(ScalarBuffer::from(values)))
1376 }
1377
1378 fn len(&self) -> usize {
1379 match self {
1380 Self::Dense(buf) => buf.len(),
1381 Self::Runs(runs) => runs.len(),
1382 }
1383 }
1384
1385 fn validate_max(&self, level_type: &str, max: u16) -> Result<()> {
1386 let invalid = match self {
1387 Self::Dense(levels) => levels
1388 .iter()
1389 .enumerate()
1390 .find_map(|(index, &value)| (value > max).then_some(("index", index, value))),
1391 Self::Runs(runs) => runs
1392 .first_value_above(max)
1393 .map(|(run, value)| ("run", run, value)),
1394 };
1395 if let Some((position_type, position, value)) = invalid {
1396 return Err(Error::invalid_input_source(
1397 format!(
1398 "Invalid {level_type} level {value} at {position_type} {position}: maximum is {max}"
1399 )
1400 .into(),
1401 ));
1402 }
1403 Ok(())
1404 }
1405
1406 fn seek_row_start(
1413 &self,
1414 cursor: &mut LevelCursor,
1415 target_row: u64,
1416 max_rep: u16,
1417 ) -> Result<usize> {
1418 let mut need = target_row.checked_sub(cursor.row).ok_or_else(|| {
1419 Error::internal(format!(
1420 "Complex all-null row ranges are not sorted: target row {target_row} follows {}",
1421 cursor.row
1422 ))
1423 })?;
1424 if need == 0 {
1425 return Ok(cursor.level);
1426 }
1427 match self {
1428 Self::Dense(buf) => {
1429 let mut level = cursor.level;
1430 while need > 0 {
1431 if level >= buf.len() {
1432 return Err(Error::internal(
1433 "Invalid complex all-null layout: repetition buffer too short",
1434 ));
1435 }
1436 if buf[level] != max_rep {
1437 return Err(Error::internal(
1438 "Invalid complex all-null layout: row did not start at max repetition level",
1439 ));
1440 }
1441 level += 1;
1442 while level < buf.len() && buf[level] != max_rep {
1443 level += 1;
1444 }
1445 need -= 1;
1446 }
1447 cursor.level = level;
1448 cursor.row = target_row;
1449 Ok(level)
1450 }
1451 Self::Runs(runs) => {
1452 let mut level = cursor.level;
1453 let mut run = cursor.run;
1454 runs.seek(&mut run, level);
1455 while need > 0 {
1456 if run.run >= runs.num_runs() {
1457 return Err(Error::internal(
1458 "Invalid complex all-null layout: repetition buffer too short",
1459 ));
1460 }
1461 if runs.value(run.run) != max_rep {
1462 return Err(Error::internal(
1463 "Invalid complex all-null layout: row did not start at max repetition level",
1464 ));
1465 }
1466 let avail = (run.end - level) as u64;
1467 if need < avail {
1468 level += need as usize;
1470 need = 0;
1471 } else {
1472 need -= avail;
1475 runs.advance(&mut run);
1476 while run.run < runs.num_runs() && runs.value(run.run) != max_rep {
1477 runs.advance(&mut run);
1478 }
1479 level = if run.run < runs.num_runs() {
1480 run.start
1481 } else {
1482 self.len()
1483 };
1484 }
1485 }
1486 cursor.level = level;
1487 cursor.row = target_row;
1488 cursor.run = run;
1489 Ok(level)
1490 }
1491 }
1492 }
1493
1494 fn count_le_cursor(
1500 &self,
1501 run_cursor: &mut RunPosition,
1502 range: Range<usize>,
1503 max: u16,
1504 ) -> (u64, RunPosition) {
1505 if range.is_empty() {
1506 return (0, *run_cursor);
1507 }
1508 match self {
1509 Self::Dense(buf) => (
1510 buf[range].iter().filter(|&&d| d <= max).count() as u64,
1511 RunPosition::default(),
1512 ),
1513 Self::Runs(runs) => {
1514 runs.seek(run_cursor, range.start);
1516 let start = *run_cursor;
1517 let mut count = 0u64;
1518 let mut current = *run_cursor;
1519 while current.run < runs.num_runs() && current.start < range.end {
1520 if runs.value(current.run) <= max {
1521 let lo = current.start.max(range.start);
1522 let hi = current.end.min(range.end);
1523 count += (hi - lo) as u64;
1524 }
1525 if current.end >= range.end {
1526 break;
1527 }
1528 runs.advance(&mut current);
1529 }
1530 *run_cursor = current;
1533 (count, start)
1534 }
1535 }
1536 }
1537
1538 fn extend_into(&self, range: Range<usize>, run: RunPosition, out: &mut Vec<u16>) {
1539 if range.is_empty() {
1540 return;
1541 }
1542 match self {
1543 Self::Dense(buf) => out.extend_from_slice(&buf[range]),
1544 Self::Runs(runs) => {
1545 let mut current = run;
1546 runs.seek(&mut current, range.start);
1547 while current.run < runs.num_runs() && current.start < range.end {
1548 let lo = current.start.max(range.start);
1549 let hi = current.end.min(range.end);
1550 if hi > lo {
1551 out.resize(out.len() + (hi - lo), runs.value(current.run));
1552 }
1553 runs.advance(&mut current);
1554 }
1555 }
1556 }
1557 }
1558
1559 #[cfg(test)]
1560 fn deep_size(&self) -> usize {
1561 self.deep_size_of_children(&mut Context::new())
1562 }
1563}
1564
1565impl DeepSizeOf for LazyLevels {
1566 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1567 match self {
1568 Self::Dense(buf) => buf.deep_size_of_children(ctx),
1569 Self::Runs(runs) => {
1570 let pointer = Arc::as_ptr(runs) as *const () as usize;
1571 if ctx.mark_seen(pointer) {
1572 std::mem::size_of_val(runs.as_ref()) + runs.deep_size()
1573 } else {
1574 0
1575 }
1576 }
1577 }
1578 }
1579}
1580
1581fn validate_complex_all_null_levels(
1582 rep: &Option<LazyLevels>,
1583 def: &Option<LazyLevels>,
1584 max_rep: u16,
1585 max_def: u16,
1586) -> Result<()> {
1587 if let Some(rep) = rep {
1588 rep.validate_max("repetition", max_rep)?;
1589 }
1590 if let Some(def) = def {
1591 def.validate_max("definition", max_def)?;
1592 }
1593 if let (Some(rep), Some(def)) = (rep, def)
1594 && rep.len() != def.len()
1595 {
1596 return Err(Error::invalid_input_source(
1597 format!(
1598 "Mismatched complex all-null level counts: repetition has {}, definition has {}",
1599 rep.len(),
1600 def.len()
1601 )
1602 .into(),
1603 ));
1604 }
1605 Ok(())
1606}
1607
1608fn expected_level_bytes(num_values: u64, level_type: &str) -> Result<usize> {
1609 usize::try_from(num_values)
1610 .ok()
1611 .and_then(|num_values| num_values.checked_mul(std::mem::size_of::<u16>()))
1612 .ok_or_else(|| {
1613 Error::invalid_input_source(
1614 format!("{level_type} level count {num_values} does not fit in memory").into(),
1615 )
1616 })
1617}
1618
1619fn dense_levels_from_block(
1620 decompressed: DataBlock,
1621 num_values: u64,
1622 level_type: &str,
1623) -> Result<LazyLevels> {
1624 let DataBlock::FixedWidth(block) = decompressed else {
1625 return Err(Error::invalid_input_source(
1626 format!("Expected fixed-width data block for {level_type} levels").into(),
1627 ));
1628 };
1629 if block.num_values != num_values {
1630 return Err(Error::invalid_input_source(
1631 format!(
1632 "Unexpected {level_type} level count after decompression: expected {num_values}, got {}",
1633 block.num_values
1634 )
1635 .into(),
1636 ));
1637 }
1638 if block.bits_per_value != 16 {
1639 return Err(Error::invalid_input_source(
1640 format!(
1641 "Unexpected {level_type} level bit width after decompression: expected 16, got {}",
1642 block.bits_per_value
1643 )
1644 .into(),
1645 ));
1646 }
1647 let expected_bytes = expected_level_bytes(num_values, level_type)?;
1648 if block.data.len() != expected_bytes {
1649 return Err(Error::invalid_input_source(
1650 format!(
1651 "Unexpected decompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1652 block.data.len()
1653 )
1654 .into(),
1655 ));
1656 }
1657 Ok(LazyLevels::Dense(block.data.borrow_to_typed_slice::<u16>()))
1658}
1659
1660#[derive(Debug)]
1661struct CachedComplexAllNullState {
1662 rep: Option<LazyLevels>,
1663 def: Option<LazyLevels>,
1664}
1665
1666impl DeepSizeOf for CachedComplexAllNullState {
1667 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1668 self.rep.deep_size_of_children(ctx) + self.def.deep_size_of_children(ctx)
1669 }
1670}
1671
1672impl CachedPageData for CachedComplexAllNullState {
1673 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1674 self
1675 }
1676}
1677
1678#[derive(Debug)]
1687pub struct ComplexAllNullScheduler {
1688 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1690 def_meaning: Arc<[DefinitionInterpretation]>,
1691 repdef: Option<Arc<CachedComplexAllNullState>>,
1692 max_rep: u16,
1693 max_def: u16,
1694 max_visible_level: u16,
1695 rep_codec: LevelCodec,
1696 def_codec: LevelCodec,
1697 num_rep_values: u64,
1698 num_def_values: u64,
1699}
1700
1701impl ComplexAllNullScheduler {
1702 pub(crate) fn new(
1703 buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1704 def_meaning: Arc<[DefinitionInterpretation]>,
1705 rep_codec: LevelCodec,
1706 def_codec: LevelCodec,
1707 num_rep_values: u64,
1708 num_def_values: u64,
1709 ) -> Self {
1710 let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
1711 let max_def = def_meaning
1712 .iter()
1713 .map(|meaning| meaning.num_def_levels())
1714 .sum::<u16>();
1715 let max_visible_level = max_visible_level(&def_meaning);
1716 Self {
1717 buffer_offsets_and_sizes,
1718 def_meaning,
1719 repdef: None,
1720 max_rep,
1721 max_def,
1722 max_visible_level,
1723 rep_codec,
1724 def_codec,
1725 num_rep_values,
1726 num_def_values,
1727 }
1728 }
1729}
1730
1731impl StructuralPageScheduler for ComplexAllNullScheduler {
1732 fn init_layout(&self) -> Result<PageInitialization> {
1733 let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
1735 let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
1736
1737 Ok(PageInitialization::new([
1738 (rep_size > 0).then_some(rep_pos..rep_pos + rep_size),
1739 (def_size > 0).then_some(def_pos..def_pos + def_size),
1740 ]))
1741 }
1742
1743 fn init_from_buffers<'a>(
1744 &'a mut self,
1745 buffers: PageInitializationBuffers,
1746 _io: &Arc<dyn EncodingsIo>,
1747 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1748 let rep_codec = self.rep_codec.clone();
1749 let def_codec = self.def_codec.clone();
1750 let num_rep_values = self.num_rep_values;
1751 let num_def_values = self.num_def_values;
1752 let max_rep = self.max_rep;
1753 let max_def = self.max_def;
1754
1755 async move {
1756 let [rep_bytes, def_bytes] = buffers.try_into_array("complex all-null")?;
1757
1758 let build_levels = |compressed_bytes: Bytes,
1761 codec: &LevelCodec,
1762 num_values: u64,
1763 level_type: &str|
1764 -> Result<LazyLevels> {
1765 match codec {
1766 LevelCodec::Uncompressed => {
1767 if num_values == 0 {
1768 if !compressed_bytes
1769 .len()
1770 .is_multiple_of(std::mem::size_of::<u16>())
1771 {
1772 return Err(Error::invalid_input_source(
1773 format!(
1774 "Unexpected uncompressed {level_type} level size: {} bytes is not divisible by {}",
1775 compressed_bytes.len(),
1776 std::mem::size_of::<u16>()
1777 )
1778 .into(),
1779 ));
1780 }
1781 } else {
1782 let expected_bytes = expected_level_bytes(num_values, level_type)?;
1783 if compressed_bytes.len() != expected_bytes {
1784 return Err(Error::invalid_input_source(
1785 format!(
1786 "Unexpected uncompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1787 compressed_bytes.len()
1788 )
1789 .into(),
1790 ));
1791 }
1792 }
1793 let buffer = LanceBuffer::from_bytes(compressed_bytes, 2);
1794 Ok(LazyLevels::Dense(buffer.borrow_to_typed_slice::<u16>()))
1795 }
1796 LevelCodec::Rle(decompressor) => {
1797 let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1798 let runs = decompressor.decode_u16_runs(frame, num_values)?;
1799 LazyLevels::from_rle_runs(runs)
1800 }
1801 LevelCodec::Block(decompressor) => {
1802 let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1803 let decompressed = decompressor.decompress(Some(frame), num_values)?;
1804 dense_levels_from_block(decompressed, num_values, level_type)
1805 }
1806 }
1807 };
1808
1809 let rep = if let Some(rep) = rep_bytes {
1810 Some(build_levels(rep, &rep_codec, num_rep_values, "repetition")?)
1811 } else {
1812 None
1813 };
1814
1815 let def = if let Some(def) = def_bytes {
1816 Some(build_levels(def, &def_codec, num_def_values, "definition")?)
1817 } else {
1818 None
1819 };
1820
1821 validate_complex_all_null_levels(&rep, &def, max_rep, max_def)?;
1822 let repdef = Arc::new(CachedComplexAllNullState { rep, def });
1823
1824 Ok(repdef as Arc<dyn CachedPageData>)
1825 }
1826 .boxed()
1827 }
1828
1829 fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
1830 self.repdef = Some(data.clone().as_arc_any().downcast().map_err(|_| {
1831 Error::invalid_input_source(
1832 "Cached complex all-null page data has an unexpected type".into(),
1833 )
1834 })?);
1835 Ok(())
1836 }
1837
1838 fn schedule_ranges(
1839 &self,
1840 ranges: &[Range<u64>],
1841 _io: &Arc<dyn EncodingsIo>,
1842 ) -> Result<Vec<PageLoadTask>> {
1843 let ranges = VecDeque::from_iter(ranges.iter().cloned());
1844 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1845 let decoder = Box::new(ComplexAllNullPageDecoder {
1846 ranges,
1847 rep: self.repdef.as_ref().unwrap().rep.clone(),
1848 def: self.repdef.as_ref().unwrap().def.clone(),
1849 num_rows,
1850 def_meaning: self.def_meaning.clone(),
1851 max_rep: self.max_rep,
1852 max_visible_level: self.max_visible_level,
1853 rep_cursor: LevelCursor::default(),
1854 def_run_cursor: RunPosition::default(),
1855 }) as Box<dyn StructuralPageDecoder>;
1856 let page_load_task = PageLoadTask {
1857 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1858 num_rows,
1859 };
1860 Ok(vec![page_load_task])
1861 }
1862}
1863
1864#[derive(Debug)]
1865pub struct ComplexAllNullPageDecoder {
1866 ranges: VecDeque<Range<u64>>,
1867 rep: Option<LazyLevels>,
1868 def: Option<LazyLevels>,
1869 num_rows: u64,
1870 def_meaning: Arc<[DefinitionInterpretation]>,
1871 max_rep: u16,
1872 max_visible_level: u16,
1873 rep_cursor: LevelCursor,
1875 def_run_cursor: RunPosition,
1877}
1878
1879impl ComplexAllNullPageDecoder {
1880 fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
1881 let mut rows_desired = num_rows;
1882 let mut ranges = Vec::with_capacity(self.ranges.len());
1883 while rows_desired > 0 {
1884 let front = self.ranges.front_mut().unwrap();
1885 let avail = front.end - front.start;
1886 if avail > rows_desired {
1887 ranges.push(front.start..front.start + rows_desired);
1888 front.start += rows_desired;
1889 rows_desired = 0;
1890 } else {
1891 ranges.push(self.ranges.pop_front().unwrap());
1892 rows_desired -= avail;
1893 }
1894 }
1895 ranges
1896 }
1897
1898 fn seek_row_start(&mut self, target_row: u64) -> Result<usize> {
1901 match &self.rep {
1902 Some(rep) => rep.seek_row_start(&mut self.rep_cursor, target_row, self.max_rep),
1903 None => {
1904 self.rep_cursor.row = target_row;
1906 self.rep_cursor.level = target_row as usize;
1907 Ok(target_row as usize)
1908 }
1909 }
1910 }
1911
1912 fn count_visible(&mut self, levels: Range<usize>) -> Result<(u64, RunPosition)> {
1915 match &self.def {
1916 Some(def) => {
1917 if levels.end > def.len() {
1918 return Err(Error::internal(
1919 "Invalid complex all-null layout: definition buffer too short",
1920 ));
1921 }
1922 Ok(def.count_le_cursor(&mut self.def_run_cursor, levels, self.max_visible_level))
1923 }
1924 None => Ok(((levels.end - levels.start) as u64, RunPosition::default())),
1925 }
1926 }
1927}
1928
1929impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1930 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1931 let drained_ranges = self.drain_ranges(num_rows);
1932 let mut level_slices: Vec<LevelSlice> = Vec::with_capacity(drained_ranges.len());
1933 let mut visible_items_total = 0;
1934
1935 for range in drained_ranges {
1940 let level_start = self.seek_row_start(range.start)?;
1941 let rep_run = self.rep_cursor.run;
1942 let level_end = self.seek_row_start(range.end)?;
1943 let (visible_items, def_run) = self.count_visible(level_start..level_end)?;
1944 visible_items_total += visible_items;
1945 if let Some(last) = level_slices.last_mut()
1946 && last.range.end == level_start
1947 {
1948 last.range.end = level_end;
1949 } else {
1950 level_slices.push(LevelSlice {
1951 range: level_start..level_end,
1952 rep_run,
1953 def_run,
1954 });
1955 }
1956 }
1957
1958 Ok(Box::new(DecodeComplexAllNullTask {
1959 level_slices,
1960 visible_items_total,
1961 rep: self.rep.clone(),
1962 def: self.def.clone(),
1963 def_meaning: self.def_meaning.clone(),
1964 max_visible_level: self.max_visible_level,
1965 }))
1966 }
1967
1968 fn num_rows(&self) -> u64 {
1969 self.num_rows
1970 }
1971}
1972
1973#[derive(Debug, Clone)]
1976struct LevelSlice {
1977 range: Range<usize>,
1978 rep_run: RunPosition,
1979 def_run: RunPosition,
1980}
1981
1982#[derive(Clone, Copy)]
1983enum LevelKind {
1984 Repetition,
1985 Definition,
1986}
1987
1988impl LevelSlice {
1989 fn run(&self, kind: LevelKind) -> RunPosition {
1990 match kind {
1991 LevelKind::Repetition => self.rep_run,
1992 LevelKind::Definition => self.def_run,
1993 }
1994 }
1995}
1996
1997#[derive(Debug)]
1998pub struct DecodeComplexAllNullTask {
1999 level_slices: Vec<LevelSlice>,
2000 visible_items_total: u64,
2001 rep: Option<LazyLevels>,
2002 def: Option<LazyLevels>,
2003 def_meaning: Arc<[DefinitionInterpretation]>,
2004 max_visible_level: u16,
2005}
2006
2007impl DecodeComplexAllNullTask {
2008 fn decode_level(&self, levels: &Option<LazyLevels>, kind: LevelKind) -> Option<Vec<u16>> {
2009 levels.as_ref().map(|levels| {
2010 let num_levels = self
2011 .level_slices
2012 .iter()
2013 .map(|slice| slice.range.end - slice.range.start)
2014 .sum();
2015 let mut referenced_levels = Vec::with_capacity(num_levels);
2016 for slice in &self.level_slices {
2017 levels.extend_into(slice.range.clone(), slice.run(kind), &mut referenced_levels);
2018 }
2019 referenced_levels
2020 })
2021 }
2022}
2023
2024impl DecodePageTask for DecodeComplexAllNullTask {
2025 fn decode(self: Box<Self>) -> Result<DecodedPage> {
2026 let rep = self.decode_level(&self.rep, LevelKind::Repetition);
2027 let def = self.decode_level(&self.def, LevelKind::Definition);
2028
2029 let num_values = if let Some(def) = &def {
2033 def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
2034 } else {
2035 self.visible_items_total
2036 };
2037
2038 let data = DataBlock::AllNull(AllNullDataBlock { num_values });
2039 let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
2040 Ok(DecodedPage {
2041 data,
2042 repdef: unraveler,
2043 })
2044 }
2045}
2046
2047#[derive(Debug, Default)]
2052pub struct SimpleAllNullScheduler {}
2053
2054impl StructuralPageScheduler for SimpleAllNullScheduler {
2055 fn needs_initialization(&self) -> bool {
2056 false
2057 }
2058
2059 fn init_layout(&self) -> Result<PageInitialization> {
2060 Ok(PageInitialization::new([]))
2061 }
2062
2063 fn init_from_buffers<'a>(
2064 &'a mut self,
2065 _buffers: PageInitializationBuffers,
2066 _io: &Arc<dyn EncodingsIo>,
2067 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2068 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
2069 }
2070
2071 fn try_load(&mut self, _cache: &Arc<dyn CachedPageData>) -> Result<()> {
2072 Ok(())
2073 }
2074
2075 fn schedule_ranges(
2076 &self,
2077 ranges: &[Range<u64>],
2078 _io: &Arc<dyn EncodingsIo>,
2079 ) -> Result<Vec<PageLoadTask>> {
2080 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
2081 let decoder =
2082 Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
2083 let page_load_task = PageLoadTask {
2084 decoder_fut: std::future::ready(Ok(decoder)).boxed(),
2085 num_rows,
2086 };
2087 Ok(vec![page_load_task])
2088 }
2089}
2090
2091#[derive(Debug)]
2094struct SimpleAllNullDecodePageTask {
2095 num_values: u64,
2096}
2097impl DecodePageTask for SimpleAllNullDecodePageTask {
2098 fn decode(self: Box<Self>) -> Result<DecodedPage> {
2099 let unraveler = RepDefUnraveler::new(
2100 None,
2101 Some(vec![1; self.num_values as usize]),
2102 Arc::new([DefinitionInterpretation::NullableItem]),
2103 self.num_values,
2104 );
2105 Ok(DecodedPage {
2106 data: DataBlock::AllNull(AllNullDataBlock {
2107 num_values: self.num_values,
2108 }),
2109 repdef: unraveler,
2110 })
2111 }
2112}
2113
2114#[derive(Debug)]
2115pub struct SimpleAllNullPageDecoder {
2116 num_rows: u64,
2117}
2118
2119impl StructuralPageDecoder for SimpleAllNullPageDecoder {
2120 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2121 Ok(Box::new(SimpleAllNullDecodePageTask {
2122 num_values: num_rows,
2123 }))
2124 }
2125
2126 fn num_rows(&self) -> u64 {
2127 self.num_rows
2128 }
2129}
2130
2131#[derive(Debug, Clone)]
2132struct MiniBlockSchedulerDictionary {
2133 dictionary_decompressor: Arc<dyn BlockDecompressor>,
2135 dictionary_buf_position_and_size: (u64, u64),
2136 dictionary_data_alignment: u64,
2137 num_dictionary_items: u64,
2138}
2139
2140#[derive(Debug)]
2142struct MiniBlockCacheableState {
2143 chunk_index: MiniBlockChunkIndex,
2145 dictionary: Option<Arc<DataBlock>>,
2147}
2148
2149impl DeepSizeOf for MiniBlockCacheableState {
2150 fn deep_size_of_children(&self, context: &mut Context) -> usize {
2151 self.chunk_index.deep_size_of_children(context)
2152 + self
2153 .dictionary
2154 .as_ref()
2155 .map(|dict| dict.data_size() as usize)
2156 .unwrap_or(0)
2157 }
2158}
2159
2160impl CachedPageData for MiniBlockCacheableState {
2161 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
2162 self
2163 }
2164}
2165
2166#[derive(Debug)]
2193pub struct MiniBlockScheduler {
2194 buffer_offsets_and_sizes: Vec<(u64, u64)>,
2196 priority: u64,
2197 items_in_page: u64,
2198 repetition_index_depth: u16,
2199 num_buffers: u64,
2200 rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
2201 def_decompressor: Option<Arc<dyn BlockDecompressor>>,
2202 value_decompressor: Arc<dyn MiniBlockDecompressor>,
2203 def_meaning: Arc<[DefinitionInterpretation]>,
2204 dictionary: Option<MiniBlockSchedulerDictionary>,
2205 page_meta: Option<Arc<MiniBlockCacheableState>>,
2207 has_large_chunk: bool,
2208}
2209
2210impl MiniBlockScheduler {
2211 fn try_new(
2212 buffer_offsets_and_sizes: &[(u64, u64)],
2213 priority: u64,
2214 items_in_page: u64,
2215 layout: &pb21::MiniBlockLayout,
2216 decompressors: &dyn DecompressionStrategy,
2217 ) -> Result<Self> {
2218 let rep_decompressor = layout
2219 .rep_compression
2220 .as_ref()
2221 .map(|rep_compression| {
2222 decompressors
2223 .create_block_decompressor(rep_compression)
2224 .map(Arc::from)
2225 })
2226 .transpose()?;
2227 let def_decompressor = layout
2228 .def_compression
2229 .as_ref()
2230 .map(|def_compression| {
2231 decompressors
2232 .create_block_decompressor(def_compression)
2233 .map(Arc::from)
2234 })
2235 .transpose()?;
2236 let def_meaning = layout
2237 .layers
2238 .iter()
2239 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2240 .collect::<Vec<_>>();
2241 let value_decompressor = decompressors.create_miniblock_decompressor(
2242 layout.value_compression.as_ref().unwrap(),
2243 decompressors,
2244 )?;
2245
2246 let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
2247 let num_dictionary_items = layout.num_dictionary_items;
2248 let dictionary_decompressor = decompressors
2249 .create_block_decompressor(dictionary_encoding)?
2250 .into();
2251 let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
2252 {
2253 Compression::Variable(_) => 4,
2254 Compression::Flat(_) => 16,
2255 Compression::General(_) => 1,
2256 Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
2257 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
2258 }
2259 _ => {
2260 return Err(Error::invalid_input_source(
2261 format!(
2262 "Unsupported mini-block dictionary encoding: {:?}",
2263 dictionary_encoding.compression.as_ref().unwrap()
2264 )
2265 .into(),
2266 ));
2267 }
2268 };
2269 Some(MiniBlockSchedulerDictionary {
2270 dictionary_decompressor,
2271 dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
2272 dictionary_data_alignment,
2273 num_dictionary_items,
2274 })
2275 } else {
2276 None
2277 };
2278
2279 Ok(Self {
2280 buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
2281 rep_decompressor,
2282 def_decompressor,
2283 value_decompressor: value_decompressor.into(),
2284 repetition_index_depth: layout.repetition_index_depth as u16,
2285 num_buffers: layout.num_buffers,
2286 priority,
2287 items_in_page,
2288 dictionary,
2289 def_meaning: def_meaning.into(),
2290 page_meta: None,
2291 has_large_chunk: layout.has_large_chunk,
2292 })
2293 }
2294
2295 fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
2296 let chunk_index = &self.page_meta.as_ref().unwrap().chunk_index;
2297 chunk_indices
2298 .iter()
2299 .map(|&chunk_idx| LoadedChunk {
2300 byte_range: chunk_index.byte_range(chunk_idx),
2301 items_in_chunk: chunk_index.items_in_chunk(chunk_idx),
2302 chunk_idx,
2303 data: LanceBuffer::empty(),
2304 })
2305 .collect()
2306 }
2307}
2308
2309#[derive(Debug, PartialEq, Eq, Clone, Copy)]
2310enum PreambleAction {
2311 Take,
2312 Skip,
2313 Absent,
2314}
2315
2316#[derive(Clone, Debug, PartialEq, Eq)]
2351struct ChunkInstructions {
2352 chunk_idx: usize,
2354 preamble: PreambleAction,
2360 rows_to_skip: u64,
2364 rows_to_take: u64,
2367 take_trailer: bool,
2374}
2375
2376#[derive(Debug, PartialEq, Eq)]
2394struct ChunkDrainInstructions {
2395 chunk_instructions: ChunkInstructions,
2396 rows_to_skip: u64,
2397 rows_to_take: u64,
2398 preamble_action: PreambleAction,
2399}
2400
2401impl ChunkInstructions {
2402 fn schedule_instructions(
2408 chunk_index: &MiniBlockChunkIndex,
2409 user_ranges: &[Range<u64>],
2410 ) -> Vec<Self> {
2411 let num_chunks = chunk_index.num_chunks();
2414 let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
2418
2419 for user_range in user_ranges {
2420 let mut rows_needed = user_range.end - user_range.start;
2421 let mut need_preamble = false;
2422
2423 let mut block_index = chunk_index.find_chunk(user_range.start);
2426
2427 let mut to_skip = user_range.start - chunk_index.first_row(block_index);
2428
2429 while rows_needed > 0 || need_preamble {
2430 if block_index >= num_chunks {
2432 log::warn!(
2433 "schedule_instructions inconsistency: block_index >= num_chunks, exiting early"
2434 );
2435 break;
2436 }
2437
2438 let starts_including_trailer = chunk_index.rows_in_chunk(block_index);
2439 let has_preamble = chunk_index.has_preamble(block_index);
2440 let has_trailer = chunk_index.has_trailer(block_index);
2441 let rows_avail = starts_including_trailer.saturating_sub(to_skip);
2442
2443 if rows_avail == 0 && to_skip == 0 {
2447 if has_preamble && need_preamble {
2449 chunk_instructions.push(Self {
2450 chunk_idx: block_index,
2451 preamble: PreambleAction::Take,
2452 rows_to_skip: 0,
2453 rows_to_take: 0,
2454 take_trailer: has_trailer,
2458 });
2459 if starts_including_trailer > 0 || block_index == num_chunks - 1 {
2463 need_preamble = false;
2464 }
2465 }
2466 block_index += 1;
2468 continue;
2469 }
2470
2471 if rows_avail == 0 && to_skip > 0 {
2475 to_skip -= starts_including_trailer;
2478 block_index += 1;
2479 continue;
2480 }
2481
2482 let rows_to_take = rows_avail.min(rows_needed);
2483 rows_needed -= rows_to_take;
2484
2485 let mut take_trailer = false;
2486 let preamble = if has_preamble {
2487 if need_preamble {
2488 PreambleAction::Take
2489 } else {
2490 PreambleAction::Skip
2491 }
2492 } else {
2493 PreambleAction::Absent
2494 };
2495
2496 if rows_to_take == rows_avail && has_trailer {
2498 take_trailer = true;
2499 need_preamble = true;
2500 } else {
2501 need_preamble = false;
2502 };
2503
2504 chunk_instructions.push(Self {
2505 preamble,
2506 chunk_idx: block_index,
2507 rows_to_skip: to_skip,
2508 rows_to_take,
2509 take_trailer,
2510 });
2511
2512 to_skip = 0;
2513 block_index += 1;
2514 }
2515 }
2516
2517 if user_ranges.len() > 1 {
2521 let mut write = 0;
2525 for read in 1..chunk_instructions.len() {
2526 let merges = {
2527 let last = &chunk_instructions[write];
2528 let candidate = &chunk_instructions[read];
2529 last.chunk_idx == candidate.chunk_idx
2530 && last.rows_to_take + last.rows_to_skip == candidate.rows_to_skip
2531 };
2532 if merges {
2533 let rows_to_take = chunk_instructions[read].rows_to_take;
2534 let take_trailer = chunk_instructions[read].take_trailer;
2535 let last = &mut chunk_instructions[write];
2536 last.rows_to_take += rows_to_take;
2537 last.take_trailer |= take_trailer;
2538 } else {
2539 write += 1;
2540 if write != read {
2541 chunk_instructions.swap(write, read);
2542 }
2543 }
2544 }
2545 chunk_instructions.truncate(write + 1);
2546 }
2547 chunk_instructions
2548 }
2549
2550 fn drain_from_instruction(
2551 &self,
2552 rows_desired: &mut u64,
2553 need_preamble: &mut bool,
2554 skip_in_chunk: &mut u64,
2555 ) -> (ChunkDrainInstructions, bool) {
2556 debug_assert!(!*need_preamble || *skip_in_chunk == 0);
2558 let rows_avail = self.rows_to_take - *skip_in_chunk;
2559 let has_preamble = self.preamble != PreambleAction::Absent;
2560 let preamble_action = match (*need_preamble, has_preamble) {
2561 (true, true) => PreambleAction::Take,
2562 (true, false) => panic!("Need preamble but there isn't one"),
2563 (false, true) => PreambleAction::Skip,
2564 (false, false) => PreambleAction::Absent,
2565 };
2566
2567 let rows_taking = if *rows_desired >= rows_avail {
2570 *need_preamble = self.take_trailer;
2578 rows_avail
2579 } else {
2580 *need_preamble = false;
2583 *rows_desired
2584 };
2585 let rows_skipped = *skip_in_chunk;
2586
2587 let consumed_chunk = if *rows_desired >= rows_avail {
2589 *rows_desired -= rows_avail;
2590 *skip_in_chunk = 0;
2591 true
2592 } else {
2593 *skip_in_chunk += *rows_desired;
2594 *rows_desired = 0;
2595 false
2596 };
2597
2598 (
2599 ChunkDrainInstructions {
2600 chunk_instructions: self.clone(),
2601 rows_to_skip: rows_skipped,
2602 rows_to_take: rows_taking,
2603 preamble_action,
2604 },
2605 consumed_chunk,
2606 )
2607 }
2608}
2609
2610enum Words {
2611 U16(ScalarBuffer<u16>),
2612 U32(ScalarBuffer<u32>),
2613}
2614
2615struct WordsIter<'a> {
2616 iter: Box<dyn Iterator<Item = u32> + 'a>,
2617}
2618
2619impl Words {
2620 pub fn len(&self) -> usize {
2621 match self {
2622 Self::U16(b) => b.len(),
2623 Self::U32(b) => b.len(),
2624 }
2625 }
2626
2627 pub fn iter(&self) -> WordsIter<'_> {
2628 match self {
2629 Self::U16(buf) => WordsIter {
2630 iter: Box::new(buf.iter().map(|&x| x as u32)),
2631 },
2632 Self::U32(buf) => WordsIter {
2633 iter: Box::new(buf.iter().copied()),
2634 },
2635 }
2636 }
2637
2638 pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
2639 let bytes_per_value = if has_large_chunk { 4 } else { 2 };
2640 assert_eq!(bytes.len() % bytes_per_value, 0);
2641 let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
2642 if has_large_chunk {
2643 Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
2644 } else {
2645 Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
2646 }
2647 }
2648}
2649
2650impl<'a> Iterator for WordsIter<'a> {
2651 type Item = u32;
2652
2653 fn next(&mut self) -> Option<Self::Item> {
2654 self.iter.next()
2655 }
2656}
2657
2658struct FlatValueCounts {
2663 logs: Vec<u8>,
2664 uniform: bool,
2665 values_per_chunk: u64,
2666 last_chunk_values: u64,
2667}
2668
2669fn analyze_value_counts(words: &Words, items_in_page: u64) -> Result<FlatValueCounts> {
2670 let num_chunks = words.len();
2671 let logs = words.iter().map(|w| (w & 0x0F) as u8).collect::<Vec<_>>();
2672 let mut counted = 0u64;
2673 for (chunk_index, &log) in logs.iter().take(num_chunks.saturating_sub(1)).enumerate() {
2674 if log == 0 {
2675 return Err(Error::corrupt_file_named(
2676 "miniblock_metadata",
2677 format!(
2678 "non-final chunk {chunk_index} of {num_chunks} has invalid log_num_values=0"
2679 ),
2680 ));
2681 }
2682 counted = counted.checked_add(1u64 << log).ok_or_else(|| {
2683 Error::corrupt_file_named(
2684 "miniblock_metadata",
2685 format!(
2686 "value count overflow at chunk {chunk_index}: counted_values={counted}, \
2687 log_num_values={log}, items_in_page={items_in_page}"
2688 ),
2689 )
2690 })?;
2691 }
2692 let last_chunk_values = items_in_page.checked_sub(counted).ok_or_else(|| {
2693 Error::corrupt_file_named(
2694 "miniblock_metadata",
2695 format!(
2696 "non-final chunks account for counted_values={counted}, exceeding \
2697 items_in_page={items_in_page}"
2698 ),
2699 )
2700 })?;
2701 if let Some(&last_log) = logs.last()
2702 && last_log != 0
2703 && (1u64 << last_log) != last_chunk_values
2704 {
2705 return Err(Error::corrupt_file_named(
2706 "miniblock_metadata",
2707 format!(
2708 "final chunk log_num_values={last_log} does not match \
2709 last_chunk_values={last_chunk_values}: counted_values={counted}, \
2710 items_in_page={items_in_page}"
2711 ),
2712 ));
2713 }
2714 let uniform = num_chunks <= 1 || logs[..num_chunks - 1].iter().all(|&log| log == logs[0]);
2715 let values_per_chunk = if num_chunks <= 1 {
2718 items_in_page.max(1)
2719 } else {
2720 1u64 << logs[0]
2721 };
2722 Ok(FlatValueCounts {
2723 logs,
2724 uniform,
2725 values_per_chunk,
2726 last_chunk_values,
2727 })
2728}
2729
2730fn flat_value_counts_iter(logs: &[u8], last_chunk_values: u64) -> impl Iterator<Item = u64> + '_ {
2733 let num_chunks = logs.len();
2734 (0..num_chunks).map(move |i| {
2735 if i + 1 < num_chunks {
2736 1u64 << logs[i]
2737 } else {
2738 last_chunk_values
2739 }
2740 })
2741}
2742
2743fn build_chunk_index(
2748 words: &Words,
2749 items_in_page: u64,
2750 base: u64,
2751 data_buf_size: u64,
2752 rep_index_bytes: Option<&[u8]>,
2753 repetition_index_depth: u16,
2754) -> Result<MiniBlockChunkIndex> {
2755 let num_chunks = words.len();
2756 let value_counts = analyze_value_counts(words, items_in_page)?;
2759
2760 let byte_starts = PrefixSums::from_deltas(
2763 words
2764 .iter()
2765 .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64),
2766 num_chunks,
2767 data_buf_size,
2768 );
2769
2770 let rows = if let Some(rep_index_data) = rep_index_bytes {
2773 assert!(rep_index_data.len() % 8 == 0);
2774 let stride = repetition_index_depth as usize + 1;
2775 let (row_starts, has_trailer) = parse_nested_rep(rep_index_data, stride);
2776 let item_counts = if value_counts.uniform {
2777 ItemCounts::Uniform {
2778 values_per_chunk: value_counts.values_per_chunk,
2779 last_chunk_values: value_counts.last_chunk_values,
2780 }
2781 } else {
2782 ItemCounts::PerChunkLog {
2783 logs: value_counts.logs,
2784 last_chunk_values: value_counts.last_chunk_values,
2785 }
2786 };
2787 RowMapping::Nested {
2788 row_starts,
2789 has_trailer,
2790 item_counts,
2791 }
2792 } else {
2793 if value_counts.uniform {
2794 RowMapping::UniformFlat {
2795 values_per_chunk: value_counts.values_per_chunk,
2796 last_chunk_values: value_counts.last_chunk_values,
2797 num_chunks,
2798 }
2799 } else {
2800 let value_starts = PrefixSums::from_deltas(
2801 flat_value_counts_iter(&value_counts.logs, value_counts.last_chunk_values),
2802 num_chunks,
2803 items_in_page,
2804 );
2805 RowMapping::Flat { value_starts }
2806 }
2807 };
2808
2809 Ok(MiniBlockChunkIndex::new(base, byte_starts, rows))
2810}
2811
2812impl StructuralPageScheduler for MiniBlockScheduler {
2813 fn init_layout(&self) -> Result<PageInitialization> {
2814 let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
2818 let dictionary_range = self.dictionary.as_ref().map(|dictionary| {
2819 dictionary.dictionary_buf_position_and_size.0
2820 ..dictionary.dictionary_buf_position_and_size.0
2821 + dictionary.dictionary_buf_position_and_size.1
2822 });
2823 let repetition_index_range = if self.repetition_index_depth > 0 {
2824 let (rep_index_pos, rep_index_size) =
2825 self.buffer_offsets_and_sizes.last().ok_or_else(|| {
2826 Error::invalid_input_source(
2827 "Mini-block page has a repetition index but no buffers".into(),
2828 )
2829 })?;
2830 Some(*rep_index_pos..*rep_index_pos + *rep_index_size)
2831 } else {
2832 None
2833 };
2834 Ok(PageInitialization::new([
2835 Some(meta_buf_position..meta_buf_position + meta_buf_size),
2836 dictionary_range,
2837 repetition_index_range,
2838 ]))
2839 }
2840
2841 fn init_from_buffers<'a>(
2842 &'a mut self,
2843 buffers: PageInitializationBuffers,
2844 _io: &Arc<dyn EncodingsIo>,
2845 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2846 let base = self.buffer_offsets_and_sizes[1].0;
2847 let data_buf_size = self.buffer_offsets_and_sizes[1].1;
2848
2849 async move {
2850 let [meta_bytes, dictionary_bytes, rep_index_bytes] =
2851 buffers.try_into_array("mini-block")?;
2852 let meta_bytes = meta_bytes.ok_or_else(|| {
2853 Error::internal("Mini-block page initialization is missing chunk metadata")
2854 })?;
2855
2856 let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
2857 let chunk_index = build_chunk_index(
2858 &words,
2859 self.items_in_page,
2860 base,
2861 data_buf_size,
2862 rep_index_bytes.as_deref(),
2863 self.repetition_index_depth,
2864 )?;
2865
2866 let dictionary = match (self.dictionary.as_ref(), dictionary_bytes) {
2868 (Some(dictionary), Some(dictionary_data)) => {
2869 Some(Arc::new(dictionary.dictionary_decompressor.decompress(
2870 Some(LanceBuffer::from_bytes(
2871 dictionary_data,
2872 dictionary.dictionary_data_alignment,
2873 )),
2874 dictionary.num_dictionary_items,
2875 )?))
2876 }
2877 (None, None) => None,
2878 _ => {
2879 return Err(Error::internal(
2880 "Mini-block page initialization dictionary slot does not match its layout",
2881 ));
2882 }
2883 };
2884
2885 let page_meta = Arc::new(MiniBlockCacheableState {
2886 chunk_index,
2887 dictionary,
2888 });
2889 Ok(page_meta as Arc<dyn CachedPageData>)
2890 }
2891 .boxed()
2892 }
2893
2894 fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
2895 self.page_meta = Some(data.clone().as_arc_any().downcast().map_err(|_| {
2896 Error::invalid_input_source("Cached mini-block page data has an unexpected type".into())
2897 })?);
2898 Ok(())
2899 }
2900
2901 fn schedule_ranges(
2902 &self,
2903 ranges: &[Range<u64>],
2904 io: &Arc<dyn EncodingsIo>,
2905 ) -> Result<Vec<PageLoadTask>> {
2906 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2907
2908 let page_meta = self.page_meta.as_ref().unwrap();
2909
2910 let chunk_instructions =
2911 ChunkInstructions::schedule_instructions(&page_meta.chunk_index, ranges);
2912
2913 debug_assert_eq!(
2914 num_rows,
2915 chunk_instructions
2916 .iter()
2917 .map(|ci| ci.rows_to_take)
2918 .sum::<u64>()
2919 );
2920
2921 let chunks_needed = chunk_instructions
2922 .iter()
2923 .map(|ci| ci.chunk_idx)
2924 .unique()
2925 .collect::<Vec<_>>();
2926
2927 let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
2928 let chunk_ranges = loaded_chunks
2929 .iter()
2930 .map(|c| c.byte_range.clone())
2931 .collect::<Vec<_>>();
2932 let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
2933
2934 let rep_decompressor = self.rep_decompressor.clone();
2935 let def_decompressor = self.def_decompressor.clone();
2936 let value_decompressor = self.value_decompressor.clone();
2937 let num_buffers = self.num_buffers;
2938 let has_large_chunk = self.has_large_chunk;
2939 let dictionary = page_meta
2940 .dictionary
2941 .as_ref()
2942 .map(|dictionary| dictionary.clone());
2943 let def_meaning = self.def_meaning.clone();
2944
2945 let res = async move {
2946 let loaded_chunk_data = loaded_chunk_data.await?;
2947 for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
2948 loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
2949 }
2950
2951 Ok(Box::new(MiniBlockDecoder {
2952 rep_decompressor,
2953 def_decompressor,
2954 value_decompressor,
2955 def_meaning,
2956 loaded_chunks: VecDeque::from_iter(loaded_chunks),
2957 instructions: VecDeque::from(chunk_instructions),
2958 offset_in_current_chunk: 0,
2959 dictionary,
2960 num_rows,
2961 num_buffers,
2962 has_large_chunk,
2963 }) as Box<dyn StructuralPageDecoder>)
2964 }
2965 .boxed();
2966 let page_load_task = PageLoadTask {
2967 decoder_fut: res,
2968 num_rows,
2969 };
2970 Ok(vec![page_load_task])
2971 }
2972}
2973
2974#[derive(Debug, Clone, Copy)]
2975struct FullZipRepIndexDetails {
2976 buf_position: u64,
2977 bytes_per_value: u64, }
2979
2980#[derive(Debug)]
2981enum PerValueDecompressor {
2982 Fixed(Arc<dyn FixedPerValueDecompressor>),
2983 Variable(Arc<dyn VariablePerValueDecompressor>),
2984}
2985
2986#[derive(Debug)]
2987struct FullZipDecodeDetails {
2988 value_decompressor: PerValueDecompressor,
2989 def_meaning: Arc<[DefinitionInterpretation]>,
2990 ctrl_word_parser: ControlWordParser,
2991 max_rep: u16,
2992 max_visible_def: u16,
2993}
2994
2995#[derive(Debug, Clone)]
3007enum FullZipReadSource {
3008 Remote(Arc<dyn EncodingsIo>),
3010 PrefetchedPage { base_offset: u64, data: LanceBuffer },
3012}
3013
3014impl FullZipReadSource {
3015 fn fetch(
3019 &self,
3020 ranges: &[Range<u64>],
3021 priority: u64,
3022 ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
3023 match self {
3024 Self::Remote(io) => {
3025 let io = io.clone();
3026 let ranges = ranges.to_vec();
3027 async move {
3028 let data = io.submit_request(ranges, priority).await?;
3029 Ok(data
3030 .into_iter()
3031 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
3032 .collect::<VecDeque<_>>())
3033 }
3034 .boxed()
3035 }
3036 Self::PrefetchedPage { base_offset, data } => {
3037 let base_offset = *base_offset;
3038 let data = data.clone();
3039 let page_end = base_offset + data.len() as u64;
3040 std::future::ready(
3041 ranges
3042 .iter()
3043 .map(|range| {
3044 if range.start > range.end
3045 || range.start < base_offset
3046 || range.end > page_end
3047 {
3048 return Err(Error::internal(format!(
3049 "Requested range {:?} is outside page range {}..{}",
3050 range, base_offset, page_end
3051 )));
3052 }
3053 let start = (range.start - base_offset) as usize;
3054 let len = (range.end - range.start) as usize;
3055 Ok(data.slice_with_length(start, len))
3056 })
3057 .collect::<Result<VecDeque<_>>>(),
3058 )
3059 .boxed()
3060 }
3061 }
3062 }
3063}
3064
3065#[derive(Debug)]
3073pub struct FullZipScheduler {
3074 data_buf_position: u64,
3075 data_buf_size: u64,
3076 rep_index: Option<FullZipRepIndexDetails>,
3077 priority: u64,
3078 rows_in_page: u64,
3079 bits_per_offset: u8,
3080 details: Arc<FullZipDecodeDetails>,
3081 cached_state: Option<Arc<FullZipCacheableState>>,
3083 enable_cache: bool,
3085}
3086
3087impl FullZipScheduler {
3088 fn try_new(
3089 buffer_offsets_and_sizes: &[(u64, u64)],
3090 priority: u64,
3091 rows_in_page: u64,
3092 layout: &pb21::FullZipLayout,
3093 decompressors: &dyn DecompressionStrategy,
3094 ) -> Result<Self> {
3095 let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
3096 let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
3097 let num_reps = rows_in_page + 1;
3098 let bytes_per_rep = len / num_reps;
3099 debug_assert_eq!(len % num_reps, 0);
3100 debug_assert!(
3101 bytes_per_rep == 1
3102 || bytes_per_rep == 2
3103 || bytes_per_rep == 4
3104 || bytes_per_rep == 8
3105 );
3106 FullZipRepIndexDetails {
3107 buf_position: *pos,
3108 bytes_per_value: bytes_per_rep,
3109 }
3110 });
3111
3112 let value_decompressor = match layout.details {
3113 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
3114 let decompressor = decompressors.create_fixed_per_value_decompressor(
3115 layout.value_compression.as_ref().unwrap(),
3116 )?;
3117 PerValueDecompressor::Fixed(decompressor.into())
3118 }
3119 Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
3120 let decompressor = decompressors.create_variable_per_value_decompressor(
3121 layout.value_compression.as_ref().unwrap(),
3122 )?;
3123 PerValueDecompressor::Variable(decompressor.into())
3124 }
3125 None => {
3126 panic!("Full-zip layout must have a `details` field");
3127 }
3128 };
3129 let ctrl_word_parser = ControlWordParser::new(
3130 layout.bits_rep.try_into().unwrap(),
3131 layout.bits_def.try_into().unwrap(),
3132 );
3133 let def_meaning = layout
3134 .layers
3135 .iter()
3136 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3137 .collect::<Vec<_>>();
3138
3139 let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
3140 let max_visible_def = max_visible_level(&def_meaning);
3141
3142 let bits_per_offset = match layout.details {
3143 Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
3144 Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
3145 bits_per_offset as u8
3146 }
3147 None => panic!("Full-zip layout must have a `details` field"),
3148 };
3149
3150 let details = Arc::new(FullZipDecodeDetails {
3151 value_decompressor,
3152 def_meaning: def_meaning.into(),
3153 ctrl_word_parser,
3154 max_rep,
3155 max_visible_def,
3156 });
3157 Ok(Self {
3158 data_buf_position,
3159 data_buf_size,
3160 rep_index,
3161 details,
3162 priority,
3163 rows_in_page,
3164 bits_per_offset,
3165 cached_state: None,
3166 enable_cache: false,
3167 })
3168 }
3169
3170 fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
3171 if ranges.is_empty() {
3172 return false;
3173 }
3174 let mut expected_start = 0;
3175 for range in ranges {
3176 if range.start != expected_start || range.end > rows_in_page || range.end < range.start
3177 {
3178 return false;
3179 }
3180 expected_start = range.end;
3181 }
3182 expected_start == rows_in_page
3183 }
3184
3185 fn create_page_load_task(
3186 io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
3187 num_rows: u64,
3188 details: Arc<FullZipDecodeDetails>,
3189 bits_per_offset: u8,
3190 ) -> PageLoadTask {
3191 let load_task = async move {
3192 let buffers = io_future.await?;
3193 let data = buffers
3194 .into_iter()
3195 .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
3196 .collect::<VecDeque<_>>();
3197 Self::create_decoder(details, data, num_rows, bits_per_offset)
3198 }
3199 .boxed();
3200 PageLoadTask {
3201 decoder_fut: load_task,
3202 num_rows,
3203 }
3204 }
3205
3206 fn create_decoder(
3208 details: Arc<FullZipDecodeDetails>,
3209 data: VecDeque<LanceBuffer>,
3210 num_rows: u64,
3211 bits_per_offset: u8,
3212 ) -> Result<Box<dyn StructuralPageDecoder>> {
3213 match &details.value_decompressor {
3214 PerValueDecompressor::Fixed(decompressor) => {
3215 let bits_per_value = decompressor.bits_per_value();
3216 if bits_per_value % 8 != 0 {
3217 return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
3218 }
3219 let bytes_per_value = bits_per_value / 8;
3220 let total_bytes_per_value =
3221 bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
3222 Ok(Box::new(FixedFullZipDecoder {
3227 details,
3228 data,
3229 num_rows,
3230 offset_in_current: 0,
3231 bytes_per_value: bytes_per_value as usize,
3232 total_bytes_per_value,
3233 }) as Box<dyn StructuralPageDecoder>)
3234 }
3235 PerValueDecompressor::Variable(_decompressor) => {
3236 Ok(Box::new(VariableFullZipDecoder::new(
3237 details,
3238 data,
3239 num_rows,
3240 bits_per_offset,
3241 bits_per_offset,
3242 )?))
3243 }
3244 }
3245 }
3246
3247 fn extract_byte_ranges_from_pairs(
3250 buffer: LanceBuffer,
3251 bytes_per_value: u64,
3252 data_buf_position: u64,
3253 ) -> Vec<Range<u64>> {
3254 ByteUnpacker::new(buffer, bytes_per_value as usize)
3255 .chunks(2)
3256 .into_iter()
3257 .map(|mut c| {
3258 let start = c.next().unwrap() + data_buf_position;
3259 let end = c.next().unwrap() + data_buf_position;
3260 start..end
3261 })
3262 .collect::<Vec<_>>()
3263 }
3264
3265 fn extract_byte_ranges_from_cached(
3268 buffer: &LanceBuffer,
3269 ranges: &[Range<u64>],
3270 bytes_per_value: u64,
3271 data_buf_position: u64,
3272 ) -> Vec<Range<u64>> {
3273 ranges
3274 .iter()
3275 .map(|r| {
3276 let start_offset = (r.start * bytes_per_value) as usize;
3277 let end_offset = (r.end * bytes_per_value) as usize;
3278
3279 let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
3280 let start_val =
3281 ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
3282 .next()
3283 .unwrap();
3284
3285 let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
3286 let end_val =
3287 ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
3288 .next()
3289 .unwrap();
3290
3291 (data_buf_position + start_val)..(data_buf_position + end_val)
3292 })
3293 .collect()
3294 }
3295
3296 fn compute_rep_index_ranges(
3298 ranges: &[Range<u64>],
3299 rep_index: &FullZipRepIndexDetails,
3300 ) -> Vec<Range<u64>> {
3301 ranges
3302 .iter()
3303 .flat_map(|r| {
3304 let first_val_start =
3305 rep_index.buf_position + (r.start * rep_index.bytes_per_value);
3306 let first_val_end = first_val_start + rep_index.bytes_per_value;
3307 let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
3308 let last_val_end = last_val_start + rep_index.bytes_per_value;
3309 [first_val_start..first_val_end, last_val_start..last_val_end]
3310 })
3311 .collect()
3312 }
3313
3314 fn schedule_ranges_rep(
3316 &self,
3317 ranges: &[Range<u64>],
3318 io: &Arc<dyn EncodingsIo>,
3319 rep_index: FullZipRepIndexDetails,
3320 ) -> Result<Vec<PageLoadTask>> {
3321 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3322 let data_buf_position = self.data_buf_position;
3323 let priority = self.priority;
3324 let details = self.details.clone();
3325 let bits_per_offset = self.bits_per_offset;
3326
3327 if Self::covers_entire_page(ranges, self.rows_in_page) {
3328 let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
3329 let page_data = io.submit_single(full_range.clone(), priority);
3330 let load_task = async move {
3331 let page_data = page_data.await?;
3332 let source = FullZipReadSource::PrefetchedPage {
3333 base_offset: full_range.start,
3334 data: LanceBuffer::from_bytes(page_data, 1),
3335 };
3336 let read_ranges = vec![full_range];
3337 let data = source.fetch(&read_ranges, priority).await?;
3338 Self::create_decoder(details, data, num_rows, bits_per_offset)
3339 }
3340 .boxed();
3341 let page_load_task = PageLoadTask {
3342 decoder_fut: load_task,
3343 num_rows,
3344 };
3345 return Ok(vec![page_load_task]);
3346 }
3347
3348 if let Some(cached_state) = &self.cached_state {
3349 let byte_ranges = Self::extract_byte_ranges_from_cached(
3350 &cached_state.rep_index_buffer,
3351 ranges,
3352 rep_index.bytes_per_value,
3353 data_buf_position,
3354 );
3355 let io_future = io.submit_request(byte_ranges, priority);
3356 let page_load_task =
3357 Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
3358 return Ok(vec![page_load_task]);
3359 }
3360
3361 let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
3362 let rep_data = io.submit_request(rep_ranges, priority);
3363 let io_clone = io.clone();
3364 let load_task = async move {
3365 let rep_data = rep_data.await?;
3366 let rep_buffer = LanceBuffer::concat(
3367 &rep_data
3368 .into_iter()
3369 .map(|d| LanceBuffer::from_bytes(d, 1))
3370 .collect::<Vec<_>>(),
3371 );
3372 let byte_ranges = Self::extract_byte_ranges_from_pairs(
3373 rep_buffer,
3374 rep_index.bytes_per_value,
3375 data_buf_position,
3376 );
3377 let source = FullZipReadSource::Remote(io_clone);
3378 let data = source.fetch(&byte_ranges, priority).await?;
3379 Self::create_decoder(details, data, num_rows, bits_per_offset)
3380 }
3381 .boxed();
3382 let page_load_task = PageLoadTask {
3383 decoder_fut: load_task,
3384 num_rows,
3385 };
3386 Ok(vec![page_load_task])
3387 }
3388
3389 fn schedule_ranges_simple(
3393 &self,
3394 ranges: &[Range<u64>],
3395 io: &Arc<dyn EncodingsIo>,
3396 ) -> Result<Vec<PageLoadTask>> {
3397 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3399
3400 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
3401 unreachable!()
3402 };
3403
3404 let bits_per_value = decompressor.bits_per_value();
3406 if !bits_per_value.is_multiple_of(8) {
3407 return Err(Error::invalid_input_source(
3408 format!(
3409 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
3410 bits_per_value
3411 )
3412 .into(),
3413 ));
3414 }
3415 let bytes_per_value = bits_per_value / 8;
3416 let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
3417 let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
3418 let byte_ranges = ranges
3419 .iter()
3420 .map(|r| {
3421 debug_assert!(r.end <= self.rows_in_page);
3422 let start = self.data_buf_position + r.start * total_bytes_per_value;
3423 let end = self.data_buf_position + r.end * total_bytes_per_value;
3424 start..end
3425 })
3426 .collect::<Vec<_>>();
3427
3428 let io_future = io.submit_request(byte_ranges, self.priority);
3429 let page_load_task = Self::create_page_load_task(
3430 io_future,
3431 num_rows,
3432 self.details.clone(),
3433 self.bits_per_offset,
3434 );
3435 Ok(vec![page_load_task])
3436 }
3437}
3438
3439#[derive(Debug)]
3441struct FullZipCacheableState {
3442 rep_index_buffer: LanceBuffer,
3444}
3445
3446impl DeepSizeOf for FullZipCacheableState {
3447 fn deep_size_of_children(&self, _context: &mut Context) -> usize {
3448 self.rep_index_buffer.len()
3449 }
3450}
3451
3452impl CachedPageData for FullZipCacheableState {
3453 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3454 self
3455 }
3456}
3457
3458impl StructuralPageScheduler for FullZipScheduler {
3459 fn needs_initialization(&self) -> bool {
3460 self.enable_cache && self.rep_index.is_some()
3461 }
3462
3463 fn init_layout(&self) -> Result<PageInitialization> {
3464 let repetition_index_range = if self.enable_cache
3465 && let Some(rep_index) = self.rep_index
3466 {
3467 let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
3468 Some(rep_index.buf_position..(rep_index.buf_position + total_size))
3469 } else {
3470 None
3471 };
3472 Ok(PageInitialization::new([repetition_index_range]))
3473 }
3474
3475 fn init_from_buffers<'a>(
3476 &'a mut self,
3477 buffers: PageInitializationBuffers,
3478 _io: &Arc<dyn EncodingsIo>,
3479 ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
3480 let [rep_index_buffer] = match buffers.try_into_array("full-zip") {
3481 Ok(buffers) => buffers,
3482 Err(error) => return std::future::ready(Err(error)).boxed(),
3483 };
3484 if let Some(rep_index_buffer) = rep_index_buffer {
3485 let state = Arc::new(FullZipCacheableState {
3486 rep_index_buffer: LanceBuffer::from_bytes(rep_index_buffer, 1),
3487 });
3488 std::future::ready(Ok(state as Arc<dyn CachedPageData>)).boxed()
3489 } else {
3490 std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
3491 }
3492 }
3493
3494 fn try_load(&mut self, cache: &Arc<dyn CachedPageData>) -> Result<()> {
3498 self.cached_state = Some(cache.clone().as_arc_any().downcast().map_err(|_| {
3499 Error::invalid_input_source("Cached full-zip page data has an unexpected type".into())
3500 })?);
3501 Ok(())
3502 }
3503
3504 fn schedule_ranges(
3505 &self,
3506 ranges: &[Range<u64>],
3507 io: &Arc<dyn EncodingsIo>,
3508 ) -> Result<Vec<PageLoadTask>> {
3509 if let Some(rep_index) = self.rep_index {
3510 self.schedule_ranges_rep(ranges, io, rep_index)
3511 } else {
3512 self.schedule_ranges_simple(ranges, io)
3513 }
3514 }
3515}
3516
3517#[derive(Debug)]
3525struct FixedFullZipDecoder {
3526 details: Arc<FullZipDecodeDetails>,
3527 data: VecDeque<LanceBuffer>,
3528 offset_in_current: usize,
3529 bytes_per_value: usize,
3530 total_bytes_per_value: usize,
3531 num_rows: u64,
3532}
3533
3534impl FixedFullZipDecoder {
3535 fn slice_next_task(&mut self, num_rows: u64) -> Result<FullZipDecodeTaskItem> {
3536 debug_assert!(num_rows > 0);
3537 let cur_buf = self.data.front_mut().ok_or_else(|| {
3538 Error::corrupt_file_named(
3539 "fixed_full_zip",
3540 format!("page data ran out with {num_rows} row(s) left to decode"),
3541 )
3542 })?;
3543 let start = self.offset_in_current;
3544 if self.details.ctrl_word_parser.has_rep() {
3545 let mut rows_started = 0;
3548 let mut num_items = 0;
3551 let bytes_per_word = self.details.ctrl_word_parser.bytes_per_word();
3552 while self.offset_in_current < cur_buf.len() {
3553 let bytes_avail = cur_buf.len() - self.offset_in_current;
3558 if bytes_avail < bytes_per_word {
3559 return Err(Error::corrupt_file_named(
3560 "fixed_full_zip",
3561 format!(
3562 "truncated control word: {bytes_avail} byte(s) remain in the page \
3563 buffer but a control word requires {bytes_per_word}"
3564 ),
3565 ));
3566 }
3567 let control = self.details.ctrl_word_parser.parse_desc(
3568 &cur_buf[self.offset_in_current..],
3569 self.details.max_rep,
3570 self.details.max_visible_def,
3571 );
3572 if control.is_new_row {
3573 if rows_started == num_rows {
3574 break;
3575 }
3576 rows_started += 1;
3577 }
3578 num_items += 1;
3579 let bytes_in_item = if control.is_visible {
3580 self.total_bytes_per_value
3581 } else {
3582 bytes_per_word
3583 };
3584 if bytes_in_item > bytes_avail {
3585 return Err(Error::corrupt_file_named(
3586 "fixed_full_zip",
3587 format!(
3588 "truncated item: {bytes_avail} byte(s) remain in the page buffer at \
3589 offset {} but the item requires {bytes_in_item}",
3590 self.offset_in_current
3591 ),
3592 ));
3593 }
3594 self.offset_in_current += bytes_in_item;
3595 }
3596
3597 let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
3598 if self.offset_in_current == cur_buf.len() {
3599 self.data.pop_front();
3600 self.offset_in_current = 0;
3601 }
3602
3603 Ok(FullZipDecodeTaskItem {
3604 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3605 data: task_slice,
3606 bits_per_value: self.bytes_per_value as u64 * 8,
3607 num_values: num_items,
3608 block_info: BlockInfo::new(),
3609 }),
3610 rows_in_buf: rows_started,
3611 })
3612 } else {
3613 let cur_buf = self.data.front_mut().unwrap();
3616 let bytes_avail = cur_buf.len() - self.offset_in_current;
3617 let offset_in_cur = self.offset_in_current;
3618
3619 let bytes_needed = num_rows as usize * self.total_bytes_per_value;
3620 let mut rows_taken = num_rows;
3621 let task_slice = if bytes_needed >= bytes_avail {
3622 self.offset_in_current = 0;
3623 rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
3624 self.data
3625 .pop_front()
3626 .unwrap()
3627 .slice_with_length(offset_in_cur, bytes_avail)
3628 } else {
3629 self.offset_in_current += bytes_needed;
3630 cur_buf.slice_with_length(offset_in_cur, bytes_needed)
3631 };
3632 Ok(FullZipDecodeTaskItem {
3633 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3634 data: task_slice,
3635 bits_per_value: self.bytes_per_value as u64 * 8,
3636 num_values: rows_taken,
3637 block_info: BlockInfo::new(),
3638 }),
3639 rows_in_buf: rows_taken,
3640 })
3641 }
3642 }
3643}
3644
3645impl StructuralPageDecoder for FixedFullZipDecoder {
3646 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3647 if self.total_bytes_per_value == 0 {
3648 return Ok(Box::new(FixedFullZipDecodeTask {
3651 details: self.details.clone(),
3652 data: vec![FullZipDecodeTaskItem {
3653 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3654 data: LanceBuffer::empty(),
3655 bits_per_value: 0,
3656 num_values: num_rows,
3657 block_info: BlockInfo::new(),
3658 }),
3659 rows_in_buf: num_rows,
3660 }],
3661 bytes_per_value: 0,
3662 num_rows: num_rows as usize,
3663 }));
3664 }
3665 let mut task_data = Vec::with_capacity(self.data.len());
3666 let mut remaining = num_rows;
3667 while remaining > 0 {
3668 let task_item = self.slice_next_task(remaining)?;
3669 remaining -= task_item.rows_in_buf;
3670 task_data.push(task_item);
3671 }
3672 Ok(Box::new(FixedFullZipDecodeTask {
3673 details: self.details.clone(),
3674 data: task_data,
3675 bytes_per_value: self.bytes_per_value,
3676 num_rows: num_rows as usize,
3677 }))
3678 }
3679
3680 fn num_rows(&self) -> u64 {
3681 self.num_rows
3682 }
3683}
3684
3685#[derive(Debug)]
3690struct VariableFullZipDecoder {
3691 details: Arc<FullZipDecodeDetails>,
3692 decompressor: Arc<dyn VariablePerValueDecompressor>,
3693 data: LanceBuffer,
3694 offsets: LanceBuffer,
3695 rep: ScalarBuffer<u16>,
3696 def: ScalarBuffer<u16>,
3697 repdef_starts: Vec<usize>,
3698 data_starts: Vec<usize>,
3699 offset_starts: Vec<usize>,
3700 visible_item_counts: Vec<u64>,
3701 bits_per_offset: u8,
3702 current_idx: usize,
3703 num_rows: u64,
3704}
3705
3706impl VariableFullZipDecoder {
3707 fn new(
3708 details: Arc<FullZipDecodeDetails>,
3709 data: VecDeque<LanceBuffer>,
3710 num_rows: u64,
3711 in_bits_per_length: u8,
3712 out_bits_per_offset: u8,
3713 ) -> Result<Self> {
3714 let decompressor = match details.value_decompressor {
3715 PerValueDecompressor::Variable(ref d) => d.clone(),
3716 _ => unreachable!(),
3717 };
3718
3719 assert_eq!(in_bits_per_length % 8, 0);
3720 assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
3721
3722 let mut decoder = Self {
3723 details,
3724 decompressor,
3725 data: LanceBuffer::empty(),
3726 offsets: LanceBuffer::empty(),
3727 rep: LanceBuffer::empty().borrow_to_typed_slice(),
3728 def: LanceBuffer::empty().borrow_to_typed_slice(),
3729 bits_per_offset: out_bits_per_offset,
3730 repdef_starts: Vec::with_capacity(num_rows as usize + 1),
3731 data_starts: Vec::with_capacity(num_rows as usize + 1),
3732 offset_starts: Vec::with_capacity(num_rows as usize + 1),
3733 visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
3734 current_idx: 0,
3735 num_rows,
3736 };
3737
3738 decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
3759
3760 Ok(decoder)
3761 }
3762
3763 fn slice_batch_data_and_rebase_offsets_typed<T>(
3764 data: &LanceBuffer,
3765 offsets: &LanceBuffer,
3766 ) -> Result<(LanceBuffer, LanceBuffer)>
3767 where
3768 T: arrow_buffer::ArrowNativeType
3769 + Copy
3770 + PartialOrd
3771 + std::ops::Sub<Output = T>
3772 + std::fmt::Display
3773 + TryInto<usize>,
3774 {
3775 let offsets_slice = offsets.borrow_to_typed_slice::<T>();
3776 let offsets_slice = offsets_slice.as_ref();
3777 if offsets_slice.is_empty() {
3778 return Err(Error::internal(
3779 "Variable offsets cannot be empty".to_string(),
3780 ));
3781 }
3782
3783 let base = offsets_slice[0];
3784 let end = *offsets_slice.last().unwrap();
3785 if end < base {
3786 return Err(Error::internal(format!(
3787 "Invalid variable offsets: end ({end}) is less than base ({base})"
3788 )));
3789 }
3790
3791 let data_start = base.try_into().map_err(|_| {
3792 Error::internal(format!("Variable offset ({base}) does not fit into usize"))
3793 })?;
3794 let data_end = end.try_into().map_err(|_| {
3795 Error::internal(format!("Variable offset ({end}) does not fit into usize"))
3796 })?;
3797 if data_end > data.len() {
3798 return Err(Error::internal(format!(
3799 "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
3800 data.len()
3801 )));
3802 }
3803
3804 let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
3805 for &offset in offsets_slice {
3806 if offset < base {
3807 return Err(Error::internal(format!(
3808 "Invalid variable offsets: offset ({offset}) is less than base ({base})"
3809 )));
3810 }
3811 rebased_offsets.push(offset - base);
3812 }
3813
3814 let sliced_data = data.slice_with_length(data_start, data_end - data_start);
3815 let sliced_data = LanceBuffer::copy_slice(&sliced_data);
3817 let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
3818 Ok((sliced_data, rebased_offsets))
3819 }
3820
3821 fn slice_batch_data_and_rebase_offsets(
3822 data: &LanceBuffer,
3823 offsets: &LanceBuffer,
3824 bits_per_offset: u8,
3825 ) -> Result<(LanceBuffer, LanceBuffer)> {
3826 match bits_per_offset {
3827 32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
3828 64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
3829 _ => Err(Error::internal(format!(
3830 "Unsupported bits_per_offset={bits_per_offset}"
3831 ))),
3832 }
3833 }
3834
3835 fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
3842 let width = bits_per_offset as usize / 8;
3843 if data.len() < width {
3844 return Err(Error::corrupt_file_named(
3845 "variable_full_zip",
3846 format!(
3847 "truncated length prefix: {} byte(s) remain in the page buffer but a \
3848 {}-bit length prefix requires {}",
3849 data.len(),
3850 bits_per_offset,
3851 width
3852 ),
3853 ));
3854 }
3855 Ok(match bits_per_offset {
3856 8 => data[0] as u64,
3857 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
3858 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
3859 64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
3860 _ => unreachable!(),
3861 })
3862 }
3863
3864 fn unzip(
3865 &mut self,
3866 data: VecDeque<LanceBuffer>,
3867 in_bits_per_length: u8,
3868 out_bits_per_offset: u8,
3869 num_rows: u64,
3870 ) -> Result<()> {
3871 let mut rep = Vec::with_capacity(num_rows as usize);
3873 let mut def = Vec::with_capacity(num_rows as usize);
3874 let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
3875
3876 let bytes_per_offset = out_bits_per_offset as usize / 8;
3879 let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
3880 let mut offsets_data = Vec::with_capacity(bytes_offsets);
3881
3882 let bytes_per_length = in_bits_per_length as usize / 8;
3883 let bytes_lengths = bytes_per_length * num_rows as usize;
3884
3885 let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
3886 let mut unzipped_data =
3889 Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
3890
3891 let mut current_offset = 0_u64;
3892 let mut visible_item_count = 0_u64;
3893 for databuf in data.into_iter() {
3894 let mut databuf = databuf.as_ref();
3895 while !databuf.is_empty() {
3896 let data_start = unzipped_data.len();
3897 let offset_start = offsets_data.len();
3898 let repdef_start = rep.len().max(def.len());
3901 let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
3903 databuf,
3904 self.details.max_rep,
3905 self.details.max_visible_def,
3906 );
3907 self.details
3908 .ctrl_word_parser
3909 .parse(databuf, &mut rep, &mut def);
3910 databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
3911
3912 if ctrl_desc.is_new_row {
3913 self.repdef_starts.push(repdef_start);
3914 self.data_starts.push(data_start);
3915 self.offset_starts.push(offset_start);
3916 self.visible_item_counts.push(visible_item_count);
3917 }
3918 if ctrl_desc.is_visible {
3919 visible_item_count += 1;
3920 if ctrl_desc.is_valid_item {
3921 let length = Self::parse_length(databuf, in_bits_per_length)?;
3922 match out_bits_per_offset {
3923 32 => offsets_data
3924 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3925 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3926 _ => unreachable!(),
3927 };
3928 databuf = &databuf[bytes_per_offset..];
3929 unzipped_data.extend_from_slice(&databuf[..length as usize]);
3930 databuf = &databuf[length as usize..];
3931 current_offset += length;
3932 } else {
3933 match out_bits_per_offset {
3935 32 => offsets_data
3936 .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3937 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3938 _ => unreachable!(),
3939 }
3940 }
3941 }
3942 }
3943 }
3944 self.repdef_starts.push(rep.len().max(def.len()));
3945 self.data_starts.push(unzipped_data.len());
3946 self.offset_starts.push(offsets_data.len());
3947 self.visible_item_counts.push(visible_item_count);
3948 match out_bits_per_offset {
3949 32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
3950 64 => offsets_data.extend_from_slice(¤t_offset.to_le_bytes()),
3951 _ => unreachable!(),
3952 };
3953 self.rep = ScalarBuffer::from(rep);
3954 self.def = ScalarBuffer::from(def);
3955 self.data = LanceBuffer::from(unzipped_data);
3956 self.offsets = LanceBuffer::from(offsets_data);
3957 Ok(())
3958 }
3959}
3960
3961impl StructuralPageDecoder for VariableFullZipDecoder {
3962 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3963 let start = self.current_idx;
3964 let end = start + num_rows as usize;
3965
3966 let offset_start = self.offset_starts[start];
3967 let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
3968 let offsets = self
3969 .offsets
3970 .slice_with_length(offset_start, offset_end - offset_start);
3971 let (data, offsets) =
3973 Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
3974
3975 let repdef_start = self.repdef_starts[start];
3976 let repdef_end = self.repdef_starts[end];
3977 let rep = if self.rep.is_empty() {
3978 self.rep.clone()
3979 } else {
3980 self.rep.slice(repdef_start, repdef_end - repdef_start)
3981 };
3982 let def = if self.def.is_empty() {
3983 self.def.clone()
3984 } else {
3985 self.def.slice(repdef_start, repdef_end - repdef_start)
3986 };
3987
3988 let visible_item_counts_start = self.visible_item_counts[start];
3989 let visible_item_counts_end = self.visible_item_counts[end];
3990 let num_visible_items = visible_item_counts_end - visible_item_counts_start;
3991
3992 self.current_idx += num_rows as usize;
3993
3994 Ok(Box::new(VariableFullZipDecodeTask {
3995 details: self.details.clone(),
3996 decompressor: self.decompressor.clone(),
3997 data,
3998 offsets,
3999 bits_per_offset: self.bits_per_offset,
4000 num_visible_items,
4001 rep,
4002 def,
4003 }))
4004 }
4005
4006 fn num_rows(&self) -> u64 {
4007 self.num_rows
4008 }
4009}
4010
4011#[derive(Debug)]
4012struct VariableFullZipDecodeTask {
4013 details: Arc<FullZipDecodeDetails>,
4014 decompressor: Arc<dyn VariablePerValueDecompressor>,
4015 data: LanceBuffer,
4016 offsets: LanceBuffer,
4017 bits_per_offset: u8,
4018 num_visible_items: u64,
4019 rep: ScalarBuffer<u16>,
4020 def: ScalarBuffer<u16>,
4021}
4022
4023impl DecodePageTask for VariableFullZipDecodeTask {
4024 fn decode(self: Box<Self>) -> Result<DecodedPage> {
4025 let block = VariableWidthBlock {
4026 data: self.data,
4027 offsets: self.offsets,
4028 bits_per_offset: self.bits_per_offset,
4029 num_values: self.num_visible_items,
4030 block_info: BlockInfo::new(),
4031 };
4032 let decomopressed = self.decompressor.decompress(block)?;
4033 let rep = if self.rep.is_empty() {
4034 None
4035 } else {
4036 Some(self.rep.to_vec())
4037 };
4038 let def = if self.def.is_empty() {
4039 None
4040 } else {
4041 Some(self.def.to_vec())
4042 };
4043 let unraveler = RepDefUnraveler::new(
4044 rep,
4045 def,
4046 self.details.def_meaning.clone(),
4047 self.num_visible_items,
4048 );
4049 Ok(DecodedPage {
4050 data: decomopressed,
4051 repdef: unraveler,
4052 })
4053 }
4054}
4055
4056#[derive(Debug)]
4057struct FullZipDecodeTaskItem {
4058 data: PerValueDataBlock,
4059 rows_in_buf: u64,
4060}
4061
4062#[derive(Debug)]
4065struct FixedFullZipDecodeTask {
4066 details: Arc<FullZipDecodeDetails>,
4067 data: Vec<FullZipDecodeTaskItem>,
4068 num_rows: usize,
4069 bytes_per_value: usize,
4070}
4071
4072impl DecodePageTask for FixedFullZipDecodeTask {
4073 fn decode(self: Box<Self>) -> Result<DecodedPage> {
4074 let estimated_size_bytes = if self.details.ctrl_word_parser.bytes_per_word() == 0 {
4075 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
4076 return Err(Error::internal(
4077 "FixedFullZipDecodeTask requires a fixed-width decompressor",
4078 ));
4079 };
4080 decompressor
4081 .decoded_size_bytes(self.num_rows as u64)
4082 .unwrap_or_else(|| {
4083 self.data
4084 .iter()
4085 .map(|task_item| task_item.data.data_size())
4086 .sum::<u64>()
4087 * 2
4088 })
4089 } else {
4090 self.data
4093 .iter()
4094 .map(|task_item| task_item.data.data_size())
4095 .sum::<u64>()
4096 * 2
4097 };
4098 let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes);
4099
4100 if self.details.ctrl_word_parser.bytes_per_word() == 0 {
4101 for task_item in self.data.into_iter() {
4105 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
4106 unreachable!()
4107 };
4108 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
4109 else {
4110 unreachable!()
4111 };
4112 debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
4113 let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
4114 data_builder.append(&decompressed, 0..task_item.rows_in_buf)?;
4115 }
4116
4117 let unraveler = RepDefUnraveler::new(
4118 None,
4119 None,
4120 self.details.def_meaning.clone(),
4121 self.num_rows as u64,
4122 );
4123
4124 Ok(DecodedPage {
4125 data: data_builder.finish(),
4126 repdef: unraveler,
4127 })
4128 } else {
4129 let mut rep = Vec::with_capacity(self.num_rows);
4131 let mut def = Vec::with_capacity(self.num_rows);
4132
4133 for task_item in self.data.into_iter() {
4134 let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
4135 unreachable!()
4136 };
4137 let mut buf_slice = fixed_data.data.as_ref();
4138 let num_values = fixed_data.num_values as usize;
4139 let mut values = Vec::with_capacity(
4142 fixed_data.data.len()
4143 - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
4144 );
4145 let mut visible_items = 0;
4146 for _ in 0..num_values {
4147 self.details
4149 .ctrl_word_parser
4150 .parse(buf_slice, &mut rep, &mut def);
4151 buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
4152
4153 let is_visible = def
4154 .last()
4155 .map(|d| *d <= self.details.max_visible_def)
4156 .unwrap_or(true);
4157 if is_visible {
4158 values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
4160 buf_slice = &buf_slice[self.bytes_per_value..];
4161 visible_items += 1;
4162 }
4163 }
4164
4165 let values_buf = LanceBuffer::from(values);
4167 let fixed_data = FixedWidthDataBlock {
4168 bits_per_value: self.bytes_per_value as u64 * 8,
4169 block_info: BlockInfo::new(),
4170 data: values_buf,
4171 num_values: visible_items,
4172 };
4173 let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
4174 else {
4175 unreachable!()
4176 };
4177 let decompressed = decompressor.decompress(fixed_data, visible_items)?;
4178 data_builder.append(&decompressed, 0..visible_items)?;
4179 }
4180
4181 let repetition = if rep.is_empty() { None } else { Some(rep) };
4182 let definition = if def.is_empty() { None } else { Some(def) };
4183
4184 let unraveler = RepDefUnraveler::new(
4185 repetition,
4186 definition,
4187 self.details.def_meaning.clone(),
4188 self.num_rows as u64,
4189 );
4190 let data = data_builder.finish();
4191
4192 Ok(DecodedPage {
4193 data,
4194 repdef: unraveler,
4195 })
4196 }
4197 }
4198}
4199
4200#[derive(Debug)]
4201struct StructuralPrimitiveFieldSchedulingJob<'a> {
4202 scheduler: &'a StructuralPrimitiveFieldScheduler,
4203 mappings: std::vec::IntoIter<PageRangeMapping>,
4204}
4205
4206impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
4207 pub fn try_new(
4208 scheduler: &'a StructuralPrimitiveFieldScheduler,
4209 ranges: Vec<Range<u64>>,
4210 ) -> Result<Self> {
4211 let mappings = map_ranges_to_pages(&scheduler.page_row_ends, &ranges)?.into_iter();
4212 Ok(Self {
4213 scheduler,
4214 mappings,
4215 })
4216 }
4217}
4218
4219impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
4220 fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
4221 let Some(first_mapping) = self.mappings.next() else {
4222 return Ok(Vec::new());
4223 };
4224 let batch = Arc::new(RequestBatch::new(context.io().clone()));
4228 let batch_io: Arc<dyn EncodingsIo> = batch.clone();
4229 let cur_path = context.current_path();
4230 let mut scan_lines = Vec::new();
4231 let mut mapping = Some(first_mapping);
4232 while let Some(cur_mapping) = mapping.take() {
4233 let cur_page = self.scheduler.built_page(cur_mapping.page_idx)?;
4234
4235 trace!(
4236 "Scheduling {} rows across {} ranges from page with {} rows (column_index={}, page_index={})",
4237 cur_mapping
4238 .ranges_in_page
4239 .iter()
4240 .map(|r| r.end - r.start)
4241 .sum::<u64>(),
4242 cur_mapping.ranges_in_page.len(),
4243 cur_page.row_range.end - cur_page.row_range.start,
4244 self.scheduler.column_index,
4245 cur_page.page_index,
4246 );
4247
4248 let page_decoders = cur_page
4249 .scheduler
4250 .schedule_ranges(&cur_mapping.ranges_in_page, &batch_io)?;
4251
4252 scan_lines.extend(page_decoders.into_iter().map(|page_load_task| {
4253 let cur_path = cur_path.clone();
4254 let page_decoder = page_load_task.decoder_fut;
4255 let unloaded_page = async move {
4256 let page_decoder = page_decoder.await?;
4257 Ok(LoadedPageShard {
4258 decoder: page_decoder,
4259 path: cur_path,
4260 })
4261 }
4262 .boxed();
4263 ScheduledScanLine {
4264 decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
4265 rows_scheduled: page_load_task.num_rows,
4266 }
4267 }));
4268
4269 if batch.pending_bytes() < IO_REQUEST_BATCH_BYTES {
4270 mapping = self.mappings.next();
4271 }
4272 }
4273 batch.flush();
4274 Ok(scan_lines)
4275 }
4276}
4277
4278#[derive(Debug)]
4279struct PageInfoAndScheduler {
4280 page_index: usize,
4281 row_range: Range<u64>,
4282 scheduler: Box<dyn StructuralPageScheduler>,
4283}
4284
4285#[derive(Debug)]
4292struct PageSchedulerSource {
4293 column_info: Arc<ColumnInfo>,
4294 decompressors: Arc<dyn DecompressionStrategy>,
4295 cache_repetition_index: bool,
4296 target_field: Field,
4297}
4298
4299#[derive(Debug)]
4304pub struct StructuralPrimitiveFieldScheduler {
4305 page_row_ends: Vec<u64>,
4308 page_schedulers: Vec<PageInfoAndScheduler>,
4312 page_source: Option<PageSchedulerSource>,
4315 column_index: u32,
4316 target_data_type_tag: Arc<str>,
4317}
4318
4319impl StructuralPrimitiveFieldScheduler {
4320 #[deprecated(
4325 since = "13.0.0",
4326 note = "builds every page scheduler eagerly; use try_new_lazy"
4327 )]
4328 pub fn try_new(
4329 column_info: &ColumnInfo,
4330 decompressors: &dyn DecompressionStrategy,
4331 cache_repetition_index: bool,
4332 target_field: &Field,
4333 ) -> Result<Self> {
4334 let page_row_ends = Self::page_row_ends(column_info)?;
4335 let mut page_schedulers = Vec::with_capacity(page_row_ends.len());
4336 let mut row_start = 0_u64;
4337 for (page_index, (page_info, &row_end)) in column_info
4338 .page_infos
4339 .iter()
4340 .zip(&page_row_ends)
4341 .enumerate()
4342 {
4343 page_schedulers.push(Self::page_info_to_scheduler(
4344 page_info,
4345 page_index,
4346 row_start..row_end,
4347 decompressors,
4348 cache_repetition_index,
4349 target_field,
4350 )?);
4351 row_start = row_end;
4352 }
4353 Ok(Self {
4354 page_row_ends,
4355 page_schedulers,
4356 page_source: None,
4357 column_index: column_info.index,
4358 target_data_type_tag: format!("{:?}", target_field.data_type()).into(),
4359 })
4360 }
4361
4362 pub fn try_new_lazy(
4369 column_info: Arc<ColumnInfo>,
4370 decompressors: Arc<dyn DecompressionStrategy>,
4371 cache_repetition_index: bool,
4372 target_field: &Field,
4373 ) -> Result<Self> {
4374 Ok(Self {
4375 page_row_ends: Self::page_row_ends(&column_info)?,
4376 page_schedulers: Vec::new(),
4377 column_index: column_info.index,
4378 target_data_type_tag: format!("{:?}", target_field.data_type()).into(),
4379 page_source: Some(PageSchedulerSource {
4380 column_info,
4381 decompressors,
4382 cache_repetition_index,
4383 target_field: target_field.clone(),
4384 }),
4385 })
4386 }
4387
4388 fn page_row_ends(column_info: &ColumnInfo) -> Result<Vec<u64>> {
4390 let mut page_row_ends = Vec::with_capacity(column_info.page_infos.len());
4391 let mut row_start = 0_u64;
4392 for (page_index, page_info) in column_info.page_infos.iter().enumerate() {
4393 let row_end = row_start.checked_add(page_info.num_rows).ok_or_else(|| {
4394 Error::invalid_input_source(
4395 format!(
4396 "Page row range overflowed for column {} at page {page_index}: start={row_start}, rows={}",
4397 column_info.index, page_info.num_rows
4398 )
4399 .into(),
4400 )
4401 })?;
4402 page_row_ends.push(row_end);
4403 row_start = row_end;
4404 }
4405 Ok(page_row_ends)
4406 }
4407
4408 #[cfg(test)]
4410 fn from_page_schedulers(
4411 page_schedulers: Vec<PageInfoAndScheduler>,
4412 column_index: u32,
4413 target_data_type_tag: Arc<str>,
4414 ) -> Self {
4415 Self {
4416 page_row_ends: page_schedulers
4417 .iter()
4418 .map(|page| page.row_range.end)
4419 .collect(),
4420 page_schedulers,
4421 page_source: None,
4422 column_index,
4423 target_data_type_tag,
4424 }
4425 }
4426
4427 fn page_row_range(&self, page_index: usize) -> Range<u64> {
4428 let start = match page_index.checked_sub(1) {
4429 Some(previous) => self.page_row_ends[previous],
4430 None => 0,
4431 };
4432 start..self.page_row_ends[page_index]
4433 }
4434
4435 fn built_page_position(&self, page_index: usize) -> Option<usize> {
4437 self.page_schedulers
4438 .binary_search_by_key(&page_index, |page| page.page_index)
4439 .ok()
4440 }
4441
4442 fn built_page(&self, page_index: usize) -> Result<&PageInfoAndScheduler> {
4443 let position = self.built_page_position(page_index).ok_or_else(|| {
4444 Error::internal(format!(
4445 "Page {page_index} of column {} was scheduled but not initialized; the scheduled ranges must be covered by the ranges passed to initialize",
4446 self.column_index
4447 ))
4448 })?;
4449 Ok(&self.page_schedulers[position])
4450 }
4451
4452 fn build_pages(&mut self, page_indices: &[usize]) -> Result<Vec<usize>> {
4455 if let Some(source) = &self.page_source {
4456 let mut missing = Vec::new();
4457 for &page_index in page_indices {
4458 if self.built_page_position(page_index).is_some() {
4459 continue;
4460 }
4461 let page_info = source
4462 .column_info
4463 .page_infos
4464 .get(page_index)
4465 .ok_or_else(|| {
4466 Error::internal(format!(
4467 "Page {page_index} is out of bounds for column {} with {} pages",
4468 self.column_index,
4469 source.column_info.page_infos.len()
4470 ))
4471 })?;
4472 missing.push(Self::page_info_to_scheduler(
4473 page_info,
4474 page_index,
4475 self.page_row_range(page_index),
4476 source.decompressors.as_ref(),
4477 source.cache_repetition_index,
4478 &source.target_field,
4479 )?);
4480 }
4481 if !missing.is_empty() {
4482 self.page_schedulers = std::mem::take(&mut self.page_schedulers)
4483 .into_iter()
4484 .merge_by(missing, |built, new| built.page_index < new.page_index)
4485 .collect();
4486 }
4487 if self.page_schedulers.len() == self.page_row_ends.len() {
4488 self.page_source = None;
4489 }
4490 }
4491 page_indices
4492 .iter()
4493 .map(|&page_index| {
4494 self.built_page_position(page_index).ok_or_else(|| {
4495 Error::internal(format!(
4496 "Page {page_index} of column {} is missing after construction",
4497 self.column_index
4498 ))
4499 })
4500 })
4501 .collect()
4502 }
4503
4504 fn page_layout_to_scheduler(
4505 page_info: &PageInfo,
4506 page_layout: &PageLayout,
4507 decompressors: &dyn DecompressionStrategy,
4508 cache_repetition_index: bool,
4509 target_field: &Field,
4510 ) -> Result<Box<dyn StructuralPageScheduler>> {
4511 use pb21::page_layout::Layout;
4512 Ok(match page_layout.layout.as_ref().expect_ok()? {
4513 Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
4514 &page_info.buffer_offsets_and_sizes,
4515 page_info.priority,
4516 mini_block.num_items,
4517 mini_block,
4518 decompressors,
4519 )?),
4520 Layout::SparseLayout(sparse_layout) => {
4521 Box::new(sparse::SparseStructuralScheduler::try_new(
4522 &page_info.buffer_offsets_and_sizes,
4523 page_info.priority,
4524 page_info.num_rows,
4525 target_field.data_type(),
4526 sparse_layout,
4527 decompressors,
4528 )?)
4529 }
4530 Layout::FullZipLayout(full_zip) => {
4531 let mut scheduler = FullZipScheduler::try_new(
4532 &page_info.buffer_offsets_and_sizes,
4533 page_info.priority,
4534 page_info.num_rows,
4535 full_zip,
4536 decompressors,
4537 )?;
4538 scheduler.enable_cache = cache_repetition_index;
4539 Box::new(scheduler)
4540 }
4541 Layout::ConstantLayout(constant_layout) => {
4542 let def_meaning = constant_layout
4543 .layers
4544 .iter()
4545 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4546 .collect::<Vec<_>>();
4547 let has_scalar_value = constant_layout.inline_value.is_some()
4548 || page_info.buffer_offsets_and_sizes.len() == 1
4549 || page_info.buffer_offsets_and_sizes.len() == 3;
4550 if has_scalar_value {
4551 Box::new(constant::ConstantPageScheduler::try_new(
4552 page_info.buffer_offsets_and_sizes.clone(),
4553 constant_layout.inline_value.clone(),
4554 target_field.data_type(),
4555 def_meaning.into(),
4556 )?) as Box<dyn StructuralPageScheduler>
4557 } else if def_meaning.len() == 1
4558 && def_meaning[0] == DefinitionInterpretation::NullableItem
4559 {
4560 Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
4561 } else {
4562 let rep_codec = LevelCodec::try_new(
4565 constant_layout.rep_compression.as_ref(),
4566 decompressors,
4567 )?;
4568 let def_codec = LevelCodec::try_new(
4569 constant_layout.def_compression.as_ref(),
4570 decompressors,
4571 )?;
4572
4573 Box::new(ComplexAllNullScheduler::new(
4574 page_info.buffer_offsets_and_sizes.clone(),
4575 def_meaning.into(),
4576 rep_codec,
4577 def_codec,
4578 constant_layout.num_rep_values,
4579 constant_layout.num_def_values,
4580 )) as Box<dyn StructuralPageScheduler>
4581 }
4582 }
4583 Layout::BlobLayout(blob) => {
4584 let inner_scheduler = Self::page_layout_to_scheduler(
4585 page_info,
4586 blob.inner_layout.as_ref().expect_ok()?.as_ref(),
4587 decompressors,
4588 cache_repetition_index,
4589 target_field,
4590 )?;
4591 let def_meaning = blob
4592 .layers
4593 .iter()
4594 .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4595 .collect::<Vec<_>>();
4596 if matches!(target_field.data_type(), DataType::Struct(_)) {
4597 Box::new(BlobDescriptionPageScheduler::new(
4599 inner_scheduler,
4600 def_meaning.into(),
4601 ))
4602 } else {
4603 Box::new(BlobPageScheduler::new(
4605 inner_scheduler,
4606 page_info.priority,
4607 page_info.num_rows,
4608 def_meaning.into(),
4609 ))
4610 }
4611 }
4612 })
4613 }
4614
4615 fn page_info_to_scheduler(
4616 page_info: &PageInfo,
4617 page_index: usize,
4618 row_range: Range<u64>,
4619 decompressors: &dyn DecompressionStrategy,
4620 cache_repetition_index: bool,
4621 target_field: &Field,
4622 ) -> Result<PageInfoAndScheduler> {
4623 let page_layout = page_info.encoding.as_structural();
4624 let scheduler = Self::page_layout_to_scheduler(
4625 page_info,
4626 page_layout,
4627 decompressors,
4628 cache_repetition_index,
4629 target_field,
4630 )?;
4631 Ok(PageInfoAndScheduler {
4632 page_index,
4633 row_range,
4634 scheduler,
4635 })
4636 }
4637}
4638
4639pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
4640 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
4641}
4642
4643pub struct NoCachedPageData;
4644
4645impl DeepSizeOf for NoCachedPageData {
4646 fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
4647 0
4648 }
4649}
4650impl CachedPageData for NoCachedPageData {
4651 fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
4652 self
4653 }
4654}
4655
4656#[doc = "Deprecated compatibility wrapper; new code uses per-page PageDataCacheKey entries."]
4657pub struct CachedFieldData {
4658 pages: Vec<Arc<dyn CachedPageData>>,
4659}
4660
4661impl DeepSizeOf for CachedFieldData {
4662 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
4663 self.pages.deep_size_of_children(ctx)
4664 }
4665}
4666
4667#[derive(Debug, Clone)]
4669pub struct FieldDataCacheKey {
4670 pub column_index: u32,
4671 pub view_tag: String,
4672}
4673
4674impl CacheKey for FieldDataCacheKey {
4675 type ValueType = CachedFieldData;
4676
4677 fn key(&self) -> std::borrow::Cow<'_, str> {
4678 format!("{}:{}", self.column_index, self.view_tag).into()
4679 }
4680
4681 fn type_name() -> &'static str {
4682 "FieldData"
4683 }
4684
4685 fn schema() -> CacheKeySchema {
4686 CacheKeySchema::new("lance.encoding.logical.primitive.field-data-key", 1)
4687 }
4688
4689 fn write_key(&self, builder: &mut KeyBuilder) {
4690 builder.write_u32(self.column_index);
4691 builder.write_str(&self.view_tag);
4692 }
4693}
4694
4695pub struct CachedPage {
4699 data: Arc<dyn CachedPageData>,
4700}
4701
4702impl DeepSizeOf for CachedPage {
4703 fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
4704 self.data.deep_size_of_children(ctx)
4705 }
4706}
4707
4708#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4718#[non_exhaustive]
4719pub enum PageCacheView {
4720 Native,
4722 BlobDescriptor,
4724 BlobPayload,
4726}
4727
4728impl PageCacheView {
4729 fn cache_tag(self) -> &'static str {
4730 match self {
4731 Self::Native => "n",
4732 Self::BlobDescriptor => "bd",
4733 Self::BlobPayload => "bp",
4734 }
4735 }
4736}
4737
4738#[derive(Debug, Clone)]
4739pub struct PageDataCacheKey {
4740 pub column_index: u32,
4741 pub page_index: u32,
4742 pub view: PageCacheView,
4743 pub target_data_type_tag: Arc<str>,
4744}
4745
4746impl CacheKey for PageDataCacheKey {
4747 type ValueType = CachedPage;
4748
4749 fn key(&self) -> std::borrow::Cow<'_, str> {
4750 format!(
4751 "v2/{}/{}/{}/{}",
4752 self.column_index,
4753 self.page_index,
4754 self.view.cache_tag(),
4755 self.target_data_type_tag
4756 )
4757 .into()
4758 }
4759
4760 fn type_name() -> &'static str {
4761 "PageData"
4762 }
4763
4764 fn schema() -> CacheKeySchema {
4765 CacheKeySchema::new("lance.encoding.logical.primitive.page-data-key", 2)
4766 }
4767
4768 fn write_key(&self, builder: &mut KeyBuilder) {
4769 builder.write_u32(self.column_index);
4770 builder.write_u32(self.page_index);
4771 builder.write_variant(match self.view {
4772 PageCacheView::Native => 0,
4773 PageCacheView::BlobDescriptor => 1,
4774 PageCacheView::BlobPayload => 2,
4775 });
4776 builder.write_str(&self.target_data_type_tag);
4777 }
4778}
4779
4780#[derive(Debug)]
4784struct PageRangeMapping {
4785 page_idx: usize,
4786 ranges_in_page: Vec<Range<u64>>,
4787}
4788
4789fn for_each_page_range(
4797 page_row_ends: &[u64],
4798 ranges: &[Range<u64>],
4799 mut visit: impl FnMut(usize, Range<u64>),
4800) -> Result<()> {
4801 for (range_idx, range) in ranges.iter().enumerate() {
4802 if range.start > range.end {
4803 return Err(Error::invalid_input_source(
4804 format!(
4805 "Requested row range {range_idx} is reversed: {}..{}",
4806 range.start, range.end
4807 )
4808 .into(),
4809 ));
4810 }
4811 if range_idx > 0 {
4812 let previous = &ranges[range_idx - 1];
4813 if previous.start >= range.start || previous.end > range.start {
4814 return Err(Error::invalid_input_source(
4815 format!(
4816 "Requested row ranges overlap or are not strictly ordered at index {range_idx}: {previous:?}, then {range:?}"
4817 )
4818 .into(),
4819 ));
4820 }
4821 }
4822 }
4823
4824 for range in ranges {
4825 let mut page_idx = page_row_ends.partition_point(|page_end| *page_end <= range.start);
4826 let mut page_start = match page_idx.checked_sub(1) {
4827 Some(previous) => page_row_ends[previous],
4828 None => 0,
4829 };
4830 while let Some(&page_end) = page_row_ends.get(page_idx) {
4831 if page_start >= range.end {
4832 break;
4833 }
4834 let start_in_page = range.start.max(page_start) - page_start;
4835 let end_in_page = range.end.min(page_end) - page_start;
4836 if start_in_page < end_in_page {
4837 visit(page_idx, start_in_page..end_in_page);
4838 }
4839 page_idx += 1;
4840 page_start = page_end;
4841 }
4842 }
4843 Ok(())
4844}
4845
4846fn map_ranges_to_pages(
4852 page_row_ends: &[u64],
4853 ranges: &[Range<u64>],
4854) -> Result<Vec<PageRangeMapping>> {
4855 let mut result: Vec<PageRangeMapping> =
4856 Vec::with_capacity(ranges.len().min(page_row_ends.len()));
4857 for_each_page_range(page_row_ends, ranges, |page_idx, range_in_page| {
4858 if let Some(mapping) = result.last_mut()
4859 && mapping.page_idx == page_idx
4860 {
4861 mapping.ranges_in_page.push(range_in_page);
4862 } else {
4863 result.push(PageRangeMapping {
4864 page_idx,
4865 ranges_in_page: vec![range_in_page],
4866 });
4867 }
4868 })?;
4869 Ok(result)
4870}
4871
4872fn pages_overlapping_ranges(page_row_ends: &[u64], ranges: &[Range<u64>]) -> Result<Vec<usize>> {
4873 let mut result = Vec::with_capacity(ranges.len().min(page_row_ends.len()));
4874 for_each_page_range(page_row_ends, ranges, |page_idx, _| {
4875 if result.last().copied() != Some(page_idx) {
4876 result.push(page_idx);
4877 }
4878 })?;
4879 Ok(result)
4880}
4881
4882#[derive(Debug)]
4883struct PageCacheMiss {
4884 page_idx: usize,
4887 cache_key: PageDataCacheKey,
4888 initialization: PageInitialization,
4889}
4890
4891async fn read_page_initialization_buffers(
4892 io: &Arc<dyn EncodingsIo>,
4893 all_ranges: Vec<Range<u64>>,
4894) -> Result<Vec<Bytes>> {
4895 let mut sorted_ranges = all_ranges.into_iter().enumerate().collect::<Vec<_>>();
4900 sorted_ranges.sort_unstable_by_key(|(_, range)| (range.start, range.end));
4901 let sorted_requests = sorted_ranges
4902 .iter()
4903 .map(|(_, range)| range.clone())
4904 .collect::<Vec<_>>();
4905 let expected_buffer_count = sorted_requests.len();
4906 let sorted_buffers = if sorted_requests.is_empty() {
4907 Vec::new()
4908 } else {
4909 io.submit_request(sorted_requests, 0).await?
4910 };
4911 if sorted_buffers.len() != expected_buffer_count {
4912 return Err(Error::invalid_input_source(
4913 format!(
4914 "Page initialization I/O returned {} buffers for {} ranges",
4915 sorted_buffers.len(),
4916 expected_buffer_count
4917 )
4918 .into(),
4919 ));
4920 }
4921 let mut buffers = vec![None; expected_buffer_count];
4922 for ((original_idx, _), buffer) in sorted_ranges.into_iter().zip(sorted_buffers) {
4923 buffers[original_idx] = Some(buffer);
4924 }
4925 buffers
4926 .into_iter()
4927 .map(|buffer| {
4928 buffer.ok_or_else(|| Error::internal("Page initialization buffer was not assigned"))
4929 })
4930 .collect()
4931}
4932
4933async fn cache_initialized_page(
4934 cache: Arc<LanceCache>,
4935 cache_key: PageDataCacheKey,
4936 data: Arc<dyn CachedPageData>,
4937) {
4938 if data
4939 .clone()
4940 .as_arc_any()
4941 .downcast::<NoCachedPageData>()
4942 .is_err()
4943 {
4944 cache
4945 .insert_with_key(&cache_key, Arc::new(CachedPage { data }))
4946 .await;
4947 }
4948}
4949
4950impl StructuralPrimitiveFieldScheduler {
4951 async fn find_cache_misses(
4955 &mut self,
4956 positions: impl ExactSizeIterator<Item = usize>,
4957 cache: &Arc<LanceCache>,
4958 ) -> Result<Vec<PageCacheMiss>> {
4959 let mut misses = Vec::with_capacity(positions.len());
4961 for page_idx in positions {
4962 let page = &self.page_schedulers[page_idx];
4963 if !page.scheduler.needs_initialization() {
4964 continue;
4965 }
4966 let page_index = u32::try_from(page.page_index).map_err(|_| {
4967 Error::invalid_input_source(
4968 format!("Page index does not fit u32: {}", page.page_index).into(),
4969 )
4970 })?;
4971 let cache_key = PageDataCacheKey {
4972 column_index: self.column_index,
4973 page_index,
4974 view: page.scheduler.cache_view(),
4975 target_data_type_tag: self.target_data_type_tag.clone(),
4976 };
4977 if let Some(cached_page) = cache.get_with_key(&cache_key).await {
4978 match self.page_schedulers[page_idx]
4979 .scheduler
4980 .try_load(&cached_page.data)
4981 {
4982 Ok(()) => continue,
4983 Err(error) => {
4984 debug!("Replacing invalid cached page data for {cache_key:?}: {error}");
4985 }
4986 }
4987 }
4988 let initialization = self.page_schedulers[page_idx].scheduler.init_layout()?;
4989 misses.push(PageCacheMiss {
4990 page_idx,
4991 cache_key,
4992 initialization,
4993 });
4994 }
4995 Ok(misses)
4996 }
4997
4998 fn initialization_ranges(misses: &[PageCacheMiss]) -> Vec<Range<u64>> {
4999 let num_ranges = misses
5000 .iter()
5001 .map(|miss| miss.initialization.num_required_ranges())
5002 .sum();
5003 let mut all_ranges = Vec::with_capacity(num_ranges);
5004 for miss in misses {
5005 all_ranges.extend(miss.initialization.required_ranges().cloned());
5006 }
5007 all_ranges
5008 }
5009
5010 async fn initialize_cache_misses(
5011 &mut self,
5012 misses: Vec<PageCacheMiss>,
5013 buffers: Vec<Bytes>,
5014 cache: Arc<LanceCache>,
5015 io: &Arc<dyn EncodingsIo>,
5016 ) -> Result<()> {
5017 let mut buffers = buffers.into_iter();
5018
5019 let mut tail = self.page_schedulers.as_mut_slice();
5025 let mut next_page = 0;
5026 let mut page_inits = FuturesOrdered::new();
5027 for miss in misses {
5028 let PageCacheMiss {
5029 page_idx,
5030 cache_key,
5031 initialization,
5032 } = miss;
5033 let relative_idx = page_idx
5034 .checked_sub(next_page)
5035 .ok_or_else(|| Error::internal("Page initialization indices are not ordered"))?;
5036 if relative_idx >= tail.len() {
5037 return Err(Error::internal(format!(
5038 "Page initialization index {page_idx} is out of bounds"
5039 )));
5040 }
5041 let (_, rest) = tail.split_at_mut(relative_idx);
5042 let (page, rest) = rest.split_first_mut().ok_or_else(|| {
5043 Error::internal(format!(
5044 "Page initialization index {page_idx} is out of bounds"
5045 ))
5046 })?;
5047 tail = rest;
5048 next_page = page_idx + 1;
5049
5050 let count = initialization.num_required_ranges();
5051 let page_buffers = buffers.by_ref().take(count).collect::<Vec<_>>();
5052 if page_buffers.len() != count {
5053 return Err(Error::invalid_input_source(
5054 format!(
5055 "Page {page_idx} received {} initialization buffers, expected {count}",
5056 page_buffers.len()
5057 )
5058 .into(),
5059 ));
5060 }
5061 let page_buffers = initialization.with_buffers(page_buffers)?;
5062 let cache = cache.clone();
5063 let scheduler = &mut page.scheduler;
5064 page_inits.push_back(async move {
5065 let cached = scheduler.init_from_buffers(page_buffers, io).await?;
5066 scheduler.try_load(&cached)?;
5067 cache_initialized_page(cache, cache_key, cached).await;
5068 Ok::<(), Error>(())
5069 });
5070 }
5071 if buffers.next().is_some() {
5072 return Err(Error::invalid_input_source(
5073 "Page initialization returned unassigned buffers".into(),
5074 ));
5075 }
5076
5077 page_inits.try_collect::<Vec<()>>().await?;
5078 Ok(())
5079 }
5080
5081 async fn initialize_pages(
5082 &mut self,
5083 requested_ranges: Option<&[Range<u64>]>,
5084 context: &SchedulerContext,
5085 ) -> Result<()> {
5086 let cache = context.cache().clone();
5087 let page_indices = match requested_ranges {
5088 None => (0..self.page_row_ends.len()).collect::<Vec<_>>(),
5089 Some(ranges) => pages_overlapping_ranges(&self.page_row_ends, ranges)?,
5090 };
5091 let positions = self.build_pages(&page_indices)?;
5092 let misses = self
5093 .find_cache_misses(positions.into_iter(), &cache)
5094 .await?;
5095
5096 if misses.is_empty() {
5097 return Ok(());
5098 }
5099
5100 let all_ranges = Self::initialization_ranges(&misses);
5103 let buffers = read_page_initialization_buffers(context.io(), all_ranges).await?;
5104 self.initialize_cache_misses(misses, buffers, cache, context.io())
5105 .await
5106 }
5107}
5108
5109impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
5110 fn initialize<'a>(
5111 &'a mut self,
5112 requested_ranges: Option<&'a [Range<u64>]>,
5113 _filter: &'a FilterExpression,
5114 context: &'a SchedulerContext,
5115 ) -> BoxFuture<'a, Result<()>> {
5116 self.initialize_pages(requested_ranges, context).boxed()
5117 }
5118
5119 fn schedule_ranges<'a>(
5120 &'a self,
5121 ranges: &[Range<u64>],
5122 _filter: &FilterExpression,
5123 ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
5124 let ranges = ranges.to_vec();
5125 Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::try_new(
5126 self, ranges,
5127 )?))
5128 }
5129}
5130
5131#[derive(Debug)]
5134pub struct StructuralCompositeDecodeArrayTask {
5135 tasks: Vec<Box<dyn DecodePageTask>>,
5136 should_validate: bool,
5137 data_type: DataType,
5138}
5139
5140impl StructuralCompositeDecodeArrayTask {
5141 fn restore_validity(
5142 array: Arc<dyn Array>,
5143 unraveler: &mut CompositeRepDefUnraveler,
5144 ) -> Result<Arc<dyn Array>> {
5145 let validity = unraveler.unravel_validity(array.len())?;
5146 let Some(validity) = validity else {
5147 return Ok(array);
5148 };
5149 if array.data_type() == &DataType::Null {
5150 return Ok(array);
5152 }
5153 if validity.len() != array.len() {
5154 return Err(Error::invalid_input_source(
5155 format!(
5156 "Structural validity has {} entries for an array with {} values",
5157 validity.len(),
5158 array.len()
5159 )
5160 .into(),
5161 ));
5162 }
5163 Ok(make_array(unsafe {
5166 array
5167 .to_data()
5168 .into_builder()
5169 .nulls(Some(validity))
5170 .build_unchecked()
5171 }))
5172 }
5173}
5174
5175impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
5176 fn decode(self: Box<Self>) -> Result<DecodedArray> {
5177 let mut arrays = Vec::with_capacity(self.tasks.len());
5178 let mut unravelers = Vec::with_capacity(self.tasks.len());
5179 let mut data_size = 0u64;
5180 for task in self.tasks {
5181 let decoded = task.decode()?;
5182 data_size += decoded.data.data_size();
5183 unravelers.push(decoded.repdef);
5184
5185 let array = make_array(
5186 decoded
5187 .data
5188 .into_arrow(self.data_type.clone(), self.should_validate)?,
5189 );
5190
5191 arrays.push(array);
5192 }
5193 let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
5194 let array = arrow_select::concat::concat(&array_refs)?;
5195 let mut repdef = CompositeRepDefUnraveler::new(unravelers);
5196
5197 let array = Self::restore_validity(array, &mut repdef)?;
5198
5199 Ok(DecodedArray {
5200 array,
5201 repdef,
5202 data_size,
5203 })
5204 }
5205}
5206
5207fn fsl_child_bitmap_bytes(data_type: &arrow_schema::DataType, rows: u64) -> u64 {
5213 match data_type {
5214 arrow_schema::DataType::FixedSizeList(child_field, dimension) => {
5215 let child_rows = rows * *dimension as u64;
5216 let own = if child_field.is_nullable() {
5217 child_rows.div_ceil(8)
5218 } else {
5219 0
5220 };
5221 own + fsl_child_bitmap_bytes(child_field.data_type(), child_rows)
5222 }
5223 _ => 0,
5224 }
5225}
5226
5227#[derive(Debug)]
5228pub struct StructuralPrimitiveFieldDecoder {
5229 field: Arc<ArrowField>,
5230 page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
5231 should_validate: bool,
5232 rows_drained_in_current: u64,
5233}
5234
5235impl StructuralPrimitiveFieldDecoder {
5236 pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
5237 Self {
5238 field: field.clone(),
5239 page_decoders: VecDeque::new(),
5240 should_validate,
5241 rows_drained_in_current: 0,
5242 }
5243 }
5244
5245 fn decoded_bytes_for_rows(&self, rows: u64) -> lance_core::Result<u64> {
5246 let mut remaining = rows;
5247 let mut total = 0u64;
5248 for (page_num, page) in self.page_decoders.iter().enumerate() {
5249 if remaining == 0 {
5250 break;
5251 }
5252 let available = if page_num == 0 {
5253 page.num_rows().saturating_sub(self.rows_drained_in_current)
5254 } else {
5255 page.num_rows()
5256 };
5257 let take = available.min(remaining);
5258 total += page.decoded_bytes(take)?;
5259 remaining -= take;
5260 }
5261 if remaining > 0 {
5262 return Err(lance_core::Error::not_supported(format!(
5263 "plan_decoded_bytes: queued pages cover only {} of {} requested rows",
5264 rows - remaining,
5265 rows,
5266 )));
5267 }
5268 Ok(total)
5269 }
5270}
5271
5272impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
5273 fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
5274 assert!(child.path.is_empty());
5275 self.page_decoders.push_back(child.decoder);
5276 Ok(())
5277 }
5278
5279 fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
5280 let mut remaining = num_rows;
5281 let mut tasks = Vec::new();
5282 while remaining > 0 {
5283 let queued_pages = self.page_decoders.len();
5284 let Some(cur_page) = self.page_decoders.front_mut() else {
5285 return Err(Error::internal(format!(
5286 "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})",
5287 self.field.name(),
5288 self.field.data_type(),
5289 num_rows,
5290 remaining,
5291 self.rows_drained_in_current,
5292 queued_pages
5293 )));
5294 };
5295 let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
5296 let to_take = num_in_page.min(remaining);
5297
5298 let task = cur_page.drain(to_take)?;
5299 tasks.push(task);
5300
5301 if to_take == num_in_page {
5302 self.page_decoders.pop_front();
5303 self.rows_drained_in_current = 0;
5304 } else {
5305 self.rows_drained_in_current += to_take;
5306 }
5307
5308 remaining -= to_take;
5309 }
5310 Ok(Box::new(StructuralCompositeDecodeArrayTask {
5311 tasks,
5312 should_validate: self.should_validate,
5313 data_type: self.field.data_type().clone(),
5314 }))
5315 }
5316
5317 fn data_type(&self) -> &DataType {
5318 self.field.data_type()
5319 }
5320
5321 fn plan_decoded_bytes(&self, rows_remaining: u64) -> lance_core::Result<[u64; 8]> {
5322 use crate::decoder::CANDIDATE_BATCH_SIZES;
5323 let data_type = self.field.data_type();
5324 let is_nullable = self.field.is_nullable();
5325
5326 if let Some(byte_width) = data_type.byte_width_opt() {
5328 let mut out = [0u64; 8];
5329 for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5330 let rows = (c as u64).min(rows_remaining);
5331 out[i] = rows * byte_width as u64;
5332 if is_nullable {
5333 out[i] += rows.div_ceil(8);
5334 }
5335 out[i] += fsl_child_bitmap_bytes(data_type, rows);
5336 }
5337 return Ok(out);
5338 }
5339
5340 if matches!(data_type, DataType::Boolean) {
5342 let mut out = [0u64; 8];
5343 for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5344 let rows = (c as u64).min(rows_remaining);
5345 out[i] = rows.div_ceil(8);
5346 if is_nullable {
5347 out[i] += rows.div_ceil(8);
5348 }
5349 }
5350 return Ok(out);
5351 }
5352
5353 if matches!(data_type, DataType::Null) {
5355 return Ok([0u64; 8]);
5356 }
5357
5358 let mut out = [0u64; 8];
5360 for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5361 let rows = (c as u64).min(rows_remaining);
5362 out[i] = self.decoded_bytes_for_rows(rows)?;
5363 if is_nullable {
5364 out[i] += rows.div_ceil(8);
5365 }
5366 }
5367 Ok(out)
5368 }
5369}
5370
5371struct SerializedFullZip {
5373 values: LanceBuffer,
5375 repetition_index: Option<LanceBuffer>,
5377}
5378
5379const MINIBLOCK_ALIGNMENT: usize = 8;
5394
5395#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5422pub(crate) enum MiniblockChunkSize {
5423 U16,
5424 U32,
5425}
5426
5427#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5428enum ComplexNullEncoding {
5429 RawLevels,
5430 CompressedLevels,
5431}
5432
5433#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5434enum FixedWidthDictionaryEncoding {
5435 Exclude64Bit,
5436 Include64Bit,
5437}
5438
5439trait PrimitivePageEncodingBehavior: Send + Sync + Debug {
5440 fn validate_field(&self, _field: &Field, _metadata: &HashMap<String, String>) -> Result<()> {
5441 Ok(())
5442 }
5443
5444 fn try_plan_pages(
5445 &self,
5446 _ctx: &PrimitivePlanContext<'_>,
5447 _arrays: &[ArrayRef],
5448 _normalized: &NormalizedStructuralPlan,
5449 _row_number: u64,
5450 _num_rows: u64,
5451 _num_values: u64,
5452 ) -> Result<Option<Vec<PrimitivePageData>>> {
5453 Ok(None)
5454 }
5455
5456 fn try_encode_page(
5457 &self,
5458 _ctx: &PrimitiveEncodeContext,
5459 page: PrimitivePageData,
5460 ) -> Result<PrimitiveEncodeAttempt> {
5461 Ok(PrimitiveEncodeAttempt::Unhandled(page))
5462 }
5463}
5464
5465#[derive(Debug, Clone)]
5468pub struct PrimitivePageEncoding {
5469 behavior: Arc<dyn PrimitivePageEncodingBehavior>,
5470}
5471
5472impl PrimitivePageEncoding {
5473 pub fn reject_sparse() -> Self {
5475 Self {
5476 behavior: Arc::new(RejectSparsePrimitiveEncoding),
5477 }
5478 }
5479
5480 pub fn constant() -> Self {
5482 Self {
5483 behavior: Arc::new(ConstantPrimitiveEncoding),
5484 }
5485 }
5486
5487 pub fn sparse(compression: Arc<dyn CompressionStrategy>) -> Self {
5489 Self {
5490 behavior: Arc::new(SparsePrimitiveEncoding { compression }),
5491 }
5492 }
5493
5494 pub fn dense_u16(compression: Arc<dyn CompressionStrategy>) -> Self {
5496 Self {
5497 behavior: Arc::new(DenseU16PrimitiveEncoding { compression }),
5498 }
5499 }
5500
5501 pub fn dense_u32(compression: Arc<dyn CompressionStrategy>) -> Self {
5503 Self {
5504 behavior: Arc::new(DenseU32PrimitiveEncoding { compression }),
5505 }
5506 }
5507}
5508
5509#[derive(Debug)]
5510struct RejectSparsePrimitiveEncoding;
5511
5512#[derive(Debug)]
5513struct ConstantPrimitiveEncoding;
5514
5515#[derive(Debug)]
5516struct SparsePrimitiveEncoding {
5517 compression: Arc<dyn CompressionStrategy>,
5518}
5519
5520#[derive(Debug)]
5521struct DenseU16PrimitiveEncoding {
5522 compression: Arc<dyn CompressionStrategy>,
5523}
5524
5525#[derive(Debug)]
5526struct DenseU32PrimitiveEncoding {
5527 compression: Arc<dyn CompressionStrategy>,
5528}
5529
5530pub struct PrimitiveStructuralEncoder {
5531 accumulation_queue: AccumulationQueue,
5533
5534 keep_original_array: bool,
5535 accumulated_repdefs: Vec<RepDefBuilder>,
5536 page_encodings: Arc<[PrimitivePageEncoding]>,
5537 column_index: u32,
5538 field: Field,
5539 encoding_metadata: Arc<HashMap<String, String>>,
5540}
5541
5542struct CompressedLevelsChunk {
5543 data: LanceBuffer,
5544 num_levels: u16,
5545}
5546
5547struct CompressedLevels {
5548 data: Vec<CompressedLevelsChunk>,
5549 compression: CompressiveEncoding,
5550 rep_index: Option<LanceBuffer>,
5551}
5552
5553struct SerializedMiniBlockPage {
5554 num_buffers: u64,
5555 data: LanceBuffer,
5556 metadata: LanceBuffer,
5557}
5558
5559#[derive(Debug, Clone, Copy)]
5560struct DictEncodingBudget {
5561 max_dict_entries: u32,
5562 max_encoded_size: usize,
5563}
5564
5565enum PrimitivePageStructure {
5566 Dense {
5567 repdef: SerializedRepDefs,
5568 single_row_miniblock_repdef_levels: Option<u64>,
5569 },
5570 Sparse {
5571 plan: sparse::SparseStructuralPlan,
5572 prepared_values: Option<sparse::writer::PreparedSparseValues>,
5573 },
5574}
5575
5576struct PrimitivePageData {
5578 arrays: Vec<ArrayRef>,
5580 structure: PrimitivePageStructure,
5582 row_number: u64,
5584 num_rows: u64,
5586}
5587
5588struct PrimitivePlanContext<'a> {
5589 column_idx: u32,
5590 field: &'a Field,
5591 encoding_metadata: &'a HashMap<String, String>,
5592}
5593
5594enum PrimitiveEncodeAttempt {
5595 Encoded(EncodedPage),
5596 Unhandled(PrimitivePageData),
5597}
5598
5599#[derive(Clone)]
5604struct PrimitiveEncodeContext {
5605 column_idx: u32,
5607 field: Field,
5608 encoding_metadata: Arc<HashMap<String, String>>,
5609 is_simple_validity: bool,
5610 has_repdef_info: bool,
5611}
5612
5613impl PrimitiveStructuralEncoder {
5614 pub fn try_new(
5615 options: &EncodingOptions,
5616 page_encodings: Arc<[PrimitivePageEncoding]>,
5617 column_index: u32,
5618 field: Field,
5619 encoding_metadata: Arc<HashMap<String, String>>,
5620 ) -> Result<Self> {
5621 for page_encoding in page_encodings.iter() {
5622 page_encoding
5623 .behavior
5624 .validate_field(&field, &encoding_metadata)?;
5625 }
5626 Ok(Self {
5627 accumulation_queue: AccumulationQueue::new(
5628 options.cache_bytes_per_column,
5629 column_index,
5630 options.keep_original_array,
5631 ),
5632 keep_original_array: options.keep_original_array,
5633 accumulated_repdefs: Vec::new(),
5634 column_index,
5635 page_encodings,
5636 field,
5637 encoding_metadata,
5638 })
5639 }
5640
5641 fn encode_page(
5642 page_encodings: &[PrimitivePageEncoding],
5643 ctx: &PrimitiveEncodeContext,
5644 mut page: PrimitivePageData,
5645 ) -> Result<EncodedPage> {
5646 for page_encoding in page_encodings {
5647 match page_encoding.behavior.try_encode_page(ctx, page)? {
5648 PrimitiveEncodeAttempt::Encoded(page) => return Ok(page),
5649 PrimitiveEncodeAttempt::Unhandled(unhandled) => page = unhandled,
5650 }
5651 }
5652 Err(Error::invalid_input_source(
5653 format!(
5654 "No primitive page encoding atom supports field '{}'",
5655 ctx.field.name
5656 )
5657 .into(),
5658 ))
5659 }
5660
5661 fn is_narrow(data_block: &DataBlock) -> bool {
5669 const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
5670
5671 if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
5672 let max_len_array = max_len_array
5673 .as_any()
5674 .downcast_ref::<PrimitiveArray<UInt64Type>>()
5675 .unwrap();
5676 if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
5677 return true;
5678 }
5679 }
5680 false
5681 }
5682
5683 fn prefers_miniblock(
5684 data_block: &DataBlock,
5685 encoding_metadata: &HashMap<String, String>,
5686 ) -> bool {
5687 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
5689 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
5690 }
5691 Self::is_narrow(data_block)
5693 }
5694
5695 fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
5696 if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
5700 return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
5701 }
5702 true
5703 }
5704
5705 fn serialize_miniblocks(
5752 miniblocks: MiniBlockCompressed,
5753 rep: Option<Vec<CompressedLevelsChunk>>,
5754 def: Option<Vec<CompressedLevelsChunk>>,
5755 miniblock_chunk_size: MiniblockChunkSize,
5756 ) -> Result<SerializedMiniBlockPage> {
5757 let bytes_rep = rep
5758 .as_ref()
5759 .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
5760 .unwrap_or(0);
5761 let bytes_def = def
5762 .as_ref()
5763 .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
5764 .unwrap_or(0);
5765 let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
5766 let mut num_buffers = miniblocks.data.len();
5767 if rep.is_some() {
5768 num_buffers += 1;
5769 }
5770 if def.is_some() {
5771 num_buffers += 1;
5772 }
5773 let max_extra = 9 * num_buffers;
5775 let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
5776 let chunk_size_bytes = match miniblock_chunk_size {
5777 MiniblockChunkSize::U16 => 2,
5778 MiniblockChunkSize::U32 => 4,
5779 };
5780 let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
5781
5782 let mut rep_iter = rep.map(|r| r.into_iter());
5783 let mut def_iter = def.map(|d| d.into_iter());
5784
5785 let mut buffer_offsets = vec![0; miniblocks.data.len()];
5786 for chunk in miniblocks.chunks {
5787 let start_pos = data_buffer.len();
5788 debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
5790
5791 let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
5792 let def = def_iter.as_mut().map(|d| d.next().unwrap());
5793
5794 let num_levels = rep
5796 .as_ref()
5797 .map(|r| r.num_levels)
5798 .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
5799 data_buffer.extend_from_slice(&num_levels.to_le_bytes());
5800
5801 if let Some(rep) = rep.as_ref() {
5803 let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
5804 Error::internal(format!(
5805 "Repetition buffer size ({} bytes) too large",
5806 rep.data.len()
5807 ))
5808 })?;
5809 data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
5810 }
5811 if let Some(def) = def.as_ref() {
5812 let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
5813 Error::internal(format!(
5814 "Definition buffer size ({} bytes) too large",
5815 def.data.len()
5816 ))
5817 })?;
5818 data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
5819 }
5820
5821 if miniblock_chunk_size == MiniblockChunkSize::U32 {
5822 for &buffer_size in &chunk.buffer_sizes {
5823 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
5824 }
5825 } else {
5826 for &buffer_size in &chunk.buffer_sizes {
5827 let buffer_size = u16::try_from(buffer_size).map_err(|_| {
5828 Error::internal(format!(
5829 "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
5830 buffer_size
5831 ))
5832 })?;
5833 data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
5834 }
5835 }
5836
5837 let add_padding = |data_buffer: &mut Vec<u8>| {
5839 let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
5840 data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
5841 };
5842 add_padding(&mut data_buffer);
5843
5844 if let Some(rep) = rep.as_ref() {
5846 data_buffer.extend_from_slice(&rep.data);
5847 add_padding(&mut data_buffer);
5848 }
5849 if let Some(def) = def.as_ref() {
5850 data_buffer.extend_from_slice(&def.data);
5851 add_padding(&mut data_buffer);
5852 }
5853 for (buffer_size, (buffer, buffer_offset)) in chunk
5854 .buffer_sizes
5855 .iter()
5856 .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
5857 {
5858 let start = *buffer_offset;
5859 let end = start + *buffer_size as usize;
5860 *buffer_offset += *buffer_size as usize;
5861 data_buffer.extend_from_slice(&buffer[start..end]);
5862 add_padding(&mut data_buffer);
5863 }
5864
5865 let chunk_bytes = data_buffer.len() - start_pos;
5866 let max_chunk_size = match miniblock_chunk_size {
5867 MiniblockChunkSize::U16 => 32 * 1024,
5868 MiniblockChunkSize::U32 => 1_u64 << 31,
5869 };
5870 if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
5871 return Err(Error::internal(format!(
5872 "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
5873 chunk_bytes, max_chunk_size
5874 )));
5875 }
5876 if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
5877 return Err(Error::internal(format!(
5878 "Mini-block chunk size {} bytes is not aligned to {} bytes",
5879 chunk_bytes, MINIBLOCK_ALIGNMENT
5880 )));
5881 }
5882 if chunk.log_num_values > 15 {
5883 return Err(Error::internal(format!(
5884 "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
5885 chunk.log_num_values
5886 )));
5887 }
5888 let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
5892 let divided_bytes_minus_one = (divided_bytes - 1) as u64;
5893
5894 let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
5895 if miniblock_chunk_size == MiniblockChunkSize::U32 {
5896 meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
5897 } else {
5898 meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
5899 }
5900 }
5901
5902 let data_buffer = LanceBuffer::from(data_buffer);
5903 let metadata_buffer = LanceBuffer::from(meta_buffer);
5904
5905 Ok(SerializedMiniBlockPage {
5906 num_buffers: miniblocks.data.len() as u64,
5907 data: data_buffer,
5908 metadata: metadata_buffer,
5909 })
5910 }
5911
5912 fn compress_levels(
5917 mut levels: RepDefSlicer<'_>,
5918 num_elements: u64,
5919 compression_strategy: &dyn CompressionStrategy,
5920 chunks: &[MiniBlockChunk],
5921 max_rep: u16,
5923 ) -> Result<CompressedLevels> {
5924 let mut rep_index = if max_rep > 0 {
5925 Vec::with_capacity(chunks.len())
5926 } else {
5927 vec![]
5928 };
5929 let num_levels = levels.num_levels() as u64;
5931 let levels_buf = levels.all_levels().clone();
5932
5933 let mut fixed_width_block = FixedWidthDataBlock {
5934 data: levels_buf,
5935 bits_per_value: 16,
5936 num_values: num_levels,
5937 block_info: BlockInfo::new(),
5938 };
5939 fixed_width_block.compute_stat();
5941
5942 let levels_block = DataBlock::FixedWidth(fixed_width_block);
5943 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
5944 let compressor =
5946 compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
5947 let mut compressor_desc = None;
5948 let mut level_chunks = Vec::with_capacity(chunks.len());
5950 let mut values_counter = 0;
5951 for (chunk_idx, chunk) in chunks.iter().enumerate() {
5952 let chunk_num_values = chunk.num_values(values_counter, num_elements);
5953 debug_assert!(chunk_num_values > 0);
5954 values_counter += chunk_num_values;
5955 let chunk_levels = if chunk_idx < chunks.len() - 1 {
5956 levels.slice_next(chunk_num_values as usize)
5957 } else {
5958 levels.slice_rest()
5959 };
5960 let num_chunk_levels = (chunk_levels.len() / 2) as u64;
5961 if max_rep > 0 {
5962 let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
5972 let rep_values = rep_values.as_ref();
5973
5974 let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
5977 let num_leftovers = if chunk_idx < chunks.len() - 1 {
5978 rep_values
5979 .iter()
5980 .rev()
5981 .position(|v| *v == max_rep)
5982 .map(|pos| pos + 1)
5984 .unwrap_or(rep_values.len())
5985 } else {
5986 0
5988 };
5989
5990 if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
5991 let rep_len = rep_index.len();
5995 if rep_index[rep_len - 1] != 0 {
5996 rep_index[rep_len - 2] += 1;
5998 rep_index[rep_len - 1] = 0;
5999 }
6000 }
6001
6002 if chunk_idx == chunks.len() - 1 {
6003 num_rows += 1;
6005 }
6006 rep_index.push(num_rows as u64);
6007 rep_index.push(num_leftovers as u64);
6008 }
6009 let mut chunk_fixed_width = FixedWidthDataBlock {
6010 data: chunk_levels,
6011 bits_per_value: 16,
6012 num_values: num_chunk_levels,
6013 block_info: BlockInfo::new(),
6014 };
6015 chunk_fixed_width.compute_stat();
6016 let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
6017 let (compressed_levels, chunk_compressor_desc) =
6018 compressor.compress(chunk_levels_block)?;
6019 if let Some(compressor_desc) = compressor_desc.as_ref() {
6020 if compressor_desc != &chunk_compressor_desc {
6021 return Err(Error::internal(
6022 "Rep/def block compressor changed encoding between chunks".to_string(),
6023 ));
6024 }
6025 } else {
6026 compressor_desc = Some(chunk_compressor_desc);
6027 }
6028 let compressed_levels = compressed_levels.ok_or_else(|| {
6029 Error::internal(
6030 "Rep/def block compressor selected a metadata-only codec".to_string(),
6031 )
6032 })?;
6033 let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
6034 Error::invalid_input_source(
6035 format!(
6036 "Mini-block cannot encode {} rep/def levels in one chunk. \
6037 This usually means a top-level row contains too much nested structure \
6038 for the current layout.",
6039 num_chunk_levels
6040 )
6041 .into(),
6042 )
6043 })?;
6044 level_chunks.push(CompressedLevelsChunk {
6045 data: compressed_levels,
6046 num_levels,
6047 });
6048 }
6049 debug_assert_eq!(levels.num_levels_remaining(), 0);
6050 let rep_index = if rep_index.is_empty() {
6051 None
6052 } else {
6053 Some(LanceBuffer::reinterpret_vec(rep_index))
6054 };
6055 Ok(CompressedLevels {
6056 data: level_chunks,
6057 compression: compressor_desc.ok_or_else(|| {
6058 Error::internal("Rep/def compression produced no chunks".to_string())
6059 })?,
6060 rep_index,
6061 })
6062 }
6063
6064 fn encode_simple_all_null(
6065 column_idx: u32,
6066 num_rows: u64,
6067 row_number: u64,
6068 ) -> Result<EncodedPage> {
6069 let description =
6070 ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
6071 Ok(EncodedPage {
6072 column_idx,
6073 data: vec![],
6074 description: PageEncoding::Structural(description),
6075 num_rows,
6076 row_number,
6077 })
6078 }
6079
6080 fn encode_complex_all_null_vals(
6081 data: &Arc<[u16]>,
6082 compression_strategy: &dyn CompressionStrategy,
6083 ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
6084 let buffer = LanceBuffer::reinterpret_slice(data.clone());
6085 let mut fixed_width_block = FixedWidthDataBlock {
6086 data: buffer,
6087 bits_per_value: 16,
6088 num_values: data.len() as u64,
6089 block_info: BlockInfo::new(),
6090 };
6091 fixed_width_block.compute_stat();
6092
6093 let levels_block = DataBlock::FixedWidth(fixed_width_block);
6094 let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
6095 let (compressed_buffer, encoding) =
6096 compress_required_block(compression_strategy, &levels_field, levels_block)?;
6097 Ok((compressed_buffer, encoding))
6098 }
6099
6100 fn encode_complex_all_null(
6104 column_idx: u32,
6105 repdef: crate::repdef::SerializedRepDefs,
6106 row_number: u64,
6107 num_rows: u64,
6108 complex_null_encoding: ComplexNullEncoding,
6109 compression_strategy: &dyn CompressionStrategy,
6110 ) -> Result<EncodedPage> {
6111 if complex_null_encoding == ComplexNullEncoding::RawLevels {
6112 let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
6113 LanceBuffer::reinterpret_slice(rep.clone())
6114 } else {
6115 LanceBuffer::empty()
6116 };
6117
6118 let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
6119 LanceBuffer::reinterpret_slice(def.clone())
6120 } else {
6121 LanceBuffer::empty()
6122 };
6123
6124 let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
6125 return Ok(EncodedPage {
6126 column_idx,
6127 data: vec![rep_bytes, def_bytes],
6128 description: PageEncoding::Structural(description),
6129 num_rows,
6130 row_number,
6131 });
6132 }
6133
6134 let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
6135 repdef.repetition_levels.as_ref()
6136 {
6137 let num_values = rep.len() as u64;
6138 let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
6139 (buffer, Some(encoding), num_values)
6140 } else {
6141 (LanceBuffer::empty(), None, 0)
6142 };
6143
6144 let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
6145 repdef.definition_levels.as_ref()
6146 {
6147 let num_values = def.len() as u64;
6148 let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
6149 (buffer, Some(encoding), num_values)
6150 } else {
6151 (LanceBuffer::empty(), None, 0)
6152 };
6153
6154 let description = ProtobufUtils21::compressed_all_null_constant_layout(
6155 &repdef.def_meaning,
6156 rep_encoding,
6157 def_encoding,
6158 num_rep_values,
6159 num_def_values,
6160 );
6161 Ok(EncodedPage {
6162 column_idx,
6163 data: vec![rep_bytes, def_bytes],
6164 description: PageEncoding::Structural(description),
6165 num_rows,
6166 row_number,
6167 })
6168 }
6169
6170 fn leaf_validity(
6171 repdef: &crate::repdef::SerializedRepDefs,
6172 num_values: usize,
6173 ) -> Result<Option<BooleanBuffer>> {
6174 let rep = repdef
6175 .repetition_levels
6176 .as_ref()
6177 .map(|rep| rep.as_ref().to_vec());
6178 let def = repdef
6179 .definition_levels
6180 .as_ref()
6181 .map(|def| def.as_ref().to_vec());
6182 let mut unraveler = RepDefUnraveler::new(
6183 rep,
6184 def,
6185 repdef.def_meaning.clone().into(),
6186 num_values as u64,
6187 );
6188 if unraveler.is_all_valid() {
6189 return Ok(None);
6190 }
6191 let mut validity = BooleanBufferBuilder::new(num_values);
6192 unraveler.unravel_validity(&mut validity)?;
6193 Ok(Some(validity.finish()))
6194 }
6195
6196 fn is_constant_values(
6197 arrays: &[ArrayRef],
6198 scalar: &ArrayRef,
6199 validity: Option<&BooleanBuffer>,
6200 ) -> Result<bool> {
6201 debug_assert_eq!(scalar.len(), 1);
6202 debug_assert_eq!(scalar.null_count(), 0);
6203
6204 match scalar.data_type() {
6205 DataType::Boolean => {
6206 let mut global_idx = 0usize;
6207 let scalar_val = scalar.as_boolean().value(0);
6208 for arr in arrays {
6209 let bool_arr = arr.as_boolean();
6210 for i in 0..arr.len() {
6211 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6212 global_idx += 1;
6213 if !is_valid {
6214 continue;
6215 }
6216 if bool_arr.value(i) != scalar_val {
6217 return Ok(false);
6218 }
6219 }
6220 }
6221 Ok(true)
6222 }
6223 DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
6224 DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
6225 DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
6226 DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
6227 data_type => {
6228 let mut global_idx = 0usize;
6229 let Some(byte_width) = data_type.byte_width_opt() else {
6230 return Ok(false);
6231 };
6232 let scalar_data = scalar.to_data();
6233 if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
6234 return Ok(false);
6235 }
6236 let scalar_bytes = scalar_data.buffers()[0].as_slice();
6237 if scalar_bytes.len() != byte_width {
6238 return Ok(false);
6239 }
6240
6241 for arr in arrays {
6242 let data = arr.to_data();
6243 if data.buffers().is_empty() {
6244 return Ok(false);
6245 }
6246 let buf = data.buffers()[0].as_slice();
6247 let base = data.offset();
6248 for i in 0..arr.len() {
6249 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6250 global_idx += 1;
6251 if !is_valid {
6252 continue;
6253 }
6254 let start = (base + i) * byte_width;
6255 if buf[start..start + byte_width] != scalar_bytes[..] {
6256 return Ok(false);
6257 }
6258 }
6259 }
6260 Ok(true)
6261 }
6262 }
6263 }
6264
6265 fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
6266 arrays: &[ArrayRef],
6267 scalar: &ArrayRef,
6268 validity: Option<&BooleanBuffer>,
6269 ) -> Result<bool> {
6270 debug_assert_eq!(scalar.len(), 1);
6271 let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
6272 let mut global_idx = 0usize;
6273 for arr in arrays {
6274 let str_arr = arr.as_string::<O>();
6275 for i in 0..arr.len() {
6276 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6277 global_idx += 1;
6278 if !is_valid {
6279 continue;
6280 }
6281 if str_arr.value(i).as_bytes() != scalar_val {
6282 return Ok(false);
6283 }
6284 }
6285 }
6286 Ok(true)
6287 }
6288
6289 fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
6290 arrays: &[ArrayRef],
6291 scalar: &ArrayRef,
6292 validity: Option<&BooleanBuffer>,
6293 ) -> Result<bool> {
6294 debug_assert_eq!(scalar.len(), 1);
6295 let scalar_val = scalar.as_binary::<O>().value(0);
6296 let mut global_idx = 0usize;
6297 for arr in arrays {
6298 let bin_arr = arr.as_binary::<O>();
6299 for i in 0..arr.len() {
6300 let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6301 global_idx += 1;
6302 if !is_valid {
6303 continue;
6304 }
6305 if bin_arr.value(i) != scalar_val {
6306 return Ok(false);
6307 }
6308 }
6309 }
6310 Ok(true)
6311 }
6312
6313 fn find_constant_scalar(
6314 arrays: &[ArrayRef],
6315 validity: Option<&BooleanBuffer>,
6316 ) -> Result<Option<ArrayRef>> {
6317 if arrays.is_empty() {
6318 return Ok(None);
6319 }
6320
6321 let global_scalar_idx = if let Some(validity) = validity {
6322 let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
6323 return Ok(None);
6324 };
6325 idx
6326 } else {
6327 0
6328 };
6329
6330 let mut idx_remaining = global_scalar_idx;
6331 let mut scalar_arr_idx = 0usize;
6332 while scalar_arr_idx < arrays.len() {
6333 let len = arrays[scalar_arr_idx].len();
6334 if idx_remaining < len {
6335 break;
6336 }
6337 idx_remaining -= len;
6338 scalar_arr_idx += 1;
6339 }
6340
6341 if scalar_arr_idx >= arrays.len() {
6342 return Ok(None);
6343 }
6344
6345 let scalar =
6346 lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
6347 if scalar.null_count() != 0 {
6348 return Ok(None);
6349 }
6350 if !Self::is_constant_values(arrays, &scalar, validity)? {
6351 return Ok(None);
6352 }
6353 Ok(Some(scalar))
6354 }
6355
6356 fn resolve_dict_values_compression_metadata(
6357 field_metadata: &HashMap<String, String>,
6358 env_compression: Option<String>,
6359 env_compression_level: Option<String>,
6360 ) -> HashMap<String, String> {
6361 let mut metadata = HashMap::new();
6362
6363 let compression = field_metadata
6364 .get(DICT_VALUES_COMPRESSION_META_KEY)
6365 .cloned()
6366 .or(env_compression)
6367 .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
6368 metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
6369
6370 if let Some(compression_level) = field_metadata
6371 .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
6372 .cloned()
6373 .or(env_compression_level)
6374 {
6375 metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
6376 }
6377
6378 metadata
6379 }
6380
6381 fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
6382 let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
6387 dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
6388 &field.metadata,
6389 env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
6390 env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
6391 );
6392 Ok(dict_values_field)
6393 }
6394
6395 #[allow(clippy::too_many_arguments)]
6396 fn encode_miniblock(
6397 column_idx: u32,
6398 field: &Field,
6399 compression_strategy: &dyn CompressionStrategy,
6400 data: DataBlock,
6401 repdef: crate::repdef::SerializedRepDefs,
6402 row_number: u64,
6403 dictionary_data: Option<DataBlock>,
6404 num_rows: u64,
6405 miniblock_chunk_size: MiniblockChunkSize,
6406 ) -> Result<EncodedPage> {
6407 if let DataBlock::AllNull(_null_block) = data {
6408 unreachable!()
6411 }
6412
6413 let num_items = data.num_values();
6414
6415 let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
6416 let common_chunk_buffers =
6417 u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some());
6418 let support_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
6419 let compression_context =
6420 MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true);
6421 let (compressed_data, value_encoding) = compressor.compress(compression_context, data)?;
6422
6423 let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
6424
6425 let mut compressed_rep = repdef
6426 .rep_slicer()
6427 .map(|rep_slicer| {
6428 Self::compress_levels(
6429 rep_slicer,
6430 num_items,
6431 compression_strategy,
6432 &compressed_data.chunks,
6433 max_rep,
6434 )
6435 })
6436 .transpose()?;
6437
6438 let (rep_index, rep_index_depth) =
6439 match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
6440 Some(rep_index) => (Some(rep_index.clone()), 1),
6441 None => (None, 0),
6442 };
6443
6444 let mut compressed_def = repdef
6445 .def_slicer()
6446 .map(|def_slicer| {
6447 Self::compress_levels(
6448 def_slicer,
6449 num_items,
6450 compression_strategy,
6451 &compressed_data.chunks,
6452 0,
6453 )
6454 })
6455 .transpose()?;
6456
6457 let rep_data = compressed_rep
6463 .as_mut()
6464 .map(|cr| std::mem::take(&mut cr.data));
6465 let def_data = compressed_def
6466 .as_mut()
6467 .map(|cd| std::mem::take(&mut cd.data));
6468
6469 let serialized =
6470 Self::serialize_miniblocks(compressed_data, rep_data, def_data, miniblock_chunk_size)?;
6471 let has_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
6472
6473 let mut data = Vec::with_capacity(4);
6475 data.push(serialized.metadata);
6476 data.push(serialized.data);
6477
6478 if let Some(dictionary_data) = dictionary_data {
6479 let num_dictionary_items = dictionary_data.num_values();
6480 let dict_values_field = Self::build_dict_values_compressor_field(field)?;
6481
6482 let (dictionary_buffer, dictionary_encoding) =
6483 compress_required_block(compression_strategy, &dict_values_field, dictionary_data)?;
6484
6485 data.push(dictionary_buffer);
6486 if let Some(rep_index) = rep_index {
6487 data.push(rep_index);
6488 }
6489
6490 let description = ProtobufUtils21::miniblock_layout(
6491 compressed_rep.map(|cr| cr.compression),
6492 compressed_def.map(|cd| cd.compression),
6493 value_encoding,
6494 rep_index_depth,
6495 serialized.num_buffers,
6496 Some((dictionary_encoding, num_dictionary_items)),
6497 &repdef.def_meaning,
6498 num_items,
6499 has_large_chunk,
6500 );
6501 Ok(EncodedPage {
6502 num_rows,
6503 column_idx,
6504 data,
6505 description: PageEncoding::Structural(description),
6506 row_number,
6507 })
6508 } else {
6509 let description = ProtobufUtils21::miniblock_layout(
6510 compressed_rep.map(|cr| cr.compression),
6511 compressed_def.map(|cd| cd.compression),
6512 value_encoding,
6513 rep_index_depth,
6514 serialized.num_buffers,
6515 None,
6516 &repdef.def_meaning,
6517 num_items,
6518 has_large_chunk,
6519 );
6520
6521 if let Some(rep_index) = rep_index {
6522 let view = rep_index.borrow_to_typed_slice::<u64>();
6523 let total = view.as_chunks::<2>().0.iter().map(|c| c[0]).sum::<u64>();
6524 debug_assert_eq!(total, num_rows);
6525
6526 data.push(rep_index);
6527 }
6528
6529 Ok(EncodedPage {
6530 num_rows,
6531 column_idx,
6532 data,
6533 description: PageEncoding::Structural(description),
6534 row_number,
6535 })
6536 }
6537 }
6538
6539 fn serialize_full_zip_fixed(
6541 fixed: FixedWidthDataBlock,
6542 mut repdef: ControlWordIterator,
6543 num_values: u64,
6544 ) -> Result<SerializedFullZip> {
6545 if !fixed.bits_per_value.is_multiple_of(8) {
6546 return Err(Error::invalid_input_source(
6547 format!(
6548 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
6549 fixed.bits_per_value
6550 )
6551 .into(),
6552 ));
6553 }
6554
6555 let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
6556 let mut zipped_data = Vec::with_capacity(len);
6557
6558 let max_rep_index_val = if repdef.has_repetition() {
6559 len as u64
6560 } else {
6561 0
6563 };
6564 let mut rep_index_builder =
6565 BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
6566
6567 let bytes_per_value = fixed.bits_per_value as usize / 8;
6568 let mut offset = 0;
6569
6570 if bytes_per_value == 0 {
6571 while let Some(control) = repdef.append_next(&mut zipped_data) {
6573 if control.is_new_row {
6574 debug_assert!(offset <= len);
6576 rep_index_builder.append_trusted(offset as u64);
6577 }
6578 offset = zipped_data.len();
6579 }
6580 } else {
6581 let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
6583 while let Some(control) = repdef.append_next(&mut zipped_data) {
6584 if control.is_new_row {
6585 debug_assert!(offset <= len);
6587 rep_index_builder.append_trusted(offset as u64);
6588 }
6589 if control.is_visible {
6590 let value = data_iter.next().unwrap();
6591 zipped_data.extend_from_slice(value);
6592 }
6593 offset = zipped_data.len();
6594 }
6595 }
6596
6597 debug_assert_eq!(zipped_data.len(), len);
6598 rep_index_builder.append_trusted(zipped_data.len() as u64);
6600
6601 let zipped_data = LanceBuffer::from(zipped_data);
6602 let rep_index = rep_index_builder.into_data();
6603 let rep_index = if rep_index.is_empty() {
6604 None
6605 } else {
6606 Some(LanceBuffer::from(rep_index))
6607 };
6608 Ok(SerializedFullZip {
6609 values: zipped_data,
6610 repetition_index: rep_index,
6611 })
6612 }
6613
6614 fn serialize_full_zip_variable(
6618 variable: VariableWidthBlock,
6619 mut repdef: ControlWordIterator,
6620 num_items: u64,
6621 ) -> Result<SerializedFullZip> {
6622 let bytes_per_offset = variable.bits_per_offset as usize / 8;
6623 if !variable.bits_per_offset.is_multiple_of(8) {
6624 return Err(Error::invalid_input_source(
6625 format!(
6626 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
6627 variable.bits_per_offset
6628 )
6629 .into(),
6630 ));
6631 }
6632 let len = variable.data.len()
6633 + repdef.bytes_per_word() * num_items as usize
6634 + bytes_per_offset * variable.num_values as usize;
6635 let mut buf = Vec::with_capacity(len);
6636
6637 let max_rep_index_val = len as u64;
6638 let mut rep_index_builder =
6639 BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
6640
6641 match bytes_per_offset {
6643 4 => {
6644 let offs = variable.offsets.borrow_to_typed_slice::<u32>();
6645 let mut rep_offset = 0;
6646 let mut windows_iter = offs.as_ref().windows(2);
6647 while let Some(control) = repdef.append_next(&mut buf) {
6648 if control.is_new_row {
6649 debug_assert!(rep_offset <= len);
6651 rep_index_builder.append_trusted(rep_offset as u64);
6652 }
6653 if control.is_visible {
6654 let window = windows_iter.next().unwrap();
6655 if control.is_valid_item {
6656 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
6657 buf.extend_from_slice(
6658 &variable.data[window[0] as usize..window[1] as usize],
6659 );
6660 }
6661 }
6662 rep_offset = buf.len();
6663 }
6664 }
6665 8 => {
6666 let offs = variable.offsets.borrow_to_typed_slice::<u64>();
6667 let mut rep_offset = 0;
6668 let mut windows_iter = offs.as_ref().windows(2);
6669 while let Some(control) = repdef.append_next(&mut buf) {
6670 if control.is_new_row {
6671 debug_assert!(rep_offset <= len);
6673 rep_index_builder.append_trusted(rep_offset as u64);
6674 }
6675 if control.is_visible {
6676 let window = windows_iter.next().unwrap();
6677 if control.is_valid_item {
6678 buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
6679 buf.extend_from_slice(
6680 &variable.data[window[0] as usize..window[1] as usize],
6681 );
6682 }
6683 }
6684 rep_offset = buf.len();
6685 }
6686 }
6687 _ => {
6688 return Err(Error::invalid_input_source(
6689 format!(
6690 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
6691 variable.bits_per_offset
6692 )
6693 .into(),
6694 ));
6695 }
6696 }
6697
6698 debug_assert!(buf.len() <= len);
6701 rep_index_builder.append_trusted(buf.len() as u64);
6703
6704 let zipped_data = LanceBuffer::from(buf);
6705 let rep_index = rep_index_builder.into_data();
6706 debug_assert!(!rep_index.is_empty());
6707 let rep_index = Some(LanceBuffer::from(rep_index));
6708 Ok(SerializedFullZip {
6709 values: zipped_data,
6710 repetition_index: rep_index,
6711 })
6712 }
6713
6714 fn serialize_full_zip(
6717 compressed_data: PerValueDataBlock,
6718 repdef: ControlWordIterator,
6719 num_items: u64,
6720 ) -> Result<SerializedFullZip> {
6721 match compressed_data {
6722 PerValueDataBlock::Fixed(fixed) => {
6723 Self::serialize_full_zip_fixed(fixed, repdef, num_items)
6724 }
6725 PerValueDataBlock::Variable(var) => {
6726 Self::serialize_full_zip_variable(var, repdef, num_items)
6727 }
6728 }
6729 }
6730
6731 fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
6732 debug_assert_eq!(fixed.bits_per_value, 1);
6733 let num_values = fixed.num_values as usize;
6734 let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
6735 let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
6736 FixedWidthDataBlock {
6737 data: LanceBuffer::from(expanded),
6738 bits_per_value: 8,
6739 num_values: fixed.num_values,
6740 block_info: BlockInfo::new(),
6741 }
6742 }
6743
6744 fn encode_full_zip(
6745 column_idx: u32,
6746 field: &Field,
6747 compression_strategy: &dyn CompressionStrategy,
6748 data: DataBlock,
6749 repdef: crate::repdef::SerializedRepDefs,
6750 row_number: u64,
6751 num_lists: u64,
6752 ) -> Result<EncodedPage> {
6753 let max_rep = repdef
6754 .repetition_levels
6755 .as_ref()
6756 .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
6757 let max_def = repdef
6758 .definition_levels
6759 .as_ref()
6760 .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
6761
6762 let (num_items, num_visible_items) =
6766 if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
6767 (rep_levels.len() as u64, data.num_values())
6770 } else {
6771 (data.num_values(), data.num_values())
6773 };
6774
6775 let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
6776
6777 let repdef_iter = build_control_word_iterator(
6778 repdef.repetition_levels.as_deref(),
6779 max_rep,
6780 repdef.definition_levels.as_deref(),
6781 max_def,
6782 max_visible_def,
6783 num_items as usize,
6784 );
6785 let bits_rep = repdef_iter.bits_rep();
6786 let bits_def = repdef_iter.bits_def();
6787
6788 let data = match data {
6790 DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
6791 DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
6792 }
6793 other => other,
6794 };
6795
6796 let compressor = compression_strategy.create_per_value(field, &data)?;
6797 let (compressed_data, value_encoding) = compressor.compress(data)?;
6798
6799 let description = match &compressed_data {
6800 PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
6801 bits_rep,
6802 bits_def,
6803 fixed.bits_per_value as u32,
6804 value_encoding,
6805 &repdef.def_meaning,
6806 num_items as u32,
6807 num_visible_items as u32,
6808 ),
6809 PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
6810 bits_rep,
6811 bits_def,
6812 variable.bits_per_offset as u32,
6813 value_encoding,
6814 &repdef.def_meaning,
6815 num_items as u32,
6816 num_visible_items as u32,
6817 ),
6818 };
6819
6820 let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
6821
6822 let data = if let Some(repindex) = zipped.repetition_index {
6823 vec![zipped.values, repindex]
6824 } else {
6825 vec![zipped.values]
6826 };
6827
6828 Ok(EncodedPage {
6829 num_rows: num_lists,
6830 column_idx,
6831 data,
6832 description: PageEncoding::Structural(description),
6833 row_number,
6834 })
6835 }
6836
6837 fn should_dictionary_encode(
6838 data_block: &DataBlock,
6839 field: &Field,
6840 fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
6841 ) -> Option<DictEncodingBudget> {
6842 const DEFAULT_SAMPLE_SIZE: usize = 4096;
6843 const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
6844
6845 match data_block {
6848 DataBlock::FixedWidth(fixed) => {
6849 if fixed.bits_per_value == 64
6850 && fixed_width_dictionary_encoding == FixedWidthDictionaryEncoding::Exclude64Bit
6851 {
6852 return None;
6853 }
6854 if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
6855 return None;
6856 }
6857 if fixed.bits_per_value % 8 != 0 {
6858 return None;
6859 }
6860 }
6861 DataBlock::VariableWidth(var) => {
6862 if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
6863 return None;
6864 }
6865 }
6866 _ => return None,
6867 }
6868
6869 let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
6871 .ok()
6872 .and_then(|val| val.parse().ok())
6873 .unwrap_or(100);
6874 if data_block.num_values() < too_small {
6875 return None;
6876 }
6877
6878 let num_values = data_block.num_values();
6879
6880 let divisor: u64 = field
6883 .metadata
6884 .get(DICT_DIVISOR_META_KEY)
6885 .and_then(|val| val.parse().ok())
6886 .or_else(|| {
6887 env::var("LANCE_ENCODING_DICT_DIVISOR")
6888 .ok()
6889 .and_then(|val| val.parse().ok())
6890 })
6891 .unwrap_or(DEFAULT_DICT_DIVISOR);
6892
6893 let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
6894 .ok()
6895 .and_then(|val| val.parse().ok())
6896 .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
6897
6898 let threshold_cardinality = num_values
6899 .checked_div(divisor.max(1))
6900 .unwrap_or(0)
6901 .min(max_cardinality);
6902 if threshold_cardinality == 0 {
6903 return None;
6904 }
6905
6906 let threshold_ratio = field
6908 .metadata
6909 .get(DICT_SIZE_RATIO_META_KEY)
6910 .and_then(|val| val.parse::<f64>().ok())
6911 .or_else(|| {
6912 env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
6913 .ok()
6914 .and_then(|val| val.parse().ok())
6915 })
6916 .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
6917
6918 if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
6919 panic!(
6920 "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
6921 threshold_ratio
6922 );
6923 }
6924
6925 let data_size = data_block.data_size();
6926 if data_size == 0 {
6927 return None;
6928 }
6929
6930 let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
6931 let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
6932
6933 if let Some(sample_unique_ratio) =
6936 Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
6937 {
6938 if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
6939 return None;
6940 }
6941
6942 let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
6943 if projected_cardinality > threshold_cardinality {
6944 return None;
6945 }
6946 }
6947
6948 let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
6949 Some(DictEncodingBudget {
6950 max_dict_entries,
6951 max_encoded_size,
6952 })
6953 }
6954
6955 fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
6963 use std::collections::HashSet;
6964
6965 const NUM_SAMPLE_BLOCKS: usize = 32;
6966 const MIN_RELIABLE_SAMPLES: usize = 1024;
6967
6968 let num_values = usize::try_from(data_block.num_values()).ok()?;
6969 if num_values == 0 {
6970 return Some(None);
6971 }
6972
6973 let sample_count = num_values.min(max_samples).max(1);
6974 if sample_count < MIN_RELIABLE_SAMPLES {
6975 return Some(None);
6976 }
6977
6978 let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
6979 let samples_per_block = (sample_count / block_count).max(1);
6980 let mut indices = Vec::with_capacity(sample_count);
6981 for block_idx in 0..block_count {
6982 let block_start = block_idx * num_values / block_count;
6983 let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
6984 let block_len = next_block_start.saturating_sub(block_start);
6985 let samples_in_block = samples_per_block.min(block_len);
6986 indices.extend((0..samples_in_block).map(|offset| block_start + offset));
6987 }
6988
6989 if indices.len() < MIN_RELIABLE_SAMPLES {
6990 return Some(None);
6991 }
6992
6993 let ratio = match data_block {
6994 DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
6995 64 => {
6996 let values = fixed.data.borrow_to_typed_slice::<u64>();
6997 let values = values.as_ref();
6998 let mut unique: HashSet<u64> =
6999 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7000 for idx in indices.iter().copied() {
7001 unique.insert(values.get(idx).copied()?);
7002 }
7003 unique.len() as f64 / indices.len() as f64
7004 }
7005 128 => {
7006 let values = fixed.data.borrow_to_typed_slice::<u128>();
7007 let values = values.as_ref();
7008 let mut unique: HashSet<u128> =
7009 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7010 for idx in indices.iter().copied() {
7011 unique.insert(values.get(idx).copied()?);
7012 }
7013 unique.len() as f64 / indices.len() as f64
7014 }
7015 _ => return Some(None),
7016 },
7017 DataBlock::VariableWidth(var) => {
7018 use xxhash_rust::xxh3::xxh3_64;
7019
7020 let mut unique: HashSet<u64> =
7022 HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7023 match var.bits_per_offset {
7024 32 => {
7025 let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
7026 let offsets: &[u32] = offsets_ref.as_ref();
7027 for i in indices.iter().copied() {
7028 let start = usize::try_from(*offsets.get(i)?).ok()?;
7029 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
7030 if start > end || end > var.data.len() {
7031 return None;
7032 }
7033 unique.insert(xxh3_64(&var.data[start..end]));
7034 }
7035 }
7036 64 => {
7037 let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
7038 let offsets: &[u64] = offsets_ref.as_ref();
7039 for i in indices.iter().copied() {
7040 let start = usize::try_from(*offsets.get(i)?).ok()?;
7041 let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
7042 if start > end || end > var.data.len() {
7043 return None;
7044 }
7045 unique.insert(xxh3_64(&var.data[start..end]));
7046 }
7047 }
7048 _ => return Some(None),
7049 }
7050 unique.len() as f64 / indices.len() as f64
7051 }
7052 _ => return Some(None),
7053 };
7054
7055 Some(Some(ratio))
7056 }
7057
7058 fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
7059 let repetition_levels = repdef
7060 .repetition_levels
7061 .as_ref()
7062 .map(|levels| levels[range.clone()].to_vec());
7063 let definition_levels = repdef
7064 .definition_levels
7065 .as_ref()
7066 .map(|levels| levels[range].to_vec());
7067 SerializedRepDefs::new_with_fixed_size_list_levels(
7068 repetition_levels,
7069 definition_levels,
7070 repdef.def_meaning.clone(),
7071 repdef.has_fixed_size_list_levels(),
7072 )
7073 }
7074
7075 fn slice_arrays(
7076 arrays: &[ArrayRef],
7077 value_start: u64,
7078 num_values: u64,
7079 ) -> Result<Vec<ArrayRef>> {
7080 if num_values == 0 {
7081 return Ok(Vec::new());
7082 }
7083
7084 let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
7085 Error::invalid_input(format!("Value start {} is too large", value_start))
7086 })?;
7087 let mut values_remaining = usize::try_from(num_values).map_err(|_| {
7088 Error::invalid_input(format!("Value count {} is too large", num_values))
7089 })?;
7090 let mut sliced = Vec::new();
7091
7092 for array in arrays {
7093 if values_to_skip >= array.len() {
7094 values_to_skip -= array.len();
7095 continue;
7096 }
7097
7098 let offset = values_to_skip;
7099 let len = (array.len() - offset).min(values_remaining);
7100 sliced.push(array.slice(offset, len));
7101 values_remaining -= len;
7102 values_to_skip = 0;
7103
7104 if values_remaining == 0 {
7105 break;
7106 }
7107 }
7108
7109 if values_remaining != 0 {
7110 return Err(Error::internal(format!(
7111 "Page split requested {} values starting at {}, but the page did not contain enough values",
7112 num_values, value_start
7113 )));
7114 }
7115
7116 Ok(sliced)
7117 }
7118
7119 fn split_pages_for_miniblock_repdef_budget(
7120 arrays: Vec<ArrayRef>,
7121 repdef: SerializedRepDefs,
7122 budget: MiniBlockRepDefBudget,
7123 row_number: u64,
7124 num_rows: u64,
7125 ) -> Result<Vec<PrimitivePageData>> {
7126 if budget == MiniBlockRepDefBudget::WithinBudget {
7127 return Ok(vec![PrimitivePageData {
7128 arrays,
7129 structure: PrimitivePageStructure::Dense {
7130 repdef,
7131 single_row_miniblock_repdef_levels: None,
7132 },
7133 row_number,
7134 num_rows,
7135 }]);
7136 }
7137 if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget {
7138 return Ok(vec![PrimitivePageData {
7139 arrays,
7140 structure: PrimitivePageStructure::Dense {
7141 repdef,
7142 single_row_miniblock_repdef_levels: Some(num_levels),
7143 },
7144 row_number,
7145 num_rows,
7146 }]);
7147 }
7148
7149 let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else {
7150 unreachable!();
7151 };
7152
7153 let mut pages = Vec::with_capacity(splits.len());
7154 for split in splits {
7155 let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
7156 let repdef = Self::slice_repdef(&repdef, split.level_range);
7157 pages.push(PrimitivePageData {
7158 arrays,
7159 structure: PrimitivePageStructure::Dense {
7160 repdef,
7161 single_row_miniblock_repdef_levels: None,
7162 },
7163 row_number: row_number + split.row_start,
7164 num_rows: split.num_rows,
7165 });
7166 }
7167 Ok(pages)
7168 }
7169
7170 fn encode_dense_page(
7171 ctx: PrimitiveEncodeContext,
7172 page: PrimitivePageData,
7173 compression_strategy: Arc<dyn CompressionStrategy>,
7174 miniblock_chunk_size: MiniblockChunkSize,
7175 complex_null_encoding: ComplexNullEncoding,
7176 fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
7177 ) -> Result<EncodedPage> {
7178 let PrimitiveEncodeContext {
7179 column_idx,
7180 field,
7181 encoding_metadata,
7182 is_simple_validity,
7183 has_repdef_info,
7184 } = ctx;
7185 let PrimitivePageData {
7186 arrays,
7187 structure,
7188 row_number,
7189 num_rows,
7190 } = page;
7191 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
7192
7193 let (repdef, single_row_miniblock_repdef_levels) = match structure {
7194 PrimitivePageStructure::Dense {
7195 repdef,
7196 single_row_miniblock_repdef_levels,
7197 } => (repdef, single_row_miniblock_repdef_levels),
7198 PrimitivePageStructure::Sparse { .. } => {
7199 unreachable!("dense atom received sparse page")
7200 }
7201 };
7202
7203 if num_values == 0 {
7204 log::debug!(
7207 "Encoding column {} with {} items ({} rows) using complex-null layout",
7208 column_idx,
7209 num_values,
7210 num_rows
7211 );
7212 return Self::encode_complex_all_null(
7213 column_idx,
7214 repdef,
7215 row_number,
7216 num_rows,
7217 complex_null_encoding,
7218 compression_strategy.as_ref(),
7219 );
7220 }
7221
7222 let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
7223 let all_null = leaf_validity
7224 .as_ref()
7225 .map(|validity| validity.count_set_bits() == 0)
7226 .unwrap_or(false);
7227
7228 if all_null {
7229 return if is_simple_validity {
7230 log::debug!(
7231 "Encoding column {} with {} items ({} rows) using simple-null layout",
7232 column_idx,
7233 num_values,
7234 num_rows
7235 );
7236 Self::encode_simple_all_null(column_idx, num_values, row_number)
7237 } else {
7238 log::debug!(
7239 "Encoding column {} with {} items ({} rows) using complex-null layout",
7240 column_idx,
7241 num_values,
7242 num_rows
7243 );
7244 Self::encode_complex_all_null(
7245 column_idx,
7246 repdef,
7247 row_number,
7248 num_rows,
7249 complex_null_encoding,
7250 compression_strategy.as_ref(),
7251 )
7252 };
7253 }
7254
7255 if let DataType::Struct(fields) = &field.data_type()
7256 && fields.is_empty()
7257 {
7258 if has_repdef_info {
7259 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()));
7260 }
7261 return Self::encode_simple_all_null(column_idx, num_values, row_number);
7264 }
7265
7266 let data_block = DataBlock::from_arrays(&arrays, num_values);
7267
7268 if let Some(num_levels) = single_row_miniblock_repdef_levels {
7269 let requested_encoding = encoding_metadata
7270 .get(STRUCTURAL_ENCODING_META_KEY)
7271 .map(|requested| requested.to_lowercase());
7272 let fullzip_error = match &data_block {
7273 DataBlock::FixedWidth(fixed)
7275 if fixed.bits_per_value != 1 && !fixed.bits_per_value.is_multiple_of(8) =>
7276 {
7277 Some(format!(
7278 "Full-zip fixed-width values must be byte aligned, got {} bits per value",
7279 fixed.bits_per_value
7280 ))
7281 }
7282 DataBlock::VariableWidth(variable)
7283 if !variable.bits_per_offset.is_multiple_of(8) =>
7284 {
7285 Some(format!(
7286 "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
7287 variable.bits_per_offset
7288 ))
7289 }
7290 DataBlock::VariableWidth(variable)
7291 if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
7292 {
7293 Some(format!(
7294 "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
7295 variable.bits_per_offset
7296 ))
7297 }
7298 DataBlock::Dictionary(_) => {
7299 Some("Full-zip does not encode dictionary data blocks directly".to_string())
7300 }
7301 DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
7302 Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
7303 Some(flat) => Some(format!(
7304 "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
7305 flat.bits_per_value
7306 )),
7307 None => Some(
7308 "Full-zip fixed-size-list capability requires a flat fixed-width child"
7309 .to_string(),
7310 ),
7311 },
7312 DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
7313 None
7314 }
7315 other => Some(format!(
7316 "Full-zip does not support value block type {}",
7317 other.name()
7318 )),
7319 };
7320 match requested_encoding.as_deref() {
7321 Some(STRUCTURAL_ENCODING_FULLZIP) => {
7322 if let Some(reason) = fullzip_error {
7323 return Err(Error::invalid_input_source(reason.into()));
7324 }
7325 return Self::encode_full_zip(
7326 column_idx,
7327 &field,
7328 compression_strategy.as_ref(),
7329 data_block,
7330 repdef,
7331 row_number,
7332 num_rows,
7333 );
7334 }
7335 Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
7336 if requested_encoding.is_none() && fullzip_error.is_none() {
7337 log::debug!(
7338 "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
7339 column_idx,
7340 num_values
7341 );
7342 return Self::encode_full_zip(
7343 column_idx,
7344 &field,
7345 compression_strategy.as_ref(),
7346 data_block,
7347 repdef,
7348 row_number,
7349 num_rows,
7350 );
7351 }
7352 return Err(Error::invalid_input_source(
7353 format!(
7354 "Mini-block cannot encode {} rep/def levels in one top-level row. \
7355 This usually means the row contains too much nested structure \
7356 for the current layout.",
7357 num_levels
7358 )
7359 .into(),
7360 ));
7361 }
7362 _ => {}
7363 }
7364 }
7365
7366 let requires_full_zip_packed_struct =
7367 if let DataBlock::Struct(ref struct_data_block) = data_block {
7368 struct_data_block.has_variable_width_child()
7369 } else {
7370 false
7371 };
7372
7373 if requires_full_zip_packed_struct {
7374 log::debug!(
7375 "Encoding column {} with {} items using full-zip packed struct layout",
7376 column_idx,
7377 num_values
7378 );
7379 return Self::encode_full_zip(
7380 column_idx,
7381 &field,
7382 compression_strategy.as_ref(),
7383 data_block,
7384 repdef,
7385 row_number,
7386 num_rows,
7387 );
7388 }
7389
7390 if let DataBlock::Dictionary(dict) = data_block {
7391 log::debug!(
7392 "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
7393 column_idx,
7394 num_values
7395 );
7396 let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
7397 indices_data_block.compute_stat();
7402 return Self::encode_miniblock(
7403 column_idx,
7404 &field,
7405 compression_strategy.as_ref(),
7406 indices_data_block,
7407 repdef,
7408 row_number,
7409 Some(dictionary_data_block),
7410 num_rows,
7411 miniblock_chunk_size,
7412 );
7413 }
7414
7415 let dict_result = Self::should_dictionary_encode(
7418 &data_block,
7419 &field,
7420 fixed_width_dictionary_encoding,
7421 )
7422 .and_then(|budget| {
7423 log::debug!(
7424 "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
7425 column_idx,
7426 num_values
7427 );
7428 dict::dictionary_encode(
7429 &data_block,
7430 budget.max_dict_entries,
7431 budget.max_encoded_size,
7432 )
7433 });
7434
7435 if let Some((indices_data_block, dictionary_data_block)) = dict_result {
7436 Self::encode_miniblock(
7437 column_idx,
7438 &field,
7439 compression_strategy.as_ref(),
7440 indices_data_block,
7441 repdef,
7442 row_number,
7443 Some(dictionary_data_block),
7444 num_rows,
7445 miniblock_chunk_size,
7446 )
7447 } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
7448 log::debug!(
7449 "Encoding column {} with {} items using mini-block layout",
7450 column_idx,
7451 num_values
7452 );
7453 Self::encode_miniblock(
7454 column_idx,
7455 &field,
7456 compression_strategy.as_ref(),
7457 data_block,
7458 repdef,
7459 row_number,
7460 None,
7461 num_rows,
7462 miniblock_chunk_size,
7463 )
7464 } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
7465 log::debug!(
7466 "Encoding column {} with {} items using full-zip layout",
7467 column_idx,
7468 num_values
7469 );
7470 Self::encode_full_zip(
7471 column_idx,
7472 &field,
7473 compression_strategy.as_ref(),
7474 data_block,
7475 repdef,
7476 row_number,
7477 num_rows,
7478 )
7479 } else {
7480 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()))
7481 }
7482 }
7483
7484 fn rebuild_empty_dictionary_chunks(arrays: Vec<ArrayRef>) -> Result<Vec<ArrayRef>> {
7491 if !arrays.iter().any(|array| {
7492 array
7493 .as_any_dictionary_opt()
7494 .is_some_and(|dictionary| dictionary.values().is_empty())
7495 }) {
7496 return Ok(arrays);
7497 }
7498
7499 let zeroed = |data_type: &DataType, len: usize| {
7500 new_null_array(data_type, len)
7501 .to_data()
7502 .into_builder()
7503 .nulls(None)
7504 .build()
7505 .map(make_array)
7506 };
7507 let rebuild =
7508 |keys: Vec<ArrayRef>, data_type: &DataType, values: ArrayRef| -> Result<ArrayRef> {
7509 let keys = keys.iter().map(|keys| keys.as_ref()).collect::<Vec<_>>();
7510 let keys = arrow_select::concat::concat(&keys)?;
7511 let data = keys
7512 .to_data()
7513 .into_builder()
7514 .data_type(data_type.clone())
7515 .child_data(vec![values.to_data()])
7516 .build()?;
7517 Ok(make_array(data))
7518 };
7519
7520 let mut rebuilt = Vec::with_capacity(arrays.len());
7521 let mut pending_keys = Vec::with_capacity(arrays.len());
7522 let mut empty_data_type = None;
7523 for array in arrays {
7524 let Some(dictionary) = array.as_any_dictionary_opt() else {
7525 return Err(Error::invalid_input_source(
7526 "Cannot mix dictionary and non-dictionary chunks".into(),
7527 ));
7528 };
7529 if dictionary.values().is_empty() {
7530 pending_keys.push(zeroed(dictionary.keys().data_type(), array.len())?);
7531 empty_data_type.get_or_insert_with(|| array.data_type().clone());
7532 } else if pending_keys.is_empty() {
7533 rebuilt.push(array);
7534 } else {
7535 pending_keys.push(make_array(dictionary.keys().to_data()));
7536 rebuilt.push(rebuild(
7537 std::mem::take(&mut pending_keys),
7538 array.data_type(),
7539 dictionary.values().clone(),
7540 )?);
7541 }
7542 }
7543
7544 if pending_keys.is_empty() {
7545 return Ok(rebuilt);
7546 }
7547 if let Some(array) = rebuilt.pop() {
7548 let Some(dictionary) = array.as_any_dictionary_opt() else {
7549 return Err(Error::invalid_input_source(
7550 "Cannot mix dictionary and non-dictionary chunks".into(),
7551 ));
7552 };
7553 let mut keys = Vec::with_capacity(pending_keys.len() + 1);
7554 keys.push(make_array(dictionary.keys().to_data()));
7555 keys.append(&mut pending_keys);
7556 rebuilt.push(rebuild(
7557 keys,
7558 array.data_type(),
7559 dictionary.values().clone(),
7560 )?);
7561 } else {
7562 let data_type = empty_data_type.ok_or_else(|| {
7563 Error::internal("Missing data type for an empty dictionary chunk")
7564 })?;
7565 let DataType::Dictionary(_, value_type) = &data_type else {
7566 return Err(Error::internal(format!(
7567 "Expected dictionary data type, got {data_type}"
7568 )));
7569 };
7570 rebuilt.push(rebuild(pending_keys, &data_type, zeroed(value_type, 1)?)?);
7571 }
7572 Ok(rebuilt)
7573 }
7574
7575 fn do_flush(
7577 &mut self,
7578 arrays: Vec<ArrayRef>,
7579 repdefs: Vec<RepDefBuilder>,
7580 row_number: u64,
7581 num_rows: u64,
7582 ) -> Result<Vec<EncodeTask>> {
7583 let arrays = Self::rebuild_empty_dictionary_chunks(arrays)?;
7584 DataBlock::validate_arrays(&arrays, &self.field.name)?;
7585 let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
7586 let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
7587 let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
7588 let normalized = RepDefBuilder::normalize(repdefs);
7589 let plan_ctx = PrimitivePlanContext {
7590 column_idx: self.column_index,
7591 field: &self.field,
7592 encoding_metadata: &self.encoding_metadata,
7593 };
7594 let mut pages = None;
7595 for page_encoding in self.page_encodings.iter() {
7596 if let Some(planned) = page_encoding.behavior.try_plan_pages(
7597 &plan_ctx,
7598 &arrays,
7599 &normalized,
7600 row_number,
7601 num_rows,
7602 num_values,
7603 )? {
7604 pages = Some(planned);
7605 break;
7606 }
7607 }
7608 let pages = pages.ok_or_else(|| {
7609 Error::invalid_input_source(
7610 format!(
7611 "No primitive page planner supports field '{}'",
7612 self.field.name
7613 )
7614 .into(),
7615 )
7616 })?;
7617
7618 let mut tasks = Vec::with_capacity(pages.len());
7619 let ctx = PrimitiveEncodeContext {
7620 column_idx: self.column_index,
7621 field: self.field.clone(),
7622 encoding_metadata: self.encoding_metadata.clone(),
7623 is_simple_validity,
7624 has_repdef_info,
7625 };
7626 for page in pages {
7627 let ctx = ctx.clone();
7628 let page_encodings = self.page_encodings.clone();
7629 let task =
7630 spawn_cpu(move || Self::encode_page(page_encodings.as_ref(), &ctx, page)).boxed();
7631 tasks.push(task);
7632 }
7633 Ok(tasks)
7634 }
7635
7636 fn extract_validity_buf(
7637 array: Arc<dyn Array>,
7638 repdef: &mut RepDefBuilder,
7639 keep_original_array: bool,
7640 ) -> Result<Arc<dyn Array>> {
7641 if let Some(validity) = array.nulls() {
7642 if keep_original_array {
7643 repdef.add_validity_bitmap(validity.clone());
7644 } else {
7645 repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
7646 }
7647 if let Some(dictionary) = array.as_any_dictionary_opt()
7651 && dictionary.values().is_empty()
7652 {
7653 return Ok(array);
7654 }
7655 let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
7656 Ok(make_array(data_no_nulls))
7657 } else {
7658 repdef.add_no_null(array.len());
7659 Ok(array)
7660 }
7661 }
7662
7663 fn extract_validity(
7664 mut array: Arc<dyn Array>,
7665 repdef: &mut RepDefBuilder,
7666 keep_original_array: bool,
7667 ) -> Result<Arc<dyn Array>> {
7668 match array.data_type() {
7669 DataType::Null => {
7670 repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
7671 Ok(array)
7672 }
7673 DataType::Dictionary(_, _) => {
7674 array = dict::normalize_dict_nulls(array)?;
7675 array = dict::clear_out_of_range_null_keys(array)?;
7676 Self::extract_validity_buf(array, repdef, keep_original_array)
7677 }
7678 _ => Self::extract_validity_buf(array, repdef, keep_original_array),
7687 }
7688 }
7689}
7690
7691impl PrimitivePageEncodingBehavior for RejectSparsePrimitiveEncoding {
7692 fn validate_field(&self, field: &Field, metadata: &HashMap<String, String>) -> Result<()> {
7693 if metadata
7694 .get(STRUCTURAL_ENCODING_META_KEY)
7695 .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE))
7696 {
7697 return Err(Error::invalid_input_source(
7698 format!(
7699 "Field '{}' requests sparse structural encoding, which is not enabled by the selected file format",
7700 field.name
7701 )
7702 .into(),
7703 ));
7704 }
7705 Ok(())
7706 }
7707}
7708
7709fn plan_dense_primitive_pages(
7710 arrays: &[ArrayRef],
7711 normalized: &NormalizedStructuralPlan,
7712 row_number: u64,
7713 num_rows: u64,
7714 num_values: u64,
7715) -> Result<Vec<PrimitivePageData>> {
7716 let (repdef, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
7717 miniblock::max_repdef_levels_per_chunk,
7718 num_rows,
7719 num_values,
7720 )?;
7721 PrimitiveStructuralEncoder::split_pages_for_miniblock_repdef_budget(
7722 arrays.to_vec(),
7723 repdef,
7724 miniblock_repdef_budget,
7725 row_number,
7726 num_rows,
7727 )
7728}
7729
7730impl PrimitivePageEncodingBehavior for DenseU16PrimitiveEncoding {
7731 fn try_plan_pages(
7732 &self,
7733 _ctx: &PrimitivePlanContext<'_>,
7734 arrays: &[ArrayRef],
7735 normalized: &NormalizedStructuralPlan,
7736 row_number: u64,
7737 num_rows: u64,
7738 num_values: u64,
7739 ) -> Result<Option<Vec<PrimitivePageData>>> {
7740 Ok(Some(plan_dense_primitive_pages(
7741 arrays, normalized, row_number, num_rows, num_values,
7742 )?))
7743 }
7744
7745 fn try_encode_page(
7746 &self,
7747 ctx: &PrimitiveEncodeContext,
7748 page: PrimitivePageData,
7749 ) -> Result<PrimitiveEncodeAttempt> {
7750 if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
7751 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7752 }
7753 Ok(PrimitiveEncodeAttempt::Encoded(
7754 PrimitiveStructuralEncoder::encode_dense_page(
7755 ctx.clone(),
7756 page,
7757 self.compression.clone(),
7758 MiniblockChunkSize::U16,
7759 ComplexNullEncoding::RawLevels,
7760 FixedWidthDictionaryEncoding::Exclude64Bit,
7761 )?,
7762 ))
7763 }
7764}
7765
7766impl PrimitivePageEncodingBehavior for DenseU32PrimitiveEncoding {
7767 fn try_plan_pages(
7768 &self,
7769 _ctx: &PrimitivePlanContext<'_>,
7770 arrays: &[ArrayRef],
7771 normalized: &NormalizedStructuralPlan,
7772 row_number: u64,
7773 num_rows: u64,
7774 num_values: u64,
7775 ) -> Result<Option<Vec<PrimitivePageData>>> {
7776 Ok(Some(plan_dense_primitive_pages(
7777 arrays, normalized, row_number, num_rows, num_values,
7778 )?))
7779 }
7780
7781 fn try_encode_page(
7782 &self,
7783 ctx: &PrimitiveEncodeContext,
7784 page: PrimitivePageData,
7785 ) -> Result<PrimitiveEncodeAttempt> {
7786 if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
7787 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7788 }
7789 Ok(PrimitiveEncodeAttempt::Encoded(
7790 PrimitiveStructuralEncoder::encode_dense_page(
7791 ctx.clone(),
7792 page,
7793 self.compression.clone(),
7794 MiniblockChunkSize::U32,
7795 ComplexNullEncoding::CompressedLevels,
7796 FixedWidthDictionaryEncoding::Include64Bit,
7797 )?,
7798 ))
7799 }
7800}
7801
7802impl PrimitivePageEncodingBehavior for SparsePrimitiveEncoding {
7803 fn try_plan_pages(
7804 &self,
7805 ctx: &PrimitivePlanContext<'_>,
7806 arrays: &[ArrayRef],
7807 normalized: &NormalizedStructuralPlan,
7808 row_number: u64,
7809 num_rows: u64,
7810 num_values: u64,
7811 ) -> Result<Option<Vec<PrimitivePageData>>> {
7812 let requested_encoding = ctx.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY);
7813 let requests_sparse = requested_encoding
7814 .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE));
7815 if requests_sparse {
7816 let plan = sparse::writer::plan(normalized, num_values)?;
7817 if sparse::writer::uses_constant_layout(&plan, ctx.field) {
7818 return Ok(None);
7819 }
7820 return Ok(Some(vec![PrimitivePageData {
7821 arrays: arrays.to_vec(),
7822 structure: PrimitivePageStructure::Sparse {
7823 plan,
7824 prepared_values: None,
7825 },
7826 row_number,
7827 num_rows,
7828 }]));
7829 }
7830
7831 let (_, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
7832 miniblock::max_repdef_levels_per_chunk,
7833 num_rows,
7834 num_values,
7835 )?;
7836 let automatic_sparse = layout::select_automatic_sparse(
7837 requested_encoding.map(String::as_str),
7838 &miniblock_repdef_budget,
7839 || {
7840 let data = DataBlock::from_arrays(arrays, num_values);
7841 if !sparse::writer::supports_value_block(&data) {
7842 return Ok(None);
7843 }
7844 let prepared_values = match sparse::writer::prepare_values(
7845 ctx.field,
7846 self.compression.as_ref(),
7847 data,
7848 MiniblockChunkSize::U32,
7849 ) {
7850 Ok(prepared_values) => prepared_values,
7851 Err(error) => {
7852 debug!(
7853 "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}",
7854 ctx.column_idx, error
7855 );
7856 return Ok(None);
7857 }
7858 };
7859 let plan = sparse::writer::plan(normalized, num_values)?;
7860 if sparse::writer::uses_constant_layout(&plan, ctx.field) {
7861 return Ok(None);
7862 }
7863 Ok(Some((plan, prepared_values)))
7864 },
7865 )?;
7866 Ok(automatic_sparse.map(|(plan, prepared_values)| {
7867 vec![PrimitivePageData {
7868 arrays: arrays.to_vec(),
7869 structure: PrimitivePageStructure::Sparse {
7870 plan,
7871 prepared_values: Some(prepared_values),
7872 },
7873 row_number,
7874 num_rows,
7875 }]
7876 }))
7877 }
7878
7879 fn try_encode_page(
7880 &self,
7881 ctx: &PrimitiveEncodeContext,
7882 page: PrimitivePageData,
7883 ) -> Result<PrimitiveEncodeAttempt> {
7884 if !matches!(&page.structure, PrimitivePageStructure::Sparse { .. }) {
7885 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7886 }
7887 let PrimitivePageData {
7888 arrays,
7889 structure:
7890 PrimitivePageStructure::Sparse {
7891 plan,
7892 prepared_values,
7893 },
7894 row_number,
7895 num_rows,
7896 } = page
7897 else {
7898 unreachable!()
7899 };
7900 let num_values = arrays.iter().map(|array| array.len() as u64).sum();
7901 log::debug!(
7902 "Encoding column {} with {} visible items ({} rows) using sparse layout",
7903 ctx.column_idx,
7904 num_values,
7905 num_rows
7906 );
7907 Ok(PrimitiveEncodeAttempt::Encoded(
7908 sparse::writer::encode_page(
7909 ctx.column_idx,
7910 &ctx.field,
7911 self.compression.as_ref(),
7912 prepared_values.map_or_else(
7913 || {
7914 sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays(
7915 &arrays, num_values,
7916 ))
7917 },
7918 sparse::writer::SparseValueInput::Prepared,
7919 ),
7920 plan,
7921 row_number,
7922 num_rows,
7923 MiniblockChunkSize::U32,
7924 )?,
7925 ))
7926 }
7927}
7928
7929impl PrimitivePageEncodingBehavior for ConstantPrimitiveEncoding {
7930 fn try_encode_page(
7931 &self,
7932 ctx: &PrimitiveEncodeContext,
7933 page: PrimitivePageData,
7934 ) -> Result<PrimitiveEncodeAttempt> {
7935 let PrimitivePageStructure::Dense { repdef, .. } = &page.structure else {
7936 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7937 };
7938 let num_values: u64 = page.arrays.iter().map(|array| array.len() as u64).sum();
7939 if num_values == 0 {
7940 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7941 }
7942 let leaf_validity = PrimitiveStructuralEncoder::leaf_validity(repdef, num_values as usize)?;
7943 if leaf_validity
7944 .as_ref()
7945 .is_some_and(|validity| validity.count_set_bits() == 0)
7946 || matches!(ctx.field.data_type(), DataType::Struct(fields) if fields.is_empty())
7947 {
7948 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7949 }
7950 let Some(scalar) =
7951 PrimitiveStructuralEncoder::find_constant_scalar(&page.arrays, leaf_validity.as_ref())?
7952 else {
7953 return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7954 };
7955 let PrimitivePageData {
7956 structure: PrimitivePageStructure::Dense { repdef, .. },
7957 row_number,
7958 num_rows,
7959 ..
7960 } = page
7961 else {
7962 unreachable!()
7963 };
7964 log::debug!(
7965 "Encoding column {} with {} items ({} rows) using constant layout",
7966 ctx.column_idx,
7967 num_values,
7968 num_rows
7969 );
7970 Ok(PrimitiveEncodeAttempt::Encoded(
7971 constant::encode_constant_page(ctx.column_idx, scalar, repdef, row_number, num_rows)?,
7972 ))
7973 }
7974}
7975
7976impl FieldEncoder for PrimitiveStructuralEncoder {
7977 fn maybe_encode(
7979 &mut self,
7980 array: ArrayRef,
7981 _external_buffers: &mut OutOfLineBuffers,
7982 mut repdef: RepDefBuilder,
7983 row_number: u64,
7984 num_rows: u64,
7985 ) -> Result<Vec<EncodeTask>> {
7986 let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
7987 self.accumulated_repdefs.push(repdef);
7988
7989 if let Some((arrays, row_number, num_rows)) =
7990 self.accumulation_queue.insert(array, row_number, num_rows)
7991 {
7992 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
7993 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
7994 } else {
7995 Ok(vec![])
7996 }
7997 }
7998
7999 fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
8001 if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
8002 let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
8003 Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
8004 } else {
8005 Ok(vec![])
8006 }
8007 }
8008
8009 fn num_columns(&self) -> u32 {
8010 1
8011 }
8012
8013 fn finish(
8014 &mut self,
8015 _external_buffers: &mut OutOfLineBuffers,
8016 ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
8017 std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
8018 }
8019}
8020
8021#[cfg(test)]
8022#[allow(clippy::single_range_in_vec_init)]
8023mod tests {
8024 use super::{
8025 CachedPage, CachedPageData, ChunkInstructions, DataBlock, DecodeMiniBlockTask,
8026 DecodePageTask, FixedFullZipDecodeTask, FixedPerValueDecompressor, FixedWidthDataBlock,
8027 FixedWidthDictionaryEncoding, FullZipCacheableState, FullZipDecodeDetails,
8028 FullZipDecodeTaskItem, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler,
8029 LazyLevels, LevelCodec, LevelCursor, LevelPlan, MiniBlockChunk, MiniBlockChunkIndex,
8030 MiniBlockCompressed, MiniblockChunkSize, PageCacheView, PageDataCacheKey,
8031 PageInfoAndScheduler, PageInitialization, PerValueDataBlock, PerValueDecompressor,
8032 PreambleAction, RunEndsBuilder, RunPosition, RunStorage, SimpleAllNullScheduler,
8033 StructuralPageScheduler, StructuralPrimitiveFieldScheduler, VariableFullZipDecoder,
8034 dense_levels_from_block, map_ranges_to_pages, pages_overlapping_ranges,
8035 validate_complex_all_null_levels,
8036 };
8037 use crate::EncodingsIo;
8038 use crate::buffer::LanceBuffer;
8039 use crate::compression::{
8040 BlockCompressor, DefaultDecompressionStrategy, MiniBlockDecompressor,
8041 };
8042 use crate::constants::{
8043 COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
8044 DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
8045 STRUCTURAL_ENCODING_MINIBLOCK,
8046 };
8047 use crate::data::BlockInfo;
8048 use crate::decoder::{
8049 FilterExpression, PageEncoding, RequestBatch, SchedulerContext, StructuralFieldDecoder,
8050 StructuralFieldScheduler,
8051 };
8052 use crate::encodings::logical::primitive::fullzip::PerValueCompressor;
8053 use crate::encodings::logical::primitive::{
8054 ChunkDrainInstructions, LoadedChunk, PrimitiveStructuralEncoder,
8055 StructuralPrimitiveFieldDecoder,
8056 };
8057 use crate::encodings::physical::rle::{RleDecompressor, RleEncoder, RleRuns, RunLengthWidth};
8058 use crate::encodings::physical::value::{ValueDecompressor, ValueEncoder};
8059 use crate::format::ProtobufUtils21;
8060 use crate::format::pb21;
8061 use crate::format::pb21::compressive_encoding::Compression;
8062 use crate::repdef::build_control_word_iterator;
8063 use crate::testing::TestEncoding;
8064 use crate::testing::{SimulatedScheduler, TestCases, check_round_trip_encoding_of_data};
8065 use arrow_array::{
8066 Array, ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int8Array,
8067 PrimitiveArray, StringArray, UInt8Array, make_array, new_null_array, types::Int32Type,
8068 };
8069 use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer};
8070 use arrow_schema::{DataType, Field as ArrowField};
8071 use std::collections::HashMap;
8072 use std::{collections::VecDeque, ops::Range, sync::Arc};
8073
8074 #[derive(Debug)]
8075 struct CacheMarker(u8);
8076
8077 impl lance_core::cache::DeepSizeOf for CacheMarker {
8078 fn deep_size_of_children(&self, _context: &mut lance_core::cache::Context) -> usize {
8079 0
8080 }
8081 }
8082
8083 impl CachedPageData for CacheMarker {
8084 fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync + 'static> {
8085 self
8086 }
8087 }
8088
8089 #[test]
8090 fn test_is_narrow() {
8091 let int8_array = Int8Array::from(vec![1, 2, 3]);
8092 let array_ref: ArrayRef = Arc::new(int8_array);
8093 let block = DataBlock::from_array(array_ref);
8094
8095 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
8096
8097 let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
8098 let block = DataBlock::from_array(string_array);
8099 assert!(PrimitiveStructuralEncoder::is_narrow(&block));
8100
8101 let string_array = StringArray::from(vec![
8102 Some("hello world".repeat(100)),
8103 Some("world".to_string()),
8104 ]);
8105 let block = DataBlock::from_array(string_array);
8106 assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
8107 }
8108
8109 fn valued_dictionary() -> ArrayRef {
8110 Arc::new(
8111 DictionaryArray::<Int32Type>::try_new(
8112 PrimitiveArray::<Int32Type>::from(vec![0]),
8113 Arc::new(StringArray::from(vec!["a"])),
8114 )
8115 .unwrap(),
8116 )
8117 }
8118
8119 fn empty_dictionary() -> ArrayRef {
8120 new_null_array(
8121 &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8122 1,
8123 )
8124 }
8125
8126 fn null_valued_dictionary() -> ArrayRef {
8127 Arc::new(
8128 DictionaryArray::<Int32Type>::try_new(
8129 PrimitiveArray::<Int32Type>::from(vec![0]),
8130 Arc::new(StringArray::from(vec![None::<&str>])),
8131 )
8132 .unwrap(),
8133 )
8134 }
8135
8136 fn sliced_empty_dictionary_with_hidden_key_payload() -> ArrayRef {
8137 let keys = PrimitiveArray::<Int32Type>::new(
8138 ScalarBuffer::from(vec![5, 7, 9]),
8139 Some(NullBuffer::new(BooleanBuffer::new_unset(3))),
8140 );
8141 let values = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
8142 let dictionary = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
8143 Arc::new(dictionary.slice(1, 1))
8144 }
8145
8146 #[rstest::rstest]
8147 #[case::empty_after_value(vec![valued_dictionary(), empty_dictionary()])]
8148 #[case::empty_before_value(vec![empty_dictionary(), valued_dictionary()])]
8149 #[case::null_value_after_value(vec![valued_dictionary(), null_valued_dictionary()])]
8150 #[case::hidden_sliced_after_value(vec![
8151 valued_dictionary(),
8152 sliced_empty_dictionary_with_hidden_key_payload(),
8153 ])]
8154 #[tokio::test]
8155 async fn test_mixed_valued_and_all_null_dictionary_chunks(#[case] dictionaries: Vec<ArrayRef>) {
8156 check_round_trip_encoding_of_data(
8157 dictionaries,
8158 &TestCases::default()
8159 .with_structural_encodings()
8160 .with_page_sizes(vec![1])
8161 .with_range(0..1)
8162 .with_range(1..2)
8163 .with_indices(vec![0])
8164 .with_indices(vec![1]),
8165 HashMap::new(),
8166 )
8167 .await;
8168 }
8169
8170 #[tokio::test]
8171 async fn test_sliced_empty_dictionary_with_hidden_key_payload() {
8172 check_round_trip_encoding_of_data(
8173 vec![sliced_empty_dictionary_with_hidden_key_payload()],
8174 &TestCases::default()
8175 .with_structural_encodings()
8176 .with_page_sizes(vec![4096]),
8177 HashMap::new(),
8178 )
8179 .await;
8180 }
8181
8182 #[test]
8183 fn test_primitive_decoder_empty_page_queue_returns_error() {
8184 let field = Arc::new(ArrowField::new("vector", DataType::Float32, true));
8185 let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8186
8187 let err = decoder.drain(1).unwrap_err();
8188 assert!(
8189 matches!(&err, lance_core::Error::Internal { .. }),
8190 "expected internal error, got: {err:?}"
8191 );
8192 let message = err.to_string();
8193 for expected in [
8194 "Primitive decoder missing page decoder",
8195 "field 'vector'",
8196 "data_type=Float32",
8197 "requested_rows=1",
8198 "remaining_rows=1",
8199 "rows_drained_in_current=0",
8200 "queued_pages=0",
8201 ] {
8202 assert!(
8203 message.contains(expected),
8204 "expected error to contain {expected:?}, got: {message}"
8205 );
8206 }
8207 }
8208
8209 #[test]
8210 fn test_plan_decoded_bytes_i32() {
8211 use crate::decoder::CANDIDATE_BATCH_SIZES;
8212
8213 let field = Arc::new(ArrowField::new("x", DataType::Int32, false));
8214 let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8215
8216 let rows_remaining = 100u64;
8217 let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap();
8218
8219 for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
8220 let rows = (candidate as u64).min(rows_remaining);
8221 assert_eq!(bytes[i], rows * 4, "candidate batch size {candidate}");
8222 }
8223 }
8224
8225 #[test]
8226 fn test_plan_decoded_bytes_i32_with_nulls() {
8227 use arrow_array::Int32Array;
8228
8229 const N: i32 = 1024;
8235 let array = Int32Array::from_iter((0..N).map(|i| if i % 2 == 0 { Some(i) } else { None }));
8236 let expected = array.get_buffer_memory_size() as u64;
8237
8238 let field = Arc::new(ArrowField::new("x", DataType::Int32, true));
8239 let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8240
8241 let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
8242 assert_eq!(
8244 bytes[5], expected,
8245 "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
8246 bytes[5],
8247 );
8248 }
8249
8250 #[test]
8251 fn test_plan_decoded_bytes_nullable_fsl_with_nullable_items() {
8252 use arrow_array::{FixedSizeListArray, Int32Array};
8261 use arrow_buffer::NullBuffer;
8262
8263 const N: usize = 1024;
8264 const DIM: i32 = 4;
8265
8266 let child_values = Int32Array::from_iter(
8267 (0..(N as i32 * DIM)).map(|i| if i % 2 == 0 { Some(i) } else { None }),
8268 );
8269 let item_field = Arc::new(ArrowField::new("item", DataType::Int32, true));
8270 let fsl_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::<Vec<_>>());
8271 let fsl_array = FixedSizeListArray::new(
8272 item_field.clone(),
8273 DIM,
8274 Arc::new(child_values),
8275 Some(fsl_nulls),
8276 );
8277 let expected = fsl_array.get_buffer_memory_size() as u64;
8278
8279 let fsl_type = DataType::FixedSizeList(item_field, DIM);
8280 let field = Arc::new(ArrowField::new("v", fsl_type, true));
8281 let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8282
8283 let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
8284 assert_eq!(
8286 bytes[5], expected,
8287 "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
8288 bytes[5],
8289 );
8290 }
8291
8292 #[test]
8293 fn test_plan_decoded_bytes_fixed_size_list_of_f32() {
8294 use crate::decoder::CANDIDATE_BATCH_SIZES;
8295
8296 const DIMENSION: i32 = 8;
8297 let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false));
8299 let fsl_type = DataType::FixedSizeList(item_field, DIMENSION);
8300 let field = Arc::new(ArrowField::new("vector", fsl_type, true));
8301 let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8302
8303 let rows_remaining = 100u64;
8304 let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap();
8305
8306 let bytes_per_row = DIMENSION as u64 * 4;
8309 for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
8310 let rows = (candidate as u64).min(rows_remaining);
8311 assert_eq!(
8312 bytes[i],
8313 rows * bytes_per_row + rows.div_ceil(8),
8314 "candidate batch size {candidate}"
8315 );
8316 }
8317 }
8318
8319 #[test]
8320 fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
8321 let fixed = FixedWidthDataBlock {
8322 data: LanceBuffer::from(vec![0_u8]),
8323 bits_per_value: 1,
8324 num_values: 8,
8325 block_info: BlockInfo::new(),
8326 };
8327 let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
8328
8329 let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
8330 else {
8331 panic!("expected full-zip to reject 1-bit fixed-width values");
8332 };
8333 assert!(
8334 err.to_string().contains("byte aligned"),
8335 "unexpected error: {err}"
8336 );
8337 }
8338
8339 fn decode_fixed_fullzip_no_levels(
8340 decompressor: Arc<dyn FixedPerValueDecompressor>,
8341 data: Vec<FullZipDecodeTaskItem>,
8342 num_rows: usize,
8343 bytes_per_value: usize,
8344 ) -> DataBlock {
8345 Box::new(FixedFullZipDecodeTask {
8346 details: Arc::new(FullZipDecodeDetails {
8347 value_decompressor: PerValueDecompressor::Fixed(decompressor),
8348 def_meaning: Arc::from([]),
8349 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 0),
8350 max_rep: 0,
8351 max_visible_def: u16::MAX,
8352 }),
8353 data,
8354 num_rows,
8355 bytes_per_value,
8356 })
8357 .decode()
8358 .unwrap()
8359 .data
8360 }
8361
8362 #[test]
8363 fn test_fixed_fullzip_decode_preallocates_exact_output_size() {
8364 #[derive(Debug)]
8365 struct IdentityFixedDecompressor;
8366
8367 impl FixedPerValueDecompressor for IdentityFixedDecompressor {
8368 fn decompress(
8369 &self,
8370 data: FixedWidthDataBlock,
8371 num_rows: u64,
8372 ) -> crate::Result<DataBlock> {
8373 assert_eq!(data.num_values, num_rows);
8374 Ok(DataBlock::FixedWidth(data))
8375 }
8376
8377 fn bits_per_value(&self) -> u64 {
8378 32
8379 }
8380
8381 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8382 num_values.checked_mul(4)
8383 }
8384 }
8385
8386 let make_item = |num_rows: u64| FullZipDecodeTaskItem {
8387 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
8388 data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]),
8389 bits_per_value: 32,
8390 num_values: num_rows,
8391 block_info: BlockInfo::new(),
8392 }),
8393 rows_in_buf: num_rows,
8394 };
8395
8396 let num_rows = 512;
8397 let decoded = decode_fixed_fullzip_no_levels(
8398 Arc::new(IdentityFixedDecompressor),
8399 vec![make_item(128), make_item(384)],
8400 num_rows,
8401 4,
8402 );
8403 let values = decoded.as_fixed_width_ref().unwrap();
8404 let expected_size = num_rows * 4;
8405 assert_eq!(values.data.len(), expected_size);
8406 assert_eq!(values.data.clone().into_buffer().capacity(), expected_size);
8407 }
8408
8409 #[test]
8410 fn test_fixed_fullzip_decode_falls_back_when_output_size_is_not_exact() {
8411 #[derive(Debug)]
8412 struct FallbackFixedDecompressor;
8413
8414 impl FixedPerValueDecompressor for FallbackFixedDecompressor {
8415 fn decompress(
8416 &self,
8417 data: FixedWidthDataBlock,
8418 num_rows: u64,
8419 ) -> crate::Result<DataBlock> {
8420 assert_eq!(data.num_values, num_rows);
8421 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
8422 data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]),
8423 bits_per_value: 32,
8424 num_values: num_rows,
8425 block_info: BlockInfo::new(),
8426 }))
8427 }
8428
8429 fn bits_per_value(&self) -> u64 {
8430 u64::MAX - 7
8433 }
8434 }
8435
8436 let num_rows = 2;
8437 let decoded = decode_fixed_fullzip_no_levels(
8438 Arc::new(FallbackFixedDecompressor),
8439 vec![FullZipDecodeTaskItem {
8440 data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
8441 data: LanceBuffer::from(vec![0_u8; num_rows * 4]),
8442 bits_per_value: 32,
8443 num_values: num_rows as u64,
8444 block_info: BlockInfo::new(),
8445 }),
8446 rows_in_buf: num_rows as u64,
8447 }],
8448 num_rows,
8449 4,
8450 );
8451 let values = decoded.as_fixed_width_ref().unwrap();
8452 assert_eq!(values.num_values, num_rows as u64);
8453 assert_eq!(values.data.len(), num_rows * 4);
8454 }
8455
8456 #[test]
8457 fn test_fixed_fullzip_real_fsl_preallocates_exact_output_size() {
8458 let num_rows = 64;
8459 let dimension = 32;
8460 let items = Arc::new(Float32Array::from_iter_values(
8461 (0..num_rows * dimension).map(|value| value as f32),
8462 ));
8463 let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false));
8464 let sample = FixedSizeListArray::new(item_field, dimension as i32, items, None);
8465
8466 let (data, compression) = PerValueCompressor::compress(
8467 &ValueEncoder::default(),
8468 DataBlock::from_array(sample.clone()),
8469 )
8470 .unwrap();
8471 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
8472 panic!("expected fixed-size-list compression");
8473 };
8474 let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
8475 let expected_size = num_rows * dimension * size_of::<f32>();
8476 assert_eq!(
8477 FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64),
8478 Some(expected_size as u64)
8479 );
8480
8481 let decoded = decode_fixed_fullzip_no_levels(
8482 Arc::new(decompressor),
8483 vec![FullZipDecodeTaskItem {
8484 data,
8485 rows_in_buf: num_rows as u64,
8486 }],
8487 num_rows,
8488 dimension * size_of::<f32>(),
8489 );
8490 let fsl = decoded.as_fixed_size_list_ref().unwrap();
8491 let values = fsl.child.as_fixed_width_ref().unwrap();
8492 assert_eq!(values.data.len(), expected_size);
8493 assert_eq!(values.data.clone().into_buffer().capacity(), expected_size);
8494
8495 let decoded_array = make_array(
8496 decoded
8497 .into_arrow(sample.data_type().clone(), true)
8498 .unwrap(),
8499 );
8500 assert_eq!(decoded_array.as_ref(), &sample);
8501 }
8502
8503 #[test]
8504 fn test_fixed_fullzip_nullable_fsl_uses_fallback_end_to_end() {
8505 #[derive(Debug)]
8506 struct NullableFslDecompressor {
8507 inner: ValueDecompressor,
8508 }
8509
8510 impl FixedPerValueDecompressor for NullableFslDecompressor {
8511 fn decompress(
8512 &self,
8513 data: FixedWidthDataBlock,
8514 num_rows: u64,
8515 ) -> crate::Result<DataBlock> {
8516 FixedPerValueDecompressor::decompress(&self.inner, data, num_rows)
8517 }
8518
8519 fn bits_per_value(&self) -> u64 {
8520 u64::MAX - 7
8523 }
8524
8525 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8526 FixedPerValueDecompressor::decoded_size_bytes(&self.inner, num_values)
8527 }
8528 }
8529
8530 let num_rows = 64;
8531 let items = Arc::new(UInt8Array::from_iter(
8532 (0..num_rows).map(|value| (value % 3 != 0).then_some(value as u8)),
8533 ));
8534 let item_field = Arc::new(ArrowField::new("item", DataType::UInt8, true));
8535 let sample = FixedSizeListArray::new(item_field, 1, items, None);
8536
8537 let (data, compression) = PerValueCompressor::compress(
8538 &ValueEncoder::default(),
8539 DataBlock::from_array(sample.clone()),
8540 )
8541 .unwrap();
8542 let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
8543 panic!("expected fixed-size-list compression");
8544 };
8545 let decompressor = NullableFslDecompressor {
8546 inner: ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(),
8547 };
8548 assert_eq!(
8549 FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64),
8550 None
8551 );
8552
8553 let decoded = decode_fixed_fullzip_no_levels(
8554 Arc::new(decompressor),
8555 vec![FullZipDecodeTaskItem {
8556 data,
8557 rows_in_buf: num_rows as u64,
8558 }],
8559 num_rows,
8560 2,
8561 );
8562 let decoded_array = make_array(
8563 decoded
8564 .into_arrow(sample.data_type().clone(), true)
8565 .unwrap(),
8566 );
8567 assert_eq!(decoded_array.as_ref(), &sample);
8568 }
8569
8570 #[test]
8571 fn test_miniblock_decode_uses_exact_fixed_width_output_size() {
8572 #[derive(Debug)]
8573 struct FixedWidthMiniBlockDecompressor;
8574
8575 impl MiniBlockDecompressor for FixedWidthMiniBlockDecompressor {
8576 fn decompress(
8577 &self,
8578 data: Vec<LanceBuffer>,
8579 num_values: u64,
8580 ) -> crate::Result<DataBlock> {
8581 assert_eq!(data.len(), 1);
8582 Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
8583 data: data.into_iter().next().unwrap(),
8584 bits_per_value: 32,
8585 num_values,
8586 block_info: BlockInfo::new(),
8587 }))
8588 }
8589
8590 fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8591 num_values.checked_mul(4)
8592 }
8593 }
8594
8595 let num_rows = 512;
8596 let expected_size = num_rows * 4;
8597 let mut chunk_data = Vec::new();
8598 chunk_data.extend_from_slice(&0_u16.to_le_bytes());
8599 chunk_data.extend_from_slice(&(expected_size as u16).to_le_bytes());
8600 let header_padding =
8601 lance_core::utils::bit::pad_bytes::<{ super::MINIBLOCK_ALIGNMENT }>(chunk_data.len());
8602 chunk_data.resize(chunk_data.len() + header_padding, 0);
8603 chunk_data.resize(chunk_data.len() + expected_size as usize, 7);
8604
8605 let task = DecodeMiniBlockTask {
8606 rep_decompressor: None,
8607 def_decompressor: None,
8608 value_decompressor: Arc::new(FixedWidthMiniBlockDecompressor),
8609 dictionary_data: None,
8610 def_meaning: Arc::from([]),
8611 num_buffers: 1,
8612 max_visible_level: 0,
8613 instructions: vec![(
8614 ChunkDrainInstructions {
8615 chunk_instructions: ChunkInstructions {
8616 chunk_idx: 0,
8617 preamble: PreambleAction::Absent,
8618 rows_to_skip: 0,
8619 rows_to_take: num_rows,
8620 take_trailer: false,
8621 },
8622 rows_to_skip: 0,
8623 rows_to_take: num_rows,
8624 preamble_action: PreambleAction::Absent,
8625 },
8626 LoadedChunk {
8627 byte_range: 0..chunk_data.len() as u64,
8628 data: LanceBuffer::from(chunk_data),
8629 items_in_chunk: num_rows,
8630 chunk_idx: 0,
8631 },
8632 )],
8633 has_large_chunk: false,
8634 };
8635
8636 let decoded = Box::new(task).decode().unwrap();
8637 let values = decoded.data.as_fixed_width_ref().unwrap();
8638 assert_eq!(values.data.len(), expected_size as usize);
8639 assert_eq!(
8640 values.data.clone().into_buffer().capacity(),
8641 expected_size as usize
8642 );
8643 }
8644
8645 #[test]
8646 fn test_map_range() {
8647 let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
8650 let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
8651 let max_visible_def = 0;
8652 let total_items = 8;
8653 let max_rep = 1;
8654
8655 let check = |range, expected_item_range, expected_level_range| {
8656 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8657 range,
8658 rep.as_ref(),
8659 def.as_ref(),
8660 max_rep,
8661 max_visible_def,
8662 total_items,
8663 PreambleAction::Absent,
8664 );
8665 assert_eq!(item_range, expected_item_range);
8666 assert_eq!(level_range, expected_level_range);
8667 };
8668
8669 check(0..1, 0..3, 0..3);
8670 check(1..2, 3..5, 3..5);
8671 check(2..3, 5..5, 5..6);
8672 check(3..4, 5..8, 6..9);
8673 check(0..2, 0..5, 0..5);
8674 check(1..3, 3..5, 3..6);
8675 check(2..4, 5..8, 5..9);
8676 check(0..3, 0..5, 0..6);
8677 check(1..4, 3..8, 3..9);
8678 check(0..4, 0..8, 0..9);
8679
8680 let rep = Some(vec![1, 1, 0, 1]);
8683 let def = Some(vec![1, 0, 0, 0]);
8684 let max_visible_def = 0;
8685 let total_items = 3;
8686
8687 let check = |range, expected_item_range, expected_level_range| {
8688 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8689 range,
8690 rep.as_ref(),
8691 def.as_ref(),
8692 max_rep,
8693 max_visible_def,
8694 total_items,
8695 PreambleAction::Absent,
8696 );
8697 assert_eq!(item_range, expected_item_range);
8698 assert_eq!(level_range, expected_level_range);
8699 };
8700
8701 check(0..1, 0..0, 0..1);
8702 check(1..2, 0..2, 1..3);
8703 check(2..3, 2..3, 3..4);
8704 check(0..2, 0..2, 0..3);
8705 check(1..3, 0..3, 1..4);
8706 check(0..3, 0..3, 0..4);
8707
8708 let rep = Some(vec![1, 1, 0, 1]);
8711 let def = Some(vec![0, 0, 0, 1]);
8712 let max_visible_def = 0;
8713 let total_items = 3;
8714
8715 let check = |range, expected_item_range, expected_level_range| {
8716 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8717 range,
8718 rep.as_ref(),
8719 def.as_ref(),
8720 max_rep,
8721 max_visible_def,
8722 total_items,
8723 PreambleAction::Absent,
8724 );
8725 assert_eq!(item_range, expected_item_range);
8726 assert_eq!(level_range, expected_level_range);
8727 };
8728
8729 check(0..1, 0..1, 0..1);
8730 check(1..2, 1..3, 1..3);
8731 check(2..3, 3..3, 3..4);
8732 check(0..2, 0..3, 0..3);
8733 check(1..3, 1..3, 1..4);
8734 check(0..3, 0..3, 0..4);
8735
8736 let rep = Some(vec![1, 0, 1, 0, 1, 0]);
8739 let def: Option<&[u16]> = None;
8740 let max_visible_def = 0;
8741 let total_items = 6;
8742
8743 let check = |range, expected_item_range, expected_level_range| {
8744 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8745 range,
8746 rep.as_ref(),
8747 def.as_ref(),
8748 max_rep,
8749 max_visible_def,
8750 total_items,
8751 PreambleAction::Absent,
8752 );
8753 assert_eq!(item_range, expected_item_range);
8754 assert_eq!(level_range, expected_level_range);
8755 };
8756
8757 check(0..1, 0..2, 0..2);
8758 check(1..2, 2..4, 2..4);
8759 check(2..3, 4..6, 4..6);
8760 check(0..2, 0..4, 0..4);
8761 check(1..3, 2..6, 2..6);
8762 check(0..3, 0..6, 0..6);
8763
8764 let rep: Option<&[u16]> = None;
8767 let def = Some(vec![0, 0, 1, 0]);
8768 let max_visible_def = 1;
8769 let total_items = 4;
8770
8771 let check = |range, expected_item_range, expected_level_range| {
8772 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8773 range,
8774 rep.as_ref(),
8775 def.as_ref(),
8776 max_rep,
8777 max_visible_def,
8778 total_items,
8779 PreambleAction::Absent,
8780 );
8781 assert_eq!(item_range, expected_item_range);
8782 assert_eq!(level_range, expected_level_range);
8783 };
8784
8785 check(0..1, 0..1, 0..1);
8786 check(1..2, 1..2, 1..2);
8787 check(2..3, 2..3, 2..3);
8788 check(0..2, 0..2, 0..2);
8789 check(1..3, 1..3, 1..3);
8790 check(0..3, 0..3, 0..3);
8791
8792 let rep = Some(vec![0, 1, 0, 1]);
8797 let def = Some(vec![0, 0, 0, 1]);
8798 let max_visible_def = 0;
8799 let total_items = 3;
8800
8801 let check = |range, expected_item_range, expected_level_range| {
8802 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8803 range,
8804 rep.as_ref(),
8805 def.as_ref(),
8806 max_rep,
8807 max_visible_def,
8808 total_items,
8809 PreambleAction::Take,
8810 );
8811 assert_eq!(item_range, expected_item_range);
8812 assert_eq!(level_range, expected_level_range);
8813 };
8814
8815 check(0..1, 0..3, 0..3);
8817 check(0..2, 0..3, 0..4);
8818
8819 let check = |range, expected_item_range, expected_level_range| {
8820 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8821 range,
8822 rep.as_ref(),
8823 def.as_ref(),
8824 max_rep,
8825 max_visible_def,
8826 total_items,
8827 PreambleAction::Skip,
8828 );
8829 assert_eq!(item_range, expected_item_range);
8830 assert_eq!(level_range, expected_level_range);
8831 };
8832
8833 check(0..1, 1..3, 1..3);
8834 check(1..2, 3..3, 3..4);
8835 check(0..2, 1..3, 1..4);
8836
8837 let rep = Some(vec![0, 1, 1, 0]);
8842 let def = Some(vec![0, 1, 0, 0]);
8843 let max_visible_def = 0;
8844 let total_items = 4;
8845
8846 let check = |range, expected_item_range, expected_level_range| {
8847 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8848 range,
8849 rep.as_ref(),
8850 def.as_ref(),
8851 max_rep,
8852 max_visible_def,
8853 total_items,
8854 PreambleAction::Take,
8855 );
8856 assert_eq!(item_range, expected_item_range);
8857 assert_eq!(level_range, expected_level_range);
8858 };
8859
8860 check(0..1, 0..1, 0..2);
8862 check(0..2, 0..3, 0..4);
8863
8864 let check = |range, expected_item_range, expected_level_range| {
8865 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8866 range,
8867 rep.as_ref(),
8868 def.as_ref(),
8869 max_rep,
8870 max_visible_def,
8871 total_items,
8872 PreambleAction::Skip,
8873 );
8874 assert_eq!(item_range, expected_item_range);
8875 assert_eq!(level_range, expected_level_range);
8876 };
8877
8878 check(0..1, 1..1, 1..2);
8880 check(1..2, 1..3, 2..4);
8881 check(0..2, 1..3, 1..4);
8882
8883 let rep = Some(vec![0, 1, 0, 1]);
8886 let def: Option<Vec<u16>> = None;
8887 let max_visible_def = 0;
8888 let total_items = 4;
8889
8890 let check = |range, expected_item_range, expected_level_range| {
8891 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8892 range,
8893 rep.as_ref(),
8894 def.as_ref(),
8895 max_rep,
8896 max_visible_def,
8897 total_items,
8898 PreambleAction::Take,
8899 );
8900 assert_eq!(item_range, expected_item_range);
8901 assert_eq!(level_range, expected_level_range);
8902 };
8903
8904 check(0..1, 0..3, 0..3);
8906 check(0..2, 0..4, 0..4);
8907
8908 let check = |range, expected_item_range, expected_level_range| {
8909 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8910 range,
8911 rep.as_ref(),
8912 def.as_ref(),
8913 max_rep,
8914 max_visible_def,
8915 total_items,
8916 PreambleAction::Skip,
8917 );
8918 assert_eq!(item_range, expected_item_range);
8919 assert_eq!(level_range, expected_level_range);
8920 };
8921
8922 check(0..1, 1..3, 1..3);
8923 check(1..2, 3..4, 3..4);
8924 check(0..2, 1..4, 1..4);
8925
8926 let rep = Some(vec![2, 1, 2, 0, 1, 2]);
8930 let def = Some(vec![0, 1, 2, 0, 0, 0]);
8931 let max_rep = 2;
8932 let max_visible_def = 0;
8933 let total_items = 4;
8934
8935 let check = |range, expected_item_range, expected_level_range| {
8936 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8937 range,
8938 rep.as_ref(),
8939 def.as_ref(),
8940 max_rep,
8941 max_visible_def,
8942 total_items,
8943 PreambleAction::Absent,
8944 );
8945 assert_eq!(item_range, expected_item_range);
8946 assert_eq!(level_range, expected_level_range);
8947 };
8948
8949 check(0..3, 0..4, 0..6);
8950 check(0..1, 0..1, 0..2);
8951 check(1..2, 1..3, 2..5);
8952 check(2..3, 3..4, 5..6);
8953
8954 let rep = Some(vec![0, 0, 1, 0, 1, 1]);
8956 let def = Some(vec![0, 1, 0, 0, 0, 0]);
8957 let max_rep = 1;
8958 let max_visible_def = 0;
8959 let total_items = 5;
8960
8961 let check = |range, expected_item_range, expected_level_range| {
8962 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8963 range,
8964 rep.as_ref(),
8965 def.as_ref(),
8966 max_rep,
8967 max_visible_def,
8968 total_items,
8969 PreambleAction::Take,
8970 );
8971 assert_eq!(item_range, expected_item_range);
8972 assert_eq!(level_range, expected_level_range);
8973 };
8974
8975 check(0..0, 0..1, 0..2);
8976 check(0..1, 0..3, 0..4);
8977 check(0..2, 0..4, 0..5);
8978
8979 let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
8982 let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
8983 let max_rep = 1;
8984 let max_visible_def = 0;
8985 let total_items = 5;
8986
8987 let check = |range, expected_item_range, expected_level_range| {
8988 let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8989 range,
8990 rep.as_ref(),
8991 def.as_ref(),
8992 max_rep,
8993 max_visible_def,
8994 total_items,
8995 PreambleAction::Skip,
8996 );
8997 assert_eq!(item_range, expected_item_range);
8998 assert_eq!(level_range, expected_level_range);
8999 };
9000
9001 check(2..3, 2..4, 5..7);
9002 }
9003
9004 #[test]
9005 fn test_slice_batch_data_and_rebase_offsets_u32() {
9006 let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
9007 let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
9008
9009 let (sliced_data, normalized_offsets) =
9010 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
9011 .unwrap();
9012
9013 assert_eq!(sliced_data.as_ref(), b"6789ab");
9014 let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
9015 assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
9016 }
9017
9018 #[test]
9019 fn test_slice_batch_data_and_rebase_offsets_u64() {
9020 let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
9021 let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
9022
9023 let (sliced_data, normalized_offsets) =
9024 VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
9025 .unwrap();
9026
9027 assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
9028 let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
9029 assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
9030 }
9031
9032 #[test]
9033 fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
9034 let data = LanceBuffer::copy_slice(b"abcd");
9035 let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
9036
9037 let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
9038 .expect_err("offset end before start should error");
9039 assert!(err.to_string().contains("less than base"));
9040 }
9041
9042 #[test]
9043 fn test_schedule_instructions() {
9044 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
9046 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9047 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9048
9049 let check = |user_ranges, expected_instructions| {
9050 let instructions = ChunkInstructions::schedule_instructions(&chunk_index, user_ranges);
9051 assert_eq!(instructions, expected_instructions);
9052 };
9053
9054 let expected_take_all = vec![
9056 ChunkInstructions {
9057 chunk_idx: 0,
9058 preamble: PreambleAction::Absent,
9059 rows_to_skip: 0,
9060 rows_to_take: 6,
9061 take_trailer: true,
9062 },
9063 ChunkInstructions {
9064 chunk_idx: 1,
9065 preamble: PreambleAction::Take,
9066 rows_to_skip: 0,
9067 rows_to_take: 2,
9068 take_trailer: false,
9069 },
9070 ChunkInstructions {
9071 chunk_idx: 2,
9072 preamble: PreambleAction::Absent,
9073 rows_to_skip: 0,
9074 rows_to_take: 5,
9075 take_trailer: true,
9076 },
9077 ChunkInstructions {
9078 chunk_idx: 3,
9079 preamble: PreambleAction::Take,
9080 rows_to_skip: 0,
9081 rows_to_take: 1,
9082 take_trailer: false,
9083 },
9084 ];
9085
9086 check(&[0..14], expected_take_all.clone());
9088
9089 check(
9091 &[
9092 0..1,
9093 1..2,
9094 2..3,
9095 3..4,
9096 4..5,
9097 5..6,
9098 6..7,
9099 7..8,
9100 8..9,
9101 9..10,
9102 10..11,
9103 11..12,
9104 12..13,
9105 13..14,
9106 ],
9107 expected_take_all,
9108 );
9109
9110 check(
9114 &[0..1, 3..4],
9115 vec![
9116 ChunkInstructions {
9117 chunk_idx: 0,
9118 preamble: PreambleAction::Absent,
9119 rows_to_skip: 0,
9120 rows_to_take: 1,
9121 take_trailer: false,
9122 },
9123 ChunkInstructions {
9124 chunk_idx: 0,
9125 preamble: PreambleAction::Absent,
9126 rows_to_skip: 3,
9127 rows_to_take: 1,
9128 take_trailer: false,
9129 },
9130 ],
9131 );
9132
9133 check(
9135 &[5..6],
9136 vec![
9137 ChunkInstructions {
9138 chunk_idx: 0,
9139 preamble: PreambleAction::Absent,
9140 rows_to_skip: 5,
9141 rows_to_take: 1,
9142 take_trailer: true,
9143 },
9144 ChunkInstructions {
9145 chunk_idx: 1,
9146 preamble: PreambleAction::Take,
9147 rows_to_skip: 0,
9148 rows_to_take: 0,
9149 take_trailer: false,
9150 },
9151 ],
9152 );
9153
9154 check(
9156 &[7..10],
9157 vec![
9158 ChunkInstructions {
9159 chunk_idx: 1,
9160 preamble: PreambleAction::Skip,
9161 rows_to_skip: 1,
9162 rows_to_take: 1,
9163 take_trailer: false,
9164 },
9165 ChunkInstructions {
9166 chunk_idx: 2,
9167 preamble: PreambleAction::Absent,
9168 rows_to_skip: 0,
9169 rows_to_take: 2,
9170 take_trailer: false,
9171 },
9172 ],
9173 );
9174 }
9175
9176 #[test]
9177 fn test_drain_instructions() {
9178 fn drain_from_instructions(
9179 instructions: &mut VecDeque<ChunkInstructions>,
9180 mut rows_desired: u64,
9181 need_preamble: &mut bool,
9182 skip_in_chunk: &mut u64,
9183 ) -> Vec<ChunkDrainInstructions> {
9184 let mut drain_instructions = Vec::with_capacity(instructions.len());
9186 while rows_desired > 0 || *need_preamble {
9187 let (next_instructions, consumed_chunk) = instructions
9188 .front()
9189 .unwrap()
9190 .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
9191 if consumed_chunk {
9192 instructions.pop_front();
9193 }
9194 drain_instructions.push(next_instructions);
9195 }
9196 drain_instructions
9197 }
9198
9199 let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
9201 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9202 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9203 let user_ranges = vec![1..7, 10..14];
9204
9205 let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
9207
9208 let mut to_drain = VecDeque::from(scheduled.clone());
9209
9210 let mut need_preamble = false;
9213 let mut skip_in_chunk = 0;
9214
9215 let next_batch =
9216 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
9217
9218 assert!(!need_preamble);
9219 assert_eq!(skip_in_chunk, 4);
9220 assert_eq!(
9221 next_batch,
9222 vec![ChunkDrainInstructions {
9223 chunk_instructions: scheduled[0].clone(),
9224 rows_to_take: 4,
9225 rows_to_skip: 0,
9226 preamble_action: PreambleAction::Absent,
9227 }]
9228 );
9229
9230 let next_batch =
9231 drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
9232
9233 assert!(!need_preamble);
9234 assert_eq!(skip_in_chunk, 2);
9235
9236 assert_eq!(
9237 next_batch,
9238 vec![
9239 ChunkDrainInstructions {
9240 chunk_instructions: scheduled[0].clone(),
9241 rows_to_take: 1,
9242 rows_to_skip: 4,
9243 preamble_action: PreambleAction::Absent,
9244 },
9245 ChunkDrainInstructions {
9246 chunk_instructions: scheduled[1].clone(),
9247 rows_to_take: 1,
9248 rows_to_skip: 0,
9249 preamble_action: PreambleAction::Take,
9250 },
9251 ChunkDrainInstructions {
9252 chunk_instructions: scheduled[2].clone(),
9253 rows_to_take: 2,
9254 rows_to_skip: 0,
9255 preamble_action: PreambleAction::Absent,
9256 }
9257 ]
9258 );
9259
9260 let next_batch =
9261 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
9262
9263 assert!(!need_preamble);
9264 assert_eq!(skip_in_chunk, 0);
9265
9266 assert_eq!(
9267 next_batch,
9268 vec![
9269 ChunkDrainInstructions {
9270 chunk_instructions: scheduled[2].clone(),
9271 rows_to_take: 1,
9272 rows_to_skip: 2,
9273 preamble_action: PreambleAction::Absent,
9274 },
9275 ChunkDrainInstructions {
9276 chunk_instructions: scheduled[3].clone(),
9277 rows_to_take: 1,
9278 rows_to_skip: 0,
9279 preamble_action: PreambleAction::Take,
9280 },
9281 ]
9282 );
9283
9284 let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
9286 let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9287 let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9288 let user_ranges = vec![0..28];
9289
9290 let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
9292
9293 let mut to_drain = VecDeque::from(scheduled.clone());
9294
9295 let mut need_preamble = false;
9298 let mut skip_in_chunk = 0;
9299
9300 let next_batch =
9301 drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
9302
9303 assert_eq!(
9304 next_batch,
9305 vec![
9306 ChunkDrainInstructions {
9307 chunk_instructions: scheduled[0].clone(),
9308 rows_to_take: 6,
9309 rows_to_skip: 0,
9310 preamble_action: PreambleAction::Absent,
9311 },
9312 ChunkDrainInstructions {
9313 chunk_instructions: scheduled[1].clone(),
9314 rows_to_take: 1,
9315 rows_to_skip: 0,
9316 preamble_action: PreambleAction::Take,
9317 },
9318 ]
9319 );
9320
9321 assert!(!need_preamble);
9322 assert_eq!(skip_in_chunk, 1);
9323
9324 let next_batch =
9327 drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
9328
9329 assert_eq!(
9330 next_batch,
9331 vec![
9332 ChunkDrainInstructions {
9333 chunk_instructions: scheduled[1].clone(),
9334 rows_to_take: 2,
9335 rows_to_skip: 1,
9336 preamble_action: PreambleAction::Skip,
9337 },
9338 ChunkDrainInstructions {
9339 chunk_instructions: scheduled[2].clone(),
9340 rows_to_take: 0,
9341 rows_to_skip: 0,
9342 preamble_action: PreambleAction::Take,
9343 },
9344 ]
9345 );
9346
9347 assert!(!need_preamble);
9348 assert_eq!(skip_in_chunk, 0);
9349 }
9350
9351 use super::chunk_index::{PrefixSums, RowMapping};
9352 use super::{MINIBLOCK_ALIGNMENT, Words, build_chunk_index};
9353 use bytes::Bytes;
9354 use lance_core::cache::{Context, DeepSizeOf};
9355 use rstest::rstest;
9356
9357 fn words_from(entries: &[(u32, u32)]) -> (Words, u64) {
9360 let mut raw = Vec::with_capacity(entries.len() * 2);
9361 let mut total = 0u64;
9362 for &(log, num_bytes) in entries {
9363 assert!(num_bytes > 0 && num_bytes % MINIBLOCK_ALIGNMENT as u32 == 0);
9364 let divided = num_bytes / MINIBLOCK_ALIGNMENT as u32 - 1;
9365 let word = (divided << 4) | log;
9366 assert!(word <= u16::MAX as u32, "test word {word} exceeds u16");
9367 raw.extend_from_slice(&(word as u16).to_le_bytes());
9368 total += num_bytes as u64;
9369 }
9370 (Words::from_bytes(Bytes::from(raw), false).unwrap(), total)
9371 }
9372
9373 fn rep_bytes_from(values: &[u64]) -> Vec<u8> {
9374 values.iter().flat_map(|v| v.to_le_bytes()).collect()
9375 }
9376
9377 #[rstest]
9378 #[case::uniform_partial_last(&[(3, 16), (3, 24), (0, 8)], 19, "uniform_flat", 8, 3)]
9381 #[case::single_chunk(&[(0, 24)], 5, "uniform_flat", 5, 5)]
9383 #[case::exact_multiple(&[(3, 16), (3, 16)], 16, "uniform_flat", 8, 8)]
9385 #[case::non_uniform(&[(4, 16), (2, 16), (0, 8)], 21, "flat", 16, 1)]
9387 fn test_flat_detection(
9388 #[case] entries: &[(u32, u32)],
9389 #[case] items_in_page: u64,
9390 #[case] expected_kind: &str,
9391 #[case] expected_first_items: u64,
9392 #[case] expected_last_items: u64,
9393 ) {
9394 let base = 100u64;
9395 let (words, data_buf_size) = words_from(entries);
9396 let index = build_chunk_index(&words, items_in_page, base, data_buf_size, None, 0).unwrap();
9397
9398 assert_eq!(index.row_mapping_debug(), expected_kind);
9399 assert_eq!(index.num_chunks(), entries.len());
9400 assert_eq!(index.items_in_chunk(0), expected_first_items);
9401 assert_eq!(index.items_in_chunk(entries.len() - 1), expected_last_items);
9402
9403 let mut expected_start = base;
9405 for (i, &(_, num_bytes)) in entries.iter().enumerate() {
9406 let range = index.byte_range(i);
9407 assert_eq!(range.start, expected_start);
9408 assert_eq!(range.end - range.start, num_bytes as u64);
9409 expected_start = range.end;
9410 }
9411 assert_eq!(expected_start, base + data_buf_size);
9412
9413 let total_items: u64 = (0..index.num_chunks())
9415 .map(|i| index.items_in_chunk(i))
9416 .sum();
9417 assert_eq!(total_items, items_in_page);
9418 }
9419
9420 #[test]
9421 fn test_nested_detection_and_axes() {
9422 let rep = rep_bytes_from(&[5, 0, 4, 0, 3, 0]);
9424
9425 let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9427 let index = build_chunk_index(&words, 10, 0, data_buf_size, Some(&rep), 1).unwrap();
9428 assert_eq!(index.row_mapping_debug(), "nested");
9429 assert_eq!(index.num_chunks(), 3);
9430 assert_eq!(index.first_row(0), 0);
9432 assert_eq!(index.rows_in_chunk(0), 5);
9433 assert_eq!(index.first_row(1), 5);
9434 assert_eq!(index.rows_in_chunk(1), 4);
9435 assert_eq!(index.first_row(2), 9);
9436 assert_eq!(index.rows_in_chunk(2), 3);
9437 assert_eq!(index.items_in_chunk(0), 4);
9439 assert_eq!(index.items_in_chunk(1), 4);
9440 assert_eq!(index.items_in_chunk(2), 2);
9441
9442 let (words_nu, dbs_nu) = words_from(&[(3, 8), (1, 8), (0, 8)]);
9444 let index_nu = build_chunk_index(&words_nu, 15, 0, dbs_nu, Some(&rep), 1).unwrap();
9445 assert_eq!(index_nu.row_mapping_debug(), "nested");
9446 assert_eq!(index_nu.items_in_chunk(0), 8);
9447 assert_eq!(index_nu.items_in_chunk(1), 2);
9448 assert_eq!(index_nu.items_in_chunk(2), 5);
9449 assert_eq!(index_nu.rows_in_chunk(0), 5);
9451 }
9452
9453 #[test]
9454 fn test_uniform_flat_matches_prefix_sum_flat() {
9455 let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (2, 8), (0, 8)]);
9457 let uniform = build_chunk_index(&words, 15, 0, data_buf_size, None, 0).unwrap();
9458 assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
9459
9460 let byte_starts = PrefixSums::from_deltas([8u64, 8, 8, 8].into_iter(), 4, 32);
9462 let value_starts = PrefixSums::from_deltas([4u64, 4, 4, 3].into_iter(), 4, 15);
9463 let flat = MiniBlockChunkIndex::new(0, byte_starts, RowMapping::Flat { value_starts });
9464 assert_eq!(flat.row_mapping_debug(), "flat");
9465
9466 for i in 0..4 {
9468 assert_eq!(uniform.byte_range(i), flat.byte_range(i));
9469 assert_eq!(uniform.items_in_chunk(i), flat.items_in_chunk(i));
9470 }
9471
9472 let range_sets: Vec<Vec<std::ops::Range<u64>>> = vec![
9474 vec![0..15],
9475 vec![0..1],
9476 vec![7..8],
9477 vec![14..15],
9478 vec![3..10],
9479 vec![0..2, 5..6, 12..15],
9480 ];
9481 for ranges in &range_sets {
9482 let from_uniform = ChunkInstructions::schedule_instructions(&uniform, ranges);
9483 let from_flat = ChunkInstructions::schedule_instructions(&flat, ranges);
9484 assert_eq!(from_uniform, from_flat, "mismatch for ranges {ranges:?}");
9485 }
9486
9487 let full = ChunkInstructions::schedule_instructions(&uniform, &[0..15]);
9489 assert_eq!(full.len(), 4);
9490 for (i, inst) in full.iter().enumerate() {
9491 assert_eq!(inst.chunk_idx, i);
9492 assert_eq!(inst.preamble, PreambleAction::Absent);
9493 assert_eq!(inst.rows_to_skip, 0);
9494 assert!(!inst.take_trailer);
9495 }
9496 assert_eq!(full.iter().map(|i| i.rows_to_take).sum::<u64>(), 15);
9497 }
9498
9499 #[test]
9500 fn test_deep_size_per_variant_below_legacy() {
9501 const LEGACY_PER_CHUNK: usize = 48;
9504 let num_chunks = 3;
9505 let heap = |index: &MiniBlockChunkIndex| index.deep_size_of_children(&mut Context::new());
9506
9507 let (uniform_words, uniform_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9508 let uniform = build_chunk_index(&uniform_words, 10, 0, uniform_dbs, None, 0).unwrap();
9509 assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
9510 assert!(heap(&uniform) < LEGACY_PER_CHUNK * num_chunks);
9511
9512 let (flat_words, flat_dbs) = words_from(&[(3, 8), (1, 8), (0, 8)]);
9513 let flat = build_chunk_index(&flat_words, 11, 0, flat_dbs, None, 0).unwrap();
9514 assert_eq!(flat.row_mapping_debug(), "flat");
9515 assert!(heap(&flat) < LEGACY_PER_CHUNK * num_chunks);
9516 assert!(heap(&flat) > heap(&uniform));
9518
9519 let rep = rep_bytes_from(&[4, 0, 3, 0, 3, 0]);
9520 let (nested_words, nested_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9521 let nested = build_chunk_index(&nested_words, 10, 0, nested_dbs, Some(&rep), 1).unwrap();
9522 assert_eq!(nested.row_mapping_debug(), "nested");
9523 assert!(heap(&nested) < LEGACY_PER_CHUNK * num_chunks);
9524 }
9525
9526 #[tokio::test]
9527 async fn test_fullzip_initialize_is_lazy() {
9528 use futures::{FutureExt, future::BoxFuture};
9529 use std::ops::Range;
9530 use std::sync::Mutex;
9531
9532 #[derive(Debug, Clone)]
9533 struct RecordingScheduler {
9534 data: bytes::Bytes,
9535 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
9536 }
9537
9538 impl RecordingScheduler {
9539 fn new(data: bytes::Bytes) -> Self {
9540 Self {
9541 data,
9542 requests: Arc::new(Mutex::new(Vec::new())),
9543 }
9544 }
9545
9546 fn requests(&self) -> Vec<Vec<Range<u64>>> {
9547 self.requests.lock().unwrap().clone()
9548 }
9549 }
9550
9551 impl crate::EncodingsIo for RecordingScheduler {
9552 fn submit_request(
9553 &self,
9554 ranges: Vec<Range<u64>>,
9555 _priority: u64,
9556 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
9557 self.requests.lock().unwrap().push(ranges.clone());
9558 let data = ranges
9559 .into_iter()
9560 .map(|range| self.data.slice(range.start as usize..range.end as usize))
9561 .collect::<Vec<_>>();
9562 std::future::ready(Ok(data)).boxed()
9563 }
9564 }
9565
9566 #[derive(Debug)]
9567 struct TestFixedDecompressor;
9568
9569 impl FixedPerValueDecompressor for TestFixedDecompressor {
9570 fn decompress(
9571 &self,
9572 _data: FixedWidthDataBlock,
9573 _num_rows: u64,
9574 ) -> crate::Result<DataBlock> {
9575 unimplemented!("Test decompressor")
9576 }
9577
9578 fn bits_per_value(&self) -> u64 {
9579 32
9580 }
9581 }
9582
9583 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
9584 0;
9585 16 * 1024
9586 ])));
9587 let mut scheduler = FullZipScheduler {
9588 data_buf_position: 0,
9589 data_buf_size: 4096,
9590 rep_index: Some(FullZipRepIndexDetails {
9591 buf_position: 1000,
9592 bytes_per_value: 4,
9593 }),
9594 priority: 0,
9595 rows_in_page: 100,
9596 bits_per_offset: 32,
9597 details: Arc::new(FullZipDecodeDetails {
9598 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
9599 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
9600 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
9601 max_rep: 0,
9602 max_visible_def: 0,
9603 }),
9604 cached_state: None,
9605 enable_cache: false,
9606 };
9607
9608 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
9609 let initialization = scheduler.init_layout().unwrap();
9610 assert!(
9611 initialization.required_ranges().next().is_none(),
9612 "FullZip should not request any ranges when caching is disabled"
9613 );
9614 let buffers = initialization.with_buffers(vec![]).unwrap();
9615 let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap();
9616
9617 assert!(
9618 cached_data
9619 .as_arc_any()
9620 .downcast_ref::<super::NoCachedPageData>()
9621 .is_some(),
9622 "FullZip initialize should not eagerly load repetition index data"
9623 );
9624 assert!(scheduler.cached_state.is_none());
9625 assert!(
9626 io.requests().is_empty(),
9627 "FullZip initialize should not issue any I/O"
9628 );
9629 }
9630
9631 #[tokio::test]
9635 async fn test_initialize_coalesces_missed_page_metadata() {
9636 use std::ops::Range;
9637
9638 use futures::FutureExt;
9639 use futures::future::BoxFuture;
9640
9641 use super::{
9642 CachedPage, CachedPageData, Error, PageCacheView, PageDataCacheKey,
9643 PageInfoAndScheduler, PageInitialization, PageInitializationBuffers, PageLoadTask,
9644 StructuralPrimitiveFieldScheduler,
9645 };
9646 use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler};
9647
9648 #[derive(Debug)]
9652 struct FakePageScheduler {
9653 meta_range: Range<u64>,
9654 }
9655 impl StructuralPageScheduler for FakePageScheduler {
9656 fn init_layout(&self) -> crate::Result<PageInitialization> {
9657 Ok(PageInitialization::new([Some(self.meta_range.clone())]))
9658 }
9659 fn init_from_buffers<'a>(
9660 &'a mut self,
9661 buffers: PageInitializationBuffers,
9662 _io: &Arc<dyn EncodingsIo>,
9663 ) -> BoxFuture<'a, crate::Result<Arc<dyn CachedPageData>>> {
9664 let [buffer] = buffers.try_into_array("fake").unwrap();
9665 let buffer = buffer.unwrap();
9666 assert_eq!(
9667 &buffer[..8],
9668 &self.meta_range.start.to_le_bytes(),
9669 "each page gets the buffer for its own range"
9670 );
9671 std::future::ready(Ok(Arc::new(CacheMarker(42)) as Arc<dyn CachedPageData>)).boxed()
9672 }
9673 fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> crate::Result<()> {
9674 data.clone()
9675 .as_arc_any()
9676 .downcast::<CacheMarker>()
9677 .map(|_| ())
9678 .map_err(|_| {
9679 Error::invalid_input_source(
9680 "Fake page cache data has an unexpected type".into(),
9681 )
9682 })
9683 }
9684 fn schedule_ranges(
9685 &self,
9686 _ranges: &[Range<u64>],
9687 _io: &Arc<dyn EncodingsIo>,
9688 ) -> crate::Result<Vec<PageLoadTask>> {
9689 unreachable!("initialize never schedules data reads");
9690 }
9691 }
9692
9693 const NUM_PAGES: u64 = 64;
9694 const META_LEN: u64 = 64;
9695 const STRIDE: u64 = 1024 * 1024;
9696 let page_schedulers = (0..NUM_PAGES)
9697 .map(|i| PageInfoAndScheduler {
9698 page_index: i as usize,
9699 row_range: (i * 100)..((i + 1) * 100),
9700 scheduler: Box::new(FakePageScheduler {
9701 meta_range: (i * STRIDE)..(i * STRIDE + META_LEN),
9702 }),
9703 })
9704 .collect();
9705 let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9706 page_schedulers,
9707 0,
9708 Arc::from("test"),
9709 );
9710
9711 let io = Arc::new(RecordingScheduler::default());
9712 let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
9714 let context = SchedulerContext::new(io.clone(), cache.clone());
9715
9716 let filter = FilterExpression::no_filter();
9719 scheduler.initialize(None, &filter, &context).await.unwrap();
9720
9721 {
9722 let requests = io.requests.lock().unwrap();
9723 assert_eq!(
9724 requests.len(),
9725 1,
9726 "an N-page miss must issue exactly one request, not one per page"
9727 );
9728 assert_eq!(
9729 requests[0].len() as u64,
9730 NUM_PAGES,
9731 "the single request must carry every missed page's metadata range"
9732 );
9733 }
9734
9735 io.requests.lock().unwrap().clear();
9736 let cache_misses_before = cache.stats().await.misses;
9737 let requested_ranges = Arc::<[Range<u64>]>::from([150..151, 6_350..6_351]);
9738 scheduler
9739 .initialize(Some(&requested_ranges), &filter, &context)
9740 .await
9741 .unwrap();
9742 assert_eq!(
9743 cache.stats().await.misses - cache_misses_before,
9744 2,
9745 "range-scoped initialization must only look up touched pages"
9746 );
9747
9748 {
9749 let requests = io.requests.lock().unwrap();
9750 assert_eq!(requests.len(), 1);
9751 assert_eq!(
9752 requests[0],
9753 vec![
9754 STRIDE..(STRIDE + META_LEN),
9755 63 * STRIDE..(63 * STRIDE + META_LEN),
9756 ],
9757 "range-scoped initialization must omit untouched pages"
9758 );
9759 }
9760
9761 io.requests.lock().unwrap().clear();
9762 let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024));
9763 let cache_key = PageDataCacheKey {
9764 column_index: 9,
9765 page_index: 0,
9766 view: PageCacheView::Native,
9767 target_data_type_tag: Arc::from("test"),
9768 };
9769 cache
9770 .insert_with_key(
9771 &cache_key,
9772 Arc::new(CachedPage {
9773 data: Arc::new(super::NoCachedPageData),
9774 }),
9775 )
9776 .await;
9777 let mut recovering_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9778 vec![PageInfoAndScheduler {
9779 page_index: 0,
9780 row_range: 0..100,
9781 scheduler: Box::new(FakePageScheduler {
9782 meta_range: 0..META_LEN,
9783 }),
9784 }],
9785 9,
9786 Arc::from("test"),
9787 );
9788 let recovering_context = SchedulerContext::new(io.clone(), cache.clone());
9789 recovering_scheduler
9790 .initialize(None, &filter, &recovering_context)
9791 .await
9792 .unwrap();
9793 assert_eq!(io.requests.lock().unwrap().as_slice(), &[vec![0..META_LEN]]);
9794
9795 io.requests.lock().unwrap().clear();
9796 let mut warm_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9797 vec![PageInfoAndScheduler {
9798 page_index: 0,
9799 row_range: 0..100,
9800 scheduler: Box::new(FakePageScheduler {
9801 meta_range: 0..META_LEN,
9802 }),
9803 }],
9804 9,
9805 Arc::from("test"),
9806 );
9807 warm_scheduler
9808 .initialize(None, &filter, &recovering_context)
9809 .await
9810 .unwrap();
9811 assert!(
9812 io.requests.lock().unwrap().is_empty(),
9813 "the rebuilt cache entry must be reusable"
9814 );
9815
9816 let mut out_of_order_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9817 vec![
9818 PageInfoAndScheduler {
9819 page_index: 0,
9820 row_range: 0..100,
9821 scheduler: Box::new(FakePageScheduler {
9822 meta_range: STRIDE..(STRIDE + META_LEN),
9823 }),
9824 },
9825 PageInfoAndScheduler {
9826 page_index: 1,
9827 row_range: 100..200,
9828 scheduler: Box::new(FakePageScheduler {
9829 meta_range: 0..META_LEN,
9830 }),
9831 },
9832 ],
9833 0,
9834 Arc::from("test"),
9835 );
9836 out_of_order_scheduler
9837 .initialize(None, &filter, &context)
9838 .await
9839 .unwrap();
9840 assert_eq!(
9841 io.requests.lock().unwrap().as_slice(),
9842 &[vec![0..META_LEN, STRIDE..(STRIDE + META_LEN)]],
9843 "metadata reads must be sorted without changing page buffer assignment"
9844 );
9845 }
9846
9847 #[tokio::test]
9848 async fn page_cache_key_distinguishes_page_view_and_target_type() {
9849 let cache = lance_core::cache::LanceCache::with_capacity(4 * 1024);
9850 let cases = [
9851 (
9852 PageDataCacheKey {
9853 column_index: 7,
9854 page_index: 0,
9855 view: PageCacheView::Native,
9856 target_data_type_tag: Arc::from("Boolean"),
9857 },
9858 10,
9859 ),
9860 (
9861 PageDataCacheKey {
9862 column_index: 7,
9863 page_index: 1,
9864 view: PageCacheView::Native,
9865 target_data_type_tag: Arc::from("Boolean"),
9866 },
9867 11,
9868 ),
9869 (
9870 PageDataCacheKey {
9871 column_index: 7,
9872 page_index: 0,
9873 view: PageCacheView::Native,
9874 target_data_type_tag: Arc::from("UInt8"),
9875 },
9876 12,
9877 ),
9878 (
9879 PageDataCacheKey {
9880 column_index: 7,
9881 page_index: 0,
9882 view: PageCacheView::BlobDescriptor,
9883 target_data_type_tag: Arc::from("Struct"),
9884 },
9885 20,
9886 ),
9887 (
9888 PageDataCacheKey {
9889 column_index: 7,
9890 page_index: 0,
9891 view: PageCacheView::BlobPayload,
9892 target_data_type_tag: Arc::from("LargeBinary"),
9893 },
9894 21,
9895 ),
9896 ];
9897
9898 for (key, marker) in &cases {
9899 cache
9900 .insert_with_key(
9901 key,
9902 Arc::new(CachedPage {
9903 data: Arc::new(CacheMarker(*marker)),
9904 }),
9905 )
9906 .await;
9907 }
9908
9909 assert_eq!(cache.size().await, cases.len());
9910 for (key, expected) in &cases {
9911 let cached = cache.get_with_key(key).await.unwrap();
9912 let marker = cached
9913 .data
9914 .clone()
9915 .as_arc_any()
9916 .downcast::<CacheMarker>()
9917 .unwrap();
9918 assert_eq!(marker.0, *expected);
9919 }
9920 }
9921
9922 #[tokio::test]
9923 async fn no_op_pages_do_not_create_cache_entries() {
9924 let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9925 vec![PageInfoAndScheduler {
9926 page_index: 0,
9927 row_range: 0..10,
9928 scheduler: Box::new(SimpleAllNullScheduler::default()),
9929 }],
9930 3,
9931 Arc::from("test"),
9932 );
9933 let io: Arc<dyn crate::EncodingsIo> =
9934 Arc::new(SimulatedScheduler::new(bytes::Bytes::new()));
9935 let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024));
9936 let context = SchedulerContext::new(io, cache.clone());
9937
9938 scheduler
9939 .initialize(None, &FilterExpression::no_filter(), &context)
9940 .await
9941 .unwrap();
9942
9943 assert_eq!(cache.size().await, 0);
9944 }
9945
9946 #[test]
9947 fn page_initialization_preserves_optional_buffer_slots() {
9948 let initialization = PageInitialization::new([None, Some(20..22), None, Some(10..11)]);
9949 assert_eq!(
9950 initialization
9951 .required_ranges()
9952 .cloned()
9953 .collect::<Vec<_>>(),
9954 [20..22, 10..11]
9955 );
9956
9957 let buffers = initialization
9958 .with_buffers(vec![
9959 bytes::Bytes::from_static(b"ab"),
9960 bytes::Bytes::from_static(b"c"),
9961 ])
9962 .unwrap();
9963 let [first, second, third, fourth] = buffers.try_into_array("test").unwrap();
9964 assert!(first.is_none());
9965 assert_eq!(second.unwrap(), bytes::Bytes::from_static(b"ab"));
9966 assert!(third.is_none());
9967 assert_eq!(fourth.unwrap(), bytes::Bytes::from_static(b"c"));
9968
9969 let error = PageInitialization::new([Some(0..1)])
9970 .with_buffers(Vec::new())
9971 .unwrap_err();
9972 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9973 assert!(error.to_string().contains("1 declared ranges"));
9974 }
9975
9976 #[tokio::test]
9977 async fn initialize_builds_only_requested_page_schedulers() {
9978 use crate::compression::DefaultDecompressionStrategy;
9979 use crate::decoder::{
9980 ColumnInfo, FilterExpression, PageEncoding, PageInfo, SchedulerContext,
9981 StructuralFieldScheduler,
9982 };
9983 use crate::format::pb;
9984 use lance_core::datatypes::Field;
9985
9986 const NUM_PAGES: u64 = 64;
9987 const ROWS_PER_PAGE: u64 = 100;
9988 let page_infos = (0..NUM_PAGES)
9991 .map(|page_index| PageInfo {
9992 num_rows: ROWS_PER_PAGE,
9993 priority: page_index * ROWS_PER_PAGE,
9994 encoding: PageEncoding::Structural(pb21::PageLayout {
9995 layout: Some(pb21::page_layout::Layout::ConstantLayout(
9996 pb21::ConstantLayout {
9997 layers: vec![pb21::RepDefLayer::RepdefNullableItem as i32],
9998 ..Default::default()
9999 },
10000 )),
10001 }),
10002 buffer_offsets_and_sizes: Arc::new([]),
10003 })
10004 .collect::<Vec<_>>();
10005 let column_info = Arc::new(ColumnInfo::new(
10006 3,
10007 Arc::from(page_infos),
10008 Vec::new(),
10009 pb::ColumnEncoding::default(),
10010 ));
10011 let field = Field::try_from(&ArrowField::new("x", DataType::Int32, true)).unwrap();
10012 #[allow(deprecated)]
10013 let eager = StructuralPrimitiveFieldScheduler::try_new(
10014 &column_info,
10015 &DefaultDecompressionStrategy {},
10016 false,
10017 &field,
10018 )
10019 .unwrap();
10020 assert_eq!(eager.page_schedulers.len(), NUM_PAGES as usize);
10021 assert!(eager.page_source.is_none());
10022
10023 let mut scheduler = StructuralPrimitiveFieldScheduler::try_new_lazy(
10024 column_info,
10025 Arc::new(DefaultDecompressionStrategy {}),
10026 false,
10027 &field,
10028 )
10029 .unwrap();
10030 assert!(
10031 scheduler.page_schedulers.is_empty(),
10032 "construction must not build page schedulers"
10033 );
10034 assert_eq!(scheduler.page_row_ends.len(), NUM_PAGES as usize);
10035
10036 let io: Arc<dyn EncodingsIo> = Arc::new(SimulatedScheduler::new(bytes::Bytes::new()));
10037 let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
10038 let mut context = SchedulerContext::new(io, cache);
10039 let filter = FilterExpression::no_filter();
10040
10041 let requested_ranges = [150..151, 6_350..6_351];
10042 scheduler
10043 .initialize(Some(&requested_ranges), &filter, &context)
10044 .await
10045 .unwrap();
10046 let built_pages = |scheduler: &StructuralPrimitiveFieldScheduler| {
10047 scheduler
10048 .page_schedulers
10049 .iter()
10050 .map(|page| page.page_index)
10051 .collect::<Vec<_>>()
10052 };
10053 assert_eq!(built_pages(&scheduler), [1, 63]);
10054 assert_eq!(scheduler.page_schedulers[1].row_range, 6_300..6_400);
10055 assert!(scheduler.page_source.is_some());
10056
10057 {
10058 let mut job = scheduler
10059 .schedule_ranges(&requested_ranges, &filter)
10060 .unwrap();
10061 assert_eq!(job.schedule_next(&mut context).unwrap().len(), 2);
10063 assert!(job.schedule_next(&mut context).unwrap().is_empty());
10064
10065 let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap();
10068 let error = job.schedule_next(&mut context).unwrap_err();
10069 assert!(matches!(error, lance_core::Error::Internal { .. }));
10070 assert!(error.to_string().contains("not initialized"));
10071 }
10072
10073 scheduler
10075 .initialize(Some(&[0..250]), &filter, &context)
10076 .await
10077 .unwrap();
10078 assert_eq!(built_pages(&scheduler), [0, 1, 2, 63]);
10079
10080 scheduler.initialize(None, &filter, &context).await.unwrap();
10082 assert_eq!(scheduler.page_schedulers.len(), NUM_PAGES as usize);
10083 assert!(scheduler.page_source.is_none());
10084 let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap();
10085 assert_eq!(job.schedule_next(&mut context).unwrap().len(), 1);
10086 }
10087
10088 #[derive(Debug, Default)]
10092 struct RecordingScheduler {
10093 requests: std::sync::Mutex<Vec<Vec<Range<u64>>>>,
10094 priorities: std::sync::Mutex<Vec<u64>>,
10095 }
10096
10097 impl crate::EncodingsIo for RecordingScheduler {
10098 fn submit_request(
10099 &self,
10100 ranges: Vec<Range<u64>>,
10101 priority: u64,
10102 ) -> futures::future::BoxFuture<'static, crate::Result<Vec<Bytes>>> {
10103 use futures::FutureExt;
10104
10105 let buffers = ranges
10106 .iter()
10107 .map(|range| {
10108 let mut buffer = vec![0u8; (range.end - range.start) as usize];
10109 let marker = range.start.to_le_bytes();
10110 let marker_len = marker.len().min(buffer.len());
10111 buffer[..marker_len].copy_from_slice(&marker[..marker_len]);
10112 Bytes::from(buffer)
10113 })
10114 .collect();
10115 self.requests.lock().unwrap().push(ranges);
10116 self.priorities.lock().unwrap().push(priority);
10117 std::future::ready(Ok(buffers)).boxed()
10118 }
10119 }
10120
10121 #[tokio::test]
10122 async fn page_read_batch_submits_pending_reads_as_one_request() {
10123 let io = Arc::new(RecordingScheduler::default());
10124 let batch = RequestBatch::new(io.clone());
10125 let first = batch.submit_request(vec![100..108, 300..308], 7);
10126 let second = batch.submit_request(vec![200..208], 3);
10127 assert_eq!(batch.pending_bytes(), 24);
10128 assert!(
10129 io.requests.lock().unwrap().is_empty(),
10130 "nothing is read before the flush"
10131 );
10132
10133 batch.flush();
10134 {
10135 let requests = io.requests.lock().unwrap();
10136 assert_eq!(requests.len(), 1, "the batch is one request");
10137 assert_eq!(
10138 requests[0],
10139 vec![100..108, 200..208, 300..308],
10140 "the batch is sorted by file offset"
10141 );
10142 }
10143 assert_eq!(
10144 io.priorities.lock().unwrap()[0],
10145 3,
10146 "the batch takes the lowest priority"
10147 );
10148 let first = first.await.unwrap();
10149 assert_eq!(first.len(), 2);
10150 assert_eq!(first[0].as_ref(), &100_u64.to_le_bytes());
10151 assert_eq!(first[1].as_ref(), &300_u64.to_le_bytes());
10152 let second = second.await.unwrap();
10153 assert_eq!(second.len(), 1);
10154 assert_eq!(second[0].as_ref(), &200_u64.to_le_bytes());
10155
10156 let late = batch.submit_request(vec![400..408], 9).await.unwrap();
10158 assert_eq!(late[0].as_ref(), &400_u64.to_le_bytes());
10159 assert_eq!(io.requests.lock().unwrap().len(), 2);
10160
10161 let io = Arc::new(RecordingScheduler::default());
10164 let batch = RequestBatch::new(io.clone());
10165 let budget = crate::decoder::IO_REQUEST_BATCH_BYTES;
10166 let first = batch.submit_request(vec![0..budget], 1);
10167 assert!(io.requests.lock().unwrap().is_empty());
10168 let second = batch.submit_request(vec![budget..(budget + 8)], 2);
10169 assert_eq!(
10170 io.requests.lock().unwrap().as_slice(),
10171 &[vec![0..budget]],
10172 "the full batch is submitted before the read that would overflow it"
10173 );
10174 batch.flush();
10175 assert_eq!(io.requests.lock().unwrap().len(), 2);
10176 assert_eq!(first.await.unwrap()[0].len(), budget as usize);
10177 assert_eq!(second.await.unwrap()[0].as_ref(), &budget.to_le_bytes());
10178 }
10179
10180 #[tokio::test]
10181 async fn request_batch_flush_keeps_collecting_and_awaiting_flushes() {
10182 let io = Arc::new(RecordingScheduler::default());
10183 let batch = RequestBatch::new(io.clone());
10184 let first = batch.submit_request(vec![100..108], 1);
10185 let second = batch.submit_request(vec![300..308], 2);
10186 batch.flush();
10187 assert_eq!(
10188 io.requests.lock().unwrap().as_slice(),
10189 &[vec![100..108, 300..308]],
10190 "a flush submits the collected reads as one request"
10191 );
10192 let third = batch.submit_request(vec![500..508], 3);
10194 let fourth = batch.submit_request(vec![700..708], 4);
10195 assert_eq!(batch.pending_bytes(), 16);
10196 assert_eq!(io.requests.lock().unwrap().len(), 1);
10197 assert_eq!(third.await.unwrap()[0].as_ref(), &500_u64.to_le_bytes());
10200 assert_eq!(
10201 io.requests.lock().unwrap().as_slice(),
10202 &[vec![100..108, 300..308], vec![500..508, 700..708]]
10203 );
10204 assert_eq!(fourth.await.unwrap()[0].as_ref(), &700_u64.to_le_bytes());
10205 assert_eq!(first.await.unwrap()[0].as_ref(), &100_u64.to_le_bytes());
10206 assert_eq!(second.await.unwrap()[0].as_ref(), &300_u64.to_le_bytes());
10207 assert_eq!(io.requests.lock().unwrap().len(), 2);
10208
10209 let (fifth, sixth) = futures::join!(
10213 async { batch.submit_request(vec![900..908], 5).await.unwrap() },
10214 async { batch.submit_request(vec![1100..1108], 6).await.unwrap() },
10215 );
10216 assert_eq!(fifth[0].as_ref(), &900_u64.to_le_bytes());
10217 assert_eq!(sixth[0].as_ref(), &1100_u64.to_le_bytes());
10218 assert_eq!(
10219 io.requests.lock().unwrap().last().unwrap(),
10220 &vec![900..908, 1100..1108],
10221 "reads submitted by sibling futures share one request"
10222 );
10223 assert_eq!(io.requests.lock().unwrap().len(), 3);
10224 }
10225
10226 #[tokio::test]
10227 async fn schedule_next_batches_page_reads_up_to_the_budget() {
10228 use super::{
10229 CachedPageData, IO_REQUEST_BATCH_BYTES, PageInfoAndScheduler, PageInitialization,
10230 PageInitializationBuffers, PageLoadTask, StructuralPageScheduler,
10231 StructuralPrimitiveFieldScheduler,
10232 };
10233 use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler};
10234
10235 #[derive(Debug)]
10239 struct FakeDataPage {
10240 data_ranges: Vec<Range<u64>>,
10241 }
10242
10243 impl StructuralPageScheduler for FakeDataPage {
10244 fn needs_initialization(&self) -> bool {
10245 false
10246 }
10247
10248 fn init_layout(&self) -> crate::Result<PageInitialization> {
10249 unreachable!("fake data pages need no initialization")
10250 }
10251
10252 fn init_from_buffers<'a>(
10253 &'a mut self,
10254 _buffers: PageInitializationBuffers,
10255 _io: &Arc<dyn EncodingsIo>,
10256 ) -> futures::future::BoxFuture<'a, crate::Result<Arc<dyn CachedPageData>>>
10257 {
10258 unreachable!("fake data pages need no initialization")
10259 }
10260
10261 fn try_load(&mut self, _data: &Arc<dyn CachedPageData>) -> crate::Result<()> {
10262 Ok(())
10263 }
10264
10265 fn schedule_ranges(
10266 &self,
10267 ranges: &[Range<u64>],
10268 io: &Arc<dyn EncodingsIo>,
10269 ) -> crate::Result<Vec<PageLoadTask>> {
10270 use futures::FutureExt;
10271
10272 let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
10273 Ok(self
10274 .data_ranges
10275 .iter()
10276 .enumerate()
10277 .map(|(shard, range)| {
10278 drop(io.submit_request(vec![range.clone()], range.start));
10279 PageLoadTask {
10280 decoder_fut: std::future::ready(Err(lance_core::Error::internal(
10281 "fake page",
10282 )))
10283 .boxed(),
10284 num_rows: if shard == 0 { num_rows } else { 0 },
10286 }
10287 })
10288 .collect())
10289 }
10290 }
10291
10292 fn scheduler_with_pages(
10293 page_bytes: u64,
10294 num_pages: u64,
10295 ) -> StructuralPrimitiveFieldScheduler {
10296 let stride = page_bytes * 2;
10299 let pages = (0..num_pages)
10300 .map(|page| PageInfoAndScheduler {
10301 page_index: page as usize,
10302 row_range: (page * 10)..((page + 1) * 10),
10303 scheduler: Box::new(FakeDataPage {
10304 data_ranges: vec![(page * stride)..(page * stride + page_bytes)],
10305 }),
10306 })
10307 .collect();
10308 StructuralPrimitiveFieldScheduler::from_page_schedulers(pages, 0, Arc::from("test"))
10309 }
10310
10311 let filter = FilterExpression::no_filter();
10312 let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
10313
10314 let io = Arc::new(RecordingScheduler::default());
10316 let scheduler = scheduler_with_pages(64, 6);
10317 let mut context = SchedulerContext::new(io.clone(), cache.clone());
10318 let mut job = scheduler.schedule_ranges(&[5..45], &filter).unwrap();
10319 let scan_lines = job.schedule_next(&mut context).unwrap();
10320 context.flush_io();
10321 assert_eq!(scan_lines.len(), 5, "one scan line per touched page");
10322 assert_eq!(
10323 scan_lines
10324 .iter()
10325 .map(|line| line.rows_scheduled)
10326 .sum::<u64>(),
10327 40
10328 );
10329 assert!(job.schedule_next(&mut context).unwrap().is_empty());
10330 {
10331 let requests = io.requests.lock().unwrap();
10332 assert_eq!(requests.len(), 1, "five pages cost one request");
10333 assert_eq!(
10334 requests[0],
10335 vec![0..64, 128..192, 256..320, 384..448, 512..576],
10336 "the untouched last page is not read"
10337 );
10338 }
10339 assert_eq!(
10340 io.priorities.lock().unwrap()[0],
10341 0,
10342 "the batch carries the first page's priority"
10343 );
10344
10345 let io = Arc::new(RecordingScheduler::default());
10347 let scheduler = scheduler_with_pages(IO_REQUEST_BATCH_BYTES, 2);
10348 let mut context = SchedulerContext::new(io.clone(), cache);
10349 let mut job = scheduler.schedule_ranges(&[0..20], &filter).unwrap();
10350 let first = job.schedule_next(&mut context).unwrap();
10351 context.flush_io();
10352 assert_eq!(first.len(), 1);
10353 assert_eq!(io.requests.lock().unwrap().len(), 1);
10354 let second = job.schedule_next(&mut context).unwrap();
10355 context.flush_io();
10356 assert_eq!(second.len(), 1);
10357 assert_eq!(io.requests.lock().unwrap().len(), 2);
10358 assert!(job.schedule_next(&mut context).unwrap().is_empty());
10359
10360 let io = Arc::new(RecordingScheduler::default());
10363 let shard = IO_REQUEST_BATCH_BYTES;
10364 let scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
10365 vec![PageInfoAndScheduler {
10366 page_index: 0,
10367 row_range: 0..10,
10368 scheduler: Box::new(FakeDataPage {
10369 data_ranges: vec![0..shard, (2 * shard)..(3 * shard)],
10370 }),
10371 }],
10372 0,
10373 Arc::from("test"),
10374 );
10375 let mut context = SchedulerContext::new(
10376 io.clone(),
10377 Arc::new(lance_core::cache::LanceCache::no_cache()),
10378 );
10379 let mut job = scheduler.schedule_ranges(&[0..10], &filter).unwrap();
10380 let scan_lines = job.schedule_next(&mut context).unwrap();
10381 context.flush_io();
10382 assert_eq!(scan_lines.len(), 2, "one scan line per shard");
10383 assert_eq!(
10384 io.requests.lock().unwrap().as_slice(),
10385 &[vec![0..shard], vec![(2 * shard)..(3 * shard)]],
10386 "each shard stays its own bounded request"
10387 );
10388 assert!(job.schedule_next(&mut context).unwrap().is_empty());
10389 }
10390
10391 #[test]
10392 fn page_range_mapping_validates_and_splits_ranges() {
10393 let pages = [10_u64, 20, 30];
10395
10396 let mappings = map_ranges_to_pages(&pages, &[8..22]).unwrap();
10397 assert_eq!(mappings.len(), 3);
10398 assert_eq!(mappings[0].page_idx, 0);
10399 assert_eq!(mappings[0].ranges_in_page, [8..10]);
10400 assert_eq!(mappings[1].page_idx, 1);
10401 assert_eq!(mappings[1].ranges_in_page, [0..10]);
10402 assert_eq!(mappings[2].page_idx, 2);
10403 assert_eq!(mappings[2].ranges_in_page, [0..2]);
10404
10405 let sparse_mappings = map_ranges_to_pages(&pages, &[1..2, 3..4, 28..29]).unwrap();
10406 assert_eq!(sparse_mappings.len(), 2);
10407 assert_eq!(sparse_mappings[0].page_idx, 0);
10408 assert_eq!(sparse_mappings[0].ranges_in_page, [1..2, 3..4]);
10409 assert_eq!(sparse_mappings[1].page_idx, 2);
10410 assert_eq!(sparse_mappings[1].ranges_in_page, [8..9]);
10411 assert_eq!(
10412 pages_overlapping_ranges(&pages, &[1..2, 3..4, 28..29]).unwrap(),
10413 [0, 2]
10414 );
10415
10416 for (ranges, expected_message) in [
10417 (vec![Range { start: 2, end: 1 }], "is reversed"),
10418 (vec![15..16, 5..6], "not strictly ordered"),
10419 (vec![5..15, 10..20], "overlap"),
10420 (vec![5..5, 5..6], "not strictly ordered"),
10421 ] {
10422 let error = map_ranges_to_pages(&pages, &ranges).unwrap_err();
10423 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
10424 assert!(error.to_string().contains(expected_message));
10425 }
10426 }
10427
10428 #[tokio::test]
10429 async fn test_fullzip_read_source_slices_prefetched_page() {
10430 let page_start = 200_u64;
10431 let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
10432 let source = FullZipReadSource::PrefetchedPage {
10433 base_offset: page_start,
10434 data: page_data,
10435 };
10436 let ranges = vec![
10437 page_start..(page_start + 3),
10438 (page_start + 4)..(page_start + 8),
10439 ];
10440 let mut data = source.fetch(&ranges, 0).await.unwrap();
10441 assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
10442 assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
10443 }
10444
10445 #[tokio::test]
10446 async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
10447 use futures::{FutureExt, future::BoxFuture};
10448 use std::ops::Range;
10449 use std::sync::Mutex;
10450
10451 #[derive(Debug, Clone)]
10452 struct RecordingScheduler {
10453 data: bytes::Bytes,
10454 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
10455 }
10456
10457 impl RecordingScheduler {
10458 fn new(data: bytes::Bytes) -> Self {
10459 Self {
10460 data,
10461 requests: Arc::new(Mutex::new(Vec::new())),
10462 }
10463 }
10464
10465 fn requests(&self) -> Vec<Vec<Range<u64>>> {
10466 self.requests.lock().unwrap().clone()
10467 }
10468 }
10469
10470 impl crate::EncodingsIo for RecordingScheduler {
10471 fn submit_request(
10472 &self,
10473 ranges: Vec<Range<u64>>,
10474 _priority: u64,
10475 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
10476 self.requests.lock().unwrap().push(ranges.clone());
10477 let data = ranges
10478 .into_iter()
10479 .map(|range| self.data.slice(range.start as usize..range.end as usize))
10480 .collect::<Vec<_>>();
10481 std::future::ready(Ok(data)).boxed()
10482 }
10483 }
10484
10485 #[derive(Debug)]
10486 struct TestFixedDecompressor;
10487
10488 impl FixedPerValueDecompressor for TestFixedDecompressor {
10489 fn decompress(
10490 &self,
10491 _data: FixedWidthDataBlock,
10492 _num_rows: u64,
10493 ) -> crate::Result<DataBlock> {
10494 unimplemented!("Test decompressor")
10495 }
10496
10497 fn bits_per_value(&self) -> u64 {
10498 32
10499 }
10500 }
10501
10502 let rows_in_page = 100_u64;
10503 let bytes_per_value = 4_u64;
10504 let rep_start = 1000_u64;
10505 let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
10506 let mut data = vec![0_u8; 16 * 1024];
10507 data[rep_start as usize..rep_start as usize + rep_size].fill(7);
10508 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
10509
10510 let mut scheduler = FullZipScheduler {
10511 data_buf_position: 0,
10512 data_buf_size: 4096,
10513 rep_index: Some(FullZipRepIndexDetails {
10514 buf_position: rep_start,
10515 bytes_per_value,
10516 }),
10517 priority: 0,
10518 rows_in_page,
10519 bits_per_offset: 32,
10520 details: Arc::new(FullZipDecodeDetails {
10521 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
10522 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
10523 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
10524 max_rep: 0,
10525 max_visible_def: 0,
10526 }),
10527 cached_state: None,
10528 enable_cache: true,
10529 };
10530
10531 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
10532 let initialization = scheduler.init_layout().unwrap();
10533 let ranges = initialization
10534 .required_ranges()
10535 .cloned()
10536 .collect::<Vec<_>>();
10537 assert_eq!(
10538 ranges,
10539 vec![rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)]
10540 );
10541 let buffers = io_dyn.submit_request(ranges, 0).await.unwrap();
10544 let buffers = initialization.with_buffers(buffers).unwrap();
10545 let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap();
10546 assert!(
10547 cached_data
10548 .clone()
10549 .as_arc_any()
10550 .downcast_ref::<FullZipCacheableState>()
10551 .is_some()
10552 );
10553 assert!(scheduler.cached_state.is_none());
10554 scheduler.try_load(&cached_data).unwrap();
10555 assert!(scheduler.cached_state.is_some());
10556 assert_eq!(
10557 io.requests(),
10558 vec![vec![
10559 rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
10560 ]]
10561 );
10562 }
10563
10564 #[tokio::test]
10565 async fn test_fullzip_full_page_bypasses_rep_index_io() {
10566 use futures::{FutureExt, future::BoxFuture};
10567 use std::ops::Range;
10568 use std::sync::Mutex;
10569
10570 #[derive(Debug, Clone)]
10571 struct RecordingScheduler {
10572 data: bytes::Bytes,
10573 requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
10574 }
10575
10576 impl RecordingScheduler {
10577 fn new(data: bytes::Bytes) -> Self {
10578 Self {
10579 data,
10580 requests: Arc::new(Mutex::new(Vec::new())),
10581 }
10582 }
10583
10584 fn requests(&self) -> Vec<Vec<Range<u64>>> {
10585 self.requests.lock().unwrap().clone()
10586 }
10587 }
10588
10589 impl crate::EncodingsIo for RecordingScheduler {
10590 fn submit_request(
10591 &self,
10592 ranges: Vec<Range<u64>>,
10593 _priority: u64,
10594 ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
10595 self.requests.lock().unwrap().push(ranges.clone());
10596 let data = ranges
10597 .into_iter()
10598 .map(|range| self.data.slice(range.start as usize..range.end as usize))
10599 .collect::<Vec<_>>();
10600 std::future::ready(Ok(data)).boxed()
10601 }
10602 }
10603
10604 #[derive(Debug)]
10605 struct TestFixedDecompressor;
10606
10607 impl FixedPerValueDecompressor for TestFixedDecompressor {
10608 fn decompress(
10609 &self,
10610 _data: FixedWidthDataBlock,
10611 _num_rows: u64,
10612 ) -> crate::Result<DataBlock> {
10613 unimplemented!("Test decompressor")
10614 }
10615
10616 fn bits_per_value(&self) -> u64 {
10617 32
10618 }
10619 }
10620
10621 let rows_in_page = 100_u64;
10622 let data_start = 256_u64;
10623 let data_size = 500_u64;
10624 let rep_start = 4096_u64;
10625 let bytes_per_value = 4_u64;
10626
10627 let mut bytes = vec![0_u8; 16 * 1024];
10628 for i in 0..=rows_in_page {
10629 let offset = (i * 5) as u32;
10630 let pos = rep_start as usize + (i * bytes_per_value) as usize;
10631 bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
10632 }
10633 let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
10634
10635 let scheduler = FullZipScheduler {
10636 data_buf_position: data_start,
10637 data_buf_size: data_size,
10638 rep_index: Some(FullZipRepIndexDetails {
10639 buf_position: rep_start,
10640 bytes_per_value,
10641 }),
10642 priority: 0,
10643 rows_in_page,
10644 bits_per_offset: 32,
10645 details: Arc::new(FullZipDecodeDetails {
10646 value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
10647 def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
10648 ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
10649 max_rep: 0,
10650 max_visible_def: 0,
10651 }),
10652 cached_state: None,
10653 enable_cache: false,
10654 };
10655
10656 let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
10657 let tasks = scheduler
10658 .schedule_ranges_rep(
10659 &[0..rows_in_page],
10660 &io_dyn,
10661 FullZipRepIndexDetails {
10662 buf_position: rep_start,
10663 bytes_per_value,
10664 },
10665 )
10666 .unwrap();
10667
10668 let requests = io.requests();
10669 assert_eq!(requests.len(), 1);
10670 assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
10671
10672 let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
10673 let requests_after_await = io.requests();
10674 assert_eq!(
10675 requests_after_await.len(),
10676 1,
10677 "full page path should not issue rep-index I/O"
10678 );
10679 }
10680
10681 #[tokio::test]
10683 async fn test_fuzz_issue_4492_empty_rep_values() {
10684 use lance_datagen::{RowCount, Seed, array, gen_batch};
10685
10686 let seed = 1823859942947654717u64;
10687 let num_rows = 2741usize;
10688
10689 let batch_gen = gen_batch().with_seed(Seed::from(seed));
10691 let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
10692 let list_generator = array::rand_list_any(base_generator, false);
10693
10694 let batch = batch_gen
10695 .anon_col(list_generator)
10696 .into_batch_rows(RowCount::from(num_rows as u64))
10697 .unwrap();
10698
10699 let list_array = batch.column(0).clone();
10700
10701 let mut metadata = HashMap::new();
10703 metadata.insert(
10704 STRUCTURAL_ENCODING_META_KEY.to_string(),
10705 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
10706 );
10707
10708 let test_cases = TestCases::default()
10709 .with_structural_encodings()
10710 .with_batch_size(100)
10711 .with_range(0..num_rows.min(500) as u64)
10712 .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
10713
10714 check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
10715 }
10716
10717 async fn test_minichunk_size_helper(
10718 string_data: Vec<Option<String>>,
10719 minichunk_size: u64,
10720 encodings: &[TestEncoding],
10721 ) {
10722 use crate::constants::MINICHUNK_SIZE_META_KEY;
10723 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
10724 use arrow_array::{ArrayRef, StringArray};
10725 use std::sync::Arc;
10726
10727 let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
10728
10729 let mut metadata = HashMap::new();
10730 metadata.insert(
10731 MINICHUNK_SIZE_META_KEY.to_string(),
10732 minichunk_size.to_string(),
10733 );
10734 metadata.insert(
10735 STRUCTURAL_ENCODING_META_KEY.to_string(),
10736 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
10737 );
10738
10739 let test_cases = TestCases::default()
10740 .with_encodings(encodings.iter().copied())
10741 .with_batch_size(1000);
10742
10743 check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
10744 }
10745
10746 #[tokio::test]
10747 async fn test_minichunk_size_roundtrip() {
10748 let mut string_data = Vec::new();
10750 for i in 0..100 {
10751 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
10752 }
10753 test_minichunk_size_helper(
10755 string_data,
10756 64,
10757 &[
10758 TestEncoding::StructuralU16,
10759 TestEncoding::StructuralU32,
10760 TestEncoding::StructuralSparse,
10761 ],
10762 )
10763 .await;
10764 }
10765
10766 #[tokio::test]
10767 async fn test_minichunk_size_128kb_v2_2() {
10768 let mut string_data = Vec::new();
10770 for i in 0..10000 {
10772 string_data.push(Some(format!("test_string_{}", i).repeat(50)));
10773 }
10774 test_minichunk_size_helper(
10775 string_data,
10776 128 * 1024,
10777 &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
10778 )
10779 .await;
10780 }
10781
10782 #[tokio::test]
10783 async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
10784 let mut string_data = Vec::new();
10785 for i in 0..10000 {
10787 string_data.push(Some(format!("t_{}", i)));
10788 }
10789 test_minichunk_size_helper(
10790 string_data,
10791 128 * 1024,
10792 &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
10793 )
10794 .await;
10795 }
10796
10797 #[tokio::test]
10798 async fn test_large_dictionary_general_compression() {
10799 use arrow_array::{ArrayRef, StringArray};
10800 use std::collections::HashMap;
10801 use std::sync::Arc;
10802
10803 let unique_values: Vec<String> = (0..100)
10806 .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
10807 .collect();
10808
10809 let repeated_strings: Vec<_> = unique_values
10811 .iter()
10812 .cycle()
10813 .take(10_000)
10814 .map(|s| Some(s.as_str()))
10815 .collect();
10816
10817 let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
10818
10819 let test_cases = TestCases::default()
10821 .with_u32_structural_encodings()
10822 .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
10823 assert_eq!(cols.len(), 1);
10824 let col = &cols[0];
10825
10826 if let Some(PageEncoding::Structural(page_layout)) =
10828 &col.final_pages.first().map(|p| &p.description)
10829 && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
10830 &page_layout.layout
10831 && let Some(dictionary_encoding) = &mini_block.dictionary
10832 {
10833 match dictionary_encoding.compression.as_ref() {
10834 Some(Compression::General(general)) => {
10835 let compression = general.compression.as_ref().unwrap();
10837 assert!(
10838 compression.scheme()
10839 == pb21::CompressionScheme::CompressionAlgorithmLz4
10840 || compression.scheme()
10841 == pb21::CompressionScheme::CompressionAlgorithmZstd,
10842 "Expected LZ4 or Zstd compression for large dictionary"
10843 );
10844 }
10845 _ => panic!("Expected General compression for large dictionary"),
10846 }
10847 }
10848 }));
10849
10850 check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
10851 }
10852
10853 fn dictionary_encoding_from_page(
10854 page: &crate::encoder::EncodedPage,
10855 ) -> &crate::format::pb21::CompressiveEncoding {
10856 let PageEncoding::Structural(layout) = &page.description else {
10857 panic!("Expected structural page encoding");
10858 };
10859 let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
10860 else {
10861 panic!("Expected mini-block layout");
10862 };
10863 layout
10864 .dictionary
10865 .as_ref()
10866 .unwrap_or_else(|| panic!("Expected dictionary encoding"))
10867 }
10868
10869 async fn encode_variable_dict_page(
10870 metadata: HashMap<String, String>,
10871 ) -> crate::encoder::EncodedPage {
10872 use arrow_array::types::Int32Type;
10873 use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
10874
10875 let values = Arc::new(StringArray::from(
10876 (0..128)
10877 .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
10878 .collect::<Vec<_>>(),
10879 )) as ArrayRef;
10880 let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
10881 let dict_array =
10882 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
10883
10884 let field = arrow_schema::Field::new(
10885 "dict_col",
10886 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
10887 false,
10888 )
10889 .with_metadata(metadata);
10890
10891 encode_first_page(field, dict_array, TestEncoding::StructuralU32).await
10892 }
10893
10894 async fn encode_auto_fixed_dict_page(
10895 metadata: HashMap<String, String>,
10896 ) -> crate::encoder::EncodedPage {
10897 use arrow_array::{ArrayRef, Decimal128Array};
10898
10899 let values = (0..20_000)
10901 .map(|i| match i % 3 {
10902 0 => 10_i128,
10903 1 => 20_i128,
10904 _ => 30_i128,
10905 })
10906 .collect::<Vec<_>>();
10907 let decimal = Decimal128Array::from_iter_values(values)
10908 .with_precision_and_scale(38, 0)
10909 .unwrap();
10910 let decimal = Arc::new(decimal) as ArrayRef;
10911
10912 let mut field_metadata = metadata;
10913 field_metadata.insert(
10915 "lance-encoding:dict-size-ratio".to_string(),
10916 "0.99".to_string(),
10917 );
10918 let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
10919 .with_metadata(field_metadata);
10920
10921 encode_first_page(field, decimal, TestEncoding::StructuralU32).await
10922 }
10923
10924 #[tokio::test]
10925 async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
10926 let page = encode_variable_dict_page(HashMap::new()).await;
10927 let dictionary_encoding = dictionary_encoding_from_page(&page);
10928 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10929 panic!("Expected General compression for dictionary values");
10930 };
10931 let compression = general.compression.as_ref().unwrap();
10932 assert_eq!(
10933 compression.scheme(),
10934 pb21::CompressionScheme::CompressionAlgorithmLz4
10935 );
10936 }
10937
10938 #[tokio::test]
10939 async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
10940 let page = encode_auto_fixed_dict_page(HashMap::new()).await;
10941 let dictionary_encoding = dictionary_encoding_from_page(&page);
10942 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10943 panic!("Expected General compression for dictionary values");
10944 };
10945 let compression = general.compression.as_ref().unwrap();
10946 assert_eq!(
10947 compression.scheme(),
10948 pb21::CompressionScheme::CompressionAlgorithmLz4
10949 );
10950 }
10951
10952 #[tokio::test]
10953 async fn test_dict_values_general_compression_zstd() {
10954 let mut metadata = HashMap::new();
10955 metadata.insert(
10956 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
10957 "zstd".to_string(),
10958 );
10959 let page = encode_variable_dict_page(metadata).await;
10960 let dictionary_encoding = dictionary_encoding_from_page(&page);
10961 let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10962 panic!("Expected General compression for dictionary values");
10963 };
10964 let compression = general.compression.as_ref().unwrap();
10965 assert_eq!(
10966 compression.scheme(),
10967 pb21::CompressionScheme::CompressionAlgorithmZstd
10968 );
10969 }
10970
10971 #[tokio::test]
10972 async fn test_dict_values_general_compression_none() {
10973 let mut metadata = HashMap::new();
10974 metadata.insert(
10975 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
10976 "none".to_string(),
10977 );
10978 let page = encode_variable_dict_page(metadata).await;
10979 let dictionary_encoding = dictionary_encoding_from_page(&page);
10980 assert!(
10981 !matches!(
10982 dictionary_encoding.compression.as_ref(),
10983 Some(Compression::General(_))
10984 ),
10985 "Expected dictionary values to avoid General compression"
10986 );
10987 }
10988
10989 #[test]
10990 fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
10991 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
10992 &HashMap::new(),
10993 None,
10994 None,
10995 );
10996 assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
10997 assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
10998 }
10999
11000 #[test]
11001 fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
11002 let field_metadata = HashMap::from([
11003 (
11004 DICT_VALUES_COMPRESSION_META_KEY.to_string(),
11005 "none".to_string(),
11006 ),
11007 (
11008 DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
11009 "7".to_string(),
11010 ),
11011 ]);
11012 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
11013 &field_metadata,
11014 Some("zstd".to_string()),
11015 Some("3".to_string()),
11016 );
11017 assert_eq!(
11018 metadata.get(COMPRESSION_META_KEY),
11019 Some(&"none".to_string()),
11020 );
11021 assert_eq!(
11022 metadata.get(COMPRESSION_LEVEL_META_KEY),
11023 Some(&"7".to_string()),
11024 );
11025 }
11026
11027 #[test]
11028 fn test_resolve_dict_values_compression_metadata_env_fallback() {
11029 let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
11030 &HashMap::new(),
11031 Some("zstd".to_string()),
11032 Some("9".to_string()),
11033 );
11034 assert_eq!(
11035 metadata.get(COMPRESSION_META_KEY),
11036 Some(&"zstd".to_string()),
11037 );
11038 assert_eq!(
11039 metadata.get(COMPRESSION_LEVEL_META_KEY),
11040 Some(&"9".to_string()),
11041 );
11042 }
11043
11044 #[tokio::test]
11045 async fn test_dictionary_encode_int64() {
11046 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
11047 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
11048 use arrow_array::{ArrayRef, Int64Array};
11049 use std::collections::HashMap;
11050 use std::sync::Arc;
11051
11052 let values = (0..1000)
11054 .map(|i| match i % 3 {
11055 0 => 10i64,
11056 1 => 20i64,
11057 _ => 30i64,
11058 })
11059 .collect::<Vec<_>>();
11060 let array = Arc::new(Int64Array::from(values)) as ArrayRef;
11061
11062 let mut metadata = HashMap::new();
11063 metadata.insert(
11064 STRUCTURAL_ENCODING_META_KEY.to_string(),
11065 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
11066 );
11067 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
11068
11069 let test_cases = TestCases::default()
11070 .with_u32_structural_encodings()
11071 .with_batch_size(1000)
11072 .with_range(0..1000)
11073 .with_indices(vec![0, 1, 10, 999])
11074 .with_expected_encoding("dictionary");
11075
11076 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
11077 }
11078
11079 #[tokio::test]
11080 async fn test_dictionary_encode_float64() {
11081 use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
11082 use crate::testing::{TestCases, check_round_trip_encoding_of_data};
11083 use arrow_array::{ArrayRef, Float64Array};
11084 use std::collections::HashMap;
11085 use std::sync::Arc;
11086
11087 let values = (0..1000)
11089 .map(|i| match i % 3 {
11090 0 => 0.1f64,
11091 1 => 0.2f64,
11092 _ => 0.3f64,
11093 })
11094 .collect::<Vec<_>>();
11095 let array = Arc::new(Float64Array::from(values)) as ArrayRef;
11096
11097 let mut metadata = HashMap::new();
11098 metadata.insert(
11099 STRUCTURAL_ENCODING_META_KEY.to_string(),
11100 STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
11101 );
11102 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
11103
11104 let test_cases = TestCases::default()
11105 .with_u32_structural_encodings()
11106 .with_batch_size(1000)
11107 .with_range(0..1000)
11108 .with_indices(vec![0, 1, 10, 999])
11109 .with_expected_encoding("dictionary");
11110
11111 check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
11112 }
11113
11114 #[test]
11115 fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
11116 let rows = 10_000;
11117 let unique_values = 2_000;
11118
11119 let dictionary_encoding =
11120 ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
11121 let layout = pb21::MiniBlockLayout {
11122 rep_compression: None,
11123 def_compression: None,
11124 value_compression: Some(ProtobufUtils21::flat(64, None)),
11125 dictionary: Some(dictionary_encoding),
11126 num_dictionary_items: unique_values,
11127 layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
11128 num_buffers: 1,
11129 repetition_index_depth: 0,
11130 num_items: rows,
11131 has_large_chunk: false,
11132 };
11133
11134 let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
11135 let scheduler = super::MiniBlockScheduler::try_new(
11136 &buffer_offsets_and_sizes,
11137 0,
11138 rows,
11139 &layout,
11140 &DefaultDecompressionStrategy::default(),
11141 )
11142 .unwrap();
11143
11144 let dictionary = scheduler.dictionary.unwrap();
11145 assert_eq!(dictionary.num_dictionary_items, unique_values);
11146 assert_eq!(
11147 dictionary.dictionary_data_alignment,
11148 crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
11149 );
11150 }
11151
11152 fn create_test_fixed_data_block(
11154 num_values: u64,
11155 cardinality: u64,
11156 bits_per_value: u64,
11157 ) -> DataBlock {
11158 assert!(cardinality > 0);
11159 assert!(cardinality <= num_values);
11160 let block_info = BlockInfo::default();
11161
11162 assert_eq!(bits_per_value % 8, 0);
11163 let data = match bits_per_value {
11164 32 => {
11165 let values = (0..num_values)
11166 .map(|i| (i % cardinality) as u32)
11167 .collect::<Vec<_>>();
11168 crate::buffer::LanceBuffer::reinterpret_vec(values)
11169 }
11170 64 => {
11171 let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
11172 crate::buffer::LanceBuffer::reinterpret_vec(values)
11173 }
11174 128 => {
11175 let values = (0..num_values)
11176 .map(|i| (i % cardinality) as u128)
11177 .collect::<Vec<_>>();
11178 crate::buffer::LanceBuffer::reinterpret_vec(values)
11179 }
11180 _ => unreachable!(),
11181 };
11182 DataBlock::FixedWidth(FixedWidthDataBlock {
11183 bits_per_value,
11184 data,
11185 num_values,
11186 block_info,
11187 })
11188 }
11189
11190 fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
11192 use arrow_array::StringArray;
11193
11194 assert!(cardinality <= num_values && cardinality > 0);
11195
11196 let mut values = Vec::with_capacity(num_values as usize);
11197 for i in 0..num_values {
11198 values.push(format!("value_{:016}", i % cardinality));
11199 }
11200
11201 let array = StringArray::from(values);
11202 DataBlock::from_array(Arc::new(array) as ArrayRef)
11203 }
11204
11205 fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
11206 use arrow_array::StringArray;
11207
11208 assert!(cardinality <= num_values && cardinality > 0);
11209
11210 let mut values = Vec::with_capacity(num_values as usize);
11211 for i in 0..num_values {
11212 let value_idx = i * cardinality / num_values;
11213 values.push(format!("value_{:016}", value_idx));
11214 }
11215
11216 Arc::new(StringArray::from(values)) as ArrayRef
11217 }
11218
11219 fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
11220 DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
11221 }
11222
11223 #[test]
11224 fn test_should_dictionary_encode() {
11225 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11226 use lance_core::datatypes::Field as LanceField;
11227
11228 let block = create_test_variable_width_block(1000, 10);
11230
11231 let mut metadata = HashMap::new();
11232 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11233 let arrow_field =
11234 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11235 let field = LanceField::try_from(&arrow_field).unwrap();
11236
11237 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11238 &block,
11239 &field,
11240 FixedWidthDictionaryEncoding::Exclude64Bit,
11241 );
11242
11243 assert!(
11244 result.is_some(),
11245 "Should use dictionary encode based on size"
11246 );
11247 }
11248
11249 #[test]
11250 fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
11251 let sample_count: usize = 4096;
11252 let num_values: u64 = 200_000;
11253 let cardinality: u64 = 8_000;
11254 let run_length = num_values / cardinality;
11255 let stride = num_values as usize / sample_count;
11256 assert!(
11257 stride > run_length as usize,
11258 "test must construct the stride > run_length case"
11259 );
11260
11261 let block = create_sorted_variable_width_block(num_values, cardinality);
11262 let sample_unique_ratio =
11263 PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
11264
11265 assert!(
11266 sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
11267 "sorted low-cardinality data must not be classified as near-unique"
11268 );
11269 }
11270
11271 #[test]
11272 fn test_should_dictionary_encode_sorted_low_cardinality() {
11273 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11274 use lance_core::datatypes::Field as LanceField;
11275
11276 let block = create_sorted_variable_width_block(200_000, 8_000);
11277
11278 let mut metadata = HashMap::new();
11279 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11280 let arrow_field =
11281 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11282 let field = LanceField::try_from(&arrow_field).unwrap();
11283
11284 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11285 &block,
11286 &field,
11287 FixedWidthDictionaryEncoding::Include64Bit,
11288 );
11289
11290 assert!(
11291 result.is_some(),
11292 "sorted low-cardinality data should reach dictionary encoding"
11293 );
11294 }
11295
11296 #[test]
11297 fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
11298 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11299 use lance_core::datatypes::Field as LanceField;
11300
11301 let num_values = 200_002;
11302 let cardinality = 100_001;
11303 let block = create_sorted_variable_width_block(num_values, cardinality);
11304
11305 let mut metadata = HashMap::new();
11306 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11307 let arrow_field =
11308 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11309 let field = LanceField::try_from(&arrow_field).unwrap();
11310
11311 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11312 &block,
11313 &field,
11314 FixedWidthDictionaryEncoding::Include64Bit,
11315 );
11316
11317 assert!(
11318 result.is_none(),
11319 "sorted high-cardinality short runs should not trigger a full dictionary probe"
11320 );
11321 }
11322
11323 #[tokio::test]
11324 async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
11325 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11326
11327 let mut metadata = HashMap::new();
11328 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11329 let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11330 let array = create_sorted_string_array(200_000, 8_000);
11331
11332 let page = encode_first_page(field, array, TestEncoding::StructuralU32).await;
11333 let _ = dictionary_encoding_from_page(&page);
11334 }
11335
11336 #[test]
11337 fn test_should_not_dictionary_encode_unsupported_bits() {
11338 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11339 use lance_core::datatypes::Field as LanceField;
11340
11341 let block = create_test_fixed_data_block(1000, 1000, 32);
11342
11343 let mut metadata = HashMap::new();
11344 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11345 let arrow_field =
11346 arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
11347 let field = LanceField::try_from(&arrow_field).unwrap();
11348
11349 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11350 &block,
11351 &field,
11352 FixedWidthDictionaryEncoding::Exclude64Bit,
11353 );
11354
11355 assert!(
11356 result.is_none(),
11357 "Should not use dictionary encode for unsupported bit width"
11358 );
11359 }
11360
11361 #[test]
11362 fn test_should_not_dictionary_encode_near_unique_sample() {
11363 use crate::constants::DICT_SIZE_RATIO_META_KEY;
11364 use lance_core::datatypes::Field as LanceField;
11365
11366 let num_values = 5000;
11367 let block = create_test_variable_width_block(num_values, num_values);
11368
11369 let mut metadata = HashMap::new();
11370 metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
11371 let arrow_field =
11372 arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11373 let field = LanceField::try_from(&arrow_field).unwrap();
11374
11375 let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11376 &block,
11377 &field,
11378 FixedWidthDictionaryEncoding::Exclude64Bit,
11379 );
11380
11381 assert!(
11382 result.is_none(),
11383 "Should not probe dictionary encoding for near-unique data"
11384 );
11385 }
11386
11387 #[test]
11388 fn test_v2_1_miniblock_serializes_log_num_values_15() {
11389 let miniblocks = MiniBlockCompressed {
11390 data: vec![LanceBuffer::from(vec![1_u8; 16])],
11391 chunks: vec![
11392 MiniBlockChunk {
11393 buffer_sizes: vec![8],
11394 log_num_values: 15,
11395 },
11396 MiniBlockChunk {
11397 buffer_sizes: vec![8],
11398 log_num_values: 0,
11399 },
11400 ],
11401 num_values: 32_769,
11402 };
11403
11404 let serialized = PrimitiveStructuralEncoder::serialize_miniblocks(
11405 miniblocks,
11406 None,
11407 None,
11408 MiniblockChunkSize::U16,
11409 )
11410 .unwrap();
11411
11412 let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
11413 assert_eq!(chunk_metadata.len(), 2);
11414 assert_eq!(
11415 chunk_metadata[0] & 0x0F,
11416 15,
11417 "V2.1 metadata should use all 4 bits for log_num_values"
11418 );
11419 }
11420
11421 async fn encode_first_page(
11422 field: arrow_schema::Field,
11423 array: ArrayRef,
11424 version: TestEncoding,
11425 ) -> crate::encoder::EncodedPage {
11426 use crate::repdef::RepDefBuilder;
11427 use crate::{
11428 encoder::{
11429 ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
11430 },
11431 testing::{create_test_field_encoder, test_encoding_strategy},
11432 };
11433
11434 let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
11435 let encoding_strategy = test_encoding_strategy(version);
11436 let mut column_index_seq = ColumnIndexSequence::default();
11437 let encoding_options = EncodingOptions {
11438 cache_bytes_per_column: 1,
11439 max_page_bytes: 32 * 1024 * 1024,
11440 keep_original_array: true,
11441 buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
11442 };
11443
11444 let mut encoder = create_test_field_encoder(
11445 encoding_strategy.as_ref(),
11446 &lance_field,
11447 &mut column_index_seq,
11448 &encoding_options,
11449 )
11450 .unwrap();
11451
11452 let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
11453 let repdef = RepDefBuilder::default();
11454 let num_rows = array.len() as u64;
11455 let mut pages = Vec::new();
11456 for task in encoder
11457 .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
11458 .unwrap()
11459 {
11460 pages.push(task.await.unwrap());
11461 }
11462 for task in encoder.flush(&mut external_buffers).unwrap() {
11463 pages.push(task.await.unwrap());
11464 }
11465 pages.into_iter().next().unwrap()
11466 }
11467
11468 #[tokio::test]
11469 async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
11470 use crate::format::pb21::page_layout::Layout;
11471
11472 let val = vec![0xABu8; 33];
11473 let arr: ArrayRef = Arc::new(
11474 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
11475 std::iter::repeat_n(Some(val.as_slice()), 256),
11476 33,
11477 )
11478 .unwrap(),
11479 );
11480 let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
11481 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11482
11483 let PageEncoding::Structural(layout) = &page.description else {
11484 panic!("Expected structural encoding");
11485 };
11486 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11487 panic!("Expected constant layout in slot 2");
11488 };
11489 assert!(layout.inline_value.is_none());
11490 assert_eq!(page.data.len(), 1);
11491
11492 let test_cases = TestCases::default()
11493 .with_encoding(TestEncoding::StructuralU32)
11494 .with_page_sizes(vec![4096]);
11495 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11496 }
11497
11498 #[tokio::test]
11499 async fn test_constant_layout_out_of_line_utf8_v2_2() {
11500 use crate::format::pb21::page_layout::Layout;
11501
11502 let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
11503 std::iter::repeat_n("hello", 512),
11504 ));
11505 let field = arrow_schema::Field::new("c", DataType::Utf8, true);
11506 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11507
11508 let PageEncoding::Structural(layout) = &page.description else {
11509 panic!("Expected structural encoding");
11510 };
11511 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11512 panic!("Expected constant layout in slot 2");
11513 };
11514 assert!(layout.inline_value.is_none());
11515 assert_eq!(page.data.len(), 1);
11516
11517 let test_cases = TestCases::default()
11518 .with_encoding(TestEncoding::StructuralU32)
11519 .with_page_sizes(vec![4096]);
11520 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11521 }
11522
11523 #[tokio::test]
11524 async fn test_constant_layout_nullable_item_v2_2() {
11525 use crate::format::pb21::page_layout::Layout;
11526
11527 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
11528 Some(7),
11529 None,
11530 Some(7),
11531 None,
11532 Some(7),
11533 ]));
11534 let field = arrow_schema::Field::new("c", DataType::Int32, true);
11535 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11536
11537 let PageEncoding::Structural(layout) = &page.description else {
11538 panic!("Expected structural encoding");
11539 };
11540 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11541 panic!("Expected constant layout in slot 2");
11542 };
11543 assert!(layout.inline_value.is_some());
11544 assert_eq!(page.data.len(), 2);
11545
11546 let test_cases = TestCases::default()
11547 .with_encoding(TestEncoding::StructuralU32)
11548 .with_page_sizes(vec![4096]);
11549 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11550 }
11551
11552 #[tokio::test]
11553 async fn test_constant_layout_list_repdef_v2_2() {
11554 use crate::format::pb21::page_layout::Layout;
11555 use arrow_array::builder::{Int32Builder, ListBuilder};
11556
11557 let mut builder = ListBuilder::new(Int32Builder::new());
11558 builder.values().append_value(7);
11559 builder.values().append_null();
11560 builder.values().append_value(7);
11561 builder.append(true);
11562
11563 builder.append(true);
11564
11565 builder.values().append_value(7);
11566 builder.append(true);
11567
11568 builder.append_null();
11569
11570 let arr: ArrayRef = Arc::new(builder.finish());
11571 let field = arrow_schema::Field::new(
11572 "c",
11573 DataType::List(Arc::new(arrow_schema::Field::new(
11574 "item",
11575 DataType::Int32,
11576 true,
11577 ))),
11578 true,
11579 );
11580 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11581
11582 let PageEncoding::Structural(layout) = &page.description else {
11583 panic!("Expected structural encoding");
11584 };
11585 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11586 panic!("Expected constant layout in slot 2");
11587 };
11588 assert!(layout.inline_value.is_some());
11589 assert_eq!(page.data.len(), 2);
11590
11591 let test_cases = TestCases::default()
11592 .with_encoding(TestEncoding::StructuralU32)
11593 .with_page_sizes(vec![4096]);
11594 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11595 }
11596
11597 #[tokio::test]
11598 async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
11599 use crate::format::pb21::page_layout::Layout;
11600 use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
11601
11602 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
11603 for _ in 0..64 {
11604 builder.values().append_value(1);
11605 builder.values().append_null();
11606 builder.values().append_value(3);
11607 builder.append(true);
11608 }
11609 let arr: ArrayRef = Arc::new(builder.finish());
11610 let field = arrow_schema::Field::new(
11611 "c",
11612 DataType::FixedSizeList(
11613 Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
11614 3,
11615 ),
11616 true,
11617 );
11618 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11619
11620 if let PageEncoding::Structural(layout) = &page.description {
11621 assert!(
11622 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
11623 "FixedSizeList should not use constant layout yet"
11624 );
11625 }
11626
11627 let test_cases = TestCases::default()
11628 .with_encoding(TestEncoding::StructuralU32)
11629 .with_page_sizes(vec![4096]);
11630 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11631 }
11632
11633 #[tokio::test]
11634 async fn test_constant_layout_not_written_before_v2_2() {
11635 use crate::format::pb21::page_layout::Layout;
11636
11637 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
11638 let field = arrow_schema::Field::new("c", DataType::Int32, true);
11639 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU16).await;
11640
11641 let PageEncoding::Structural(layout) = &page.description else {
11642 return;
11643 };
11644 assert!(
11645 !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
11646 "Should not emit constant layout before v2.2"
11647 );
11648
11649 let test_cases = TestCases::default()
11650 .with_encoding(TestEncoding::StructuralU16)
11651 .with_page_sizes(vec![4096]);
11652 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11653 }
11654
11655 #[tokio::test]
11656 async fn test_all_null_constant_layout_still_works_v2_2() {
11657 use crate::format::pb21::page_layout::Layout;
11658
11659 let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
11660 let field = arrow_schema::Field::new("c", DataType::Int32, true);
11661 let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11662
11663 let PageEncoding::Structural(layout) = &page.description else {
11664 panic!("Expected structural encoding");
11665 };
11666 let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11667 panic!("Expected layout in slot 2");
11668 };
11669 assert!(layout.inline_value.is_none());
11670 assert_eq!(page.data.len(), 0);
11671
11672 let test_cases = TestCases::default()
11673 .with_encoding(TestEncoding::StructuralU32)
11674 .with_page_sizes(vec![4096]);
11675 check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11676 }
11677
11678 fn hand_built_dictionary_with_out_of_range_null_keys() -> ArrayRef {
11679 use arrow_array::{DictionaryArray, Int32Array, types::Int32Type};
11680 use arrow_buffer::NullBuffer;
11681
11682 let keys = Int32Array::new(
11683 vec![0, 7, 7].into(),
11684 Some(NullBuffer::from(vec![true, false, false])),
11685 );
11686 let values = Arc::new(StringArray::from(vec!["a"]));
11687 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef
11688 }
11689
11690 fn concatenated_dictionary_with_out_of_range_null_keys() -> ArrayRef {
11691 use arrow_array::{builder::StringDictionaryBuilder, new_null_array, types::Int32Type};
11692
11693 let mut builder = StringDictionaryBuilder::<Int32Type>::new();
11694 builder.append_value("a");
11695 for _ in 0..7 {
11696 builder.append_null();
11697 }
11698 let valued = Arc::new(builder.finish()) as ArrayRef;
11699 let all_null = new_null_array(valued.data_type(), 8);
11700 arrow_select::concat::concat(&[valued.as_ref(), all_null.as_ref()]).unwrap()
11701 }
11702
11703 #[rstest::rstest]
11704 #[case::hand_built(hand_built_dictionary_with_out_of_range_null_keys(), 1)]
11705 #[case::concatenated(concatenated_dictionary_with_out_of_range_null_keys(), 8)]
11706 #[tokio::test]
11707 async fn test_dictionary_out_of_range_null_keys_round_trip(
11708 #[case] dictionary: ArrayRef,
11709 #[case] split_at: usize,
11710 ) {
11711 let num_rows = dictionary.len();
11712 let dictionaries = vec![
11713 dictionary.slice(0, split_at),
11714 dictionary.slice(split_at, num_rows - split_at),
11715 ];
11716 let test_cases = TestCases::default()
11717 .with_encoding(TestEncoding::StructuralU32)
11718 .with_page_sizes(vec![1])
11719 .with_range(split_at as u64..num_rows as u64)
11720 .with_indices(vec![split_at as u64]);
11721
11722 check_round_trip_encoding_of_data(dictionaries, &test_cases, HashMap::new()).await;
11723 }
11724
11725 #[test]
11726 fn test_encode_decode_complex_all_null_vals_roundtrip() {
11727 use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy};
11728
11729 let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
11730
11731 let compression_strategy = crate::testing::test_compression_strategy(
11732 TestEncoding::StructuralU16,
11733 crate::compression_config::CompressionParams::default(),
11734 );
11735 let decompression_strategy = DefaultDecompressionStrategy::default();
11736
11737 let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
11738 &values,
11739 compression_strategy.as_ref(),
11740 )
11741 .unwrap();
11742
11743 let decompressor = decompression_strategy
11744 .create_block_decompressor(&encoding)
11745 .unwrap();
11746 let decompressed = decompressor
11747 .decompress(Some(compressed_buf), values.len() as u64)
11748 .unwrap();
11749 let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
11750 assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
11751 assert_eq!(decompressed_fixed_width.bits_per_value, 16);
11752 let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
11753 assert_eq!(rep_result.as_ref(), values.as_ref());
11754 }
11755
11756 #[tokio::test]
11757 async fn test_complex_all_null_compression_gated_by_version() {
11758 use crate::format::pb21::page_layout::Layout;
11759 use arrow_array::ListArray;
11760
11761 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11762 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
11763 );
11764 let arr: ArrayRef = Arc::new(list_array);
11765 let field = arrow_schema::Field::new(
11766 "c",
11767 DataType::List(Arc::new(arrow_schema::Field::new(
11768 "item",
11769 DataType::Int32,
11770 true,
11771 ))),
11772 true,
11773 );
11774
11775 let page_v21 =
11776 encode_first_page(field.clone(), arr.clone(), TestEncoding::StructuralU16).await;
11777 let PageEncoding::Structural(layout_v21) = &page_v21.description else {
11778 panic!("Expected structural encoding");
11779 };
11780 let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
11781 panic!("Expected constant layout");
11782 };
11783 assert!(layout_v21.rep_compression.is_none());
11784 assert!(layout_v21.def_compression.is_none());
11785 assert_eq!(layout_v21.num_rep_values, 0);
11786 assert_eq!(layout_v21.num_def_values, 0);
11787
11788 let page_v22 = encode_first_page(field, arr, TestEncoding::StructuralU32).await;
11789 let PageEncoding::Structural(layout_v22) = &page_v22.description else {
11790 panic!("Expected structural encoding");
11791 };
11792 let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
11793 panic!("Expected constant layout");
11794 };
11795 assert!(layout_v22.def_compression.is_some());
11796 assert!(layout_v22.num_def_values > 0);
11797 }
11798
11799 #[tokio::test]
11800 async fn test_complex_all_null_round_trip() {
11801 use arrow_array::ListArray;
11802
11803 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11804 (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
11805 );
11806
11807 let test_cases = TestCases::default().with_u32_structural_encodings();
11808 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
11809 .await;
11810 }
11811
11812 #[tokio::test]
11813 async fn test_complex_all_null_constant_def_round_trip() {
11814 use arrow_array::ListArray;
11815
11816 let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11819 (0..5000).map(|_| None::<Vec<Option<i32>>>),
11820 );
11821
11822 let test_cases = TestCases::default().with_u32_structural_encodings();
11823 check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
11824 .await;
11825 }
11826
11827 fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer {
11828 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
11829 data: LanceBuffer::reinterpret_slice(Arc::from(levels)),
11830 bits_per_value: 16,
11831 num_values: levels.len() as u64,
11832 block_info: BlockInfo::new(),
11833 });
11834 BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block)
11835 .unwrap()
11836 .0
11837 .unwrap()
11838 }
11839
11840 fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns {
11841 let frame = encoded_u16_frame(levels, run_length_width);
11842 RleDecompressor::with_run_length_width(16, run_length_width)
11843 .decode_u16_runs(frame, levels.len() as u64)
11844 .unwrap()
11845 }
11846
11847 #[test]
11848 fn miniblock_levels_use_one_count_for_mixed_codecs() {
11849 let actual_num_levels = usize::from(u16::MAX) + 8;
11850 let levels = vec![1_u16; actual_num_levels];
11851 let rep_frame = encoded_u16_frame(&levels, RunLengthWidth::U8);
11852 let rep_decompressor = RleDecompressor::new(16);
11853 let def_frame = LanceBuffer::reinterpret_slice(Arc::from(levels.clone()));
11854 let def_decompressor = ValueDecompressor::from_flat(&pb21::Flat {
11855 bits_per_value: 16,
11856 data: None,
11857 });
11858
11859 let num_levels = DecodeMiniBlockTask::resolve_num_levels(
11860 Some(&rep_decompressor),
11861 Some(&rep_frame),
11862 Some(&def_decompressor),
11863 Some(&def_frame),
11864 7,
11865 )
11866 .unwrap();
11867 assert_eq!(num_levels, actual_num_levels as u64);
11868
11869 let rep =
11870 DecodeMiniBlockTask::decode_levels(&rep_decompressor, rep_frame, num_levels).unwrap();
11871 let def =
11872 DecodeMiniBlockTask::decode_levels(&def_decompressor, def_frame, num_levels).unwrap();
11873 assert_eq!(rep.as_ref(), levels);
11874 assert_eq!(def, rep);
11875 }
11876
11877 #[test]
11878 fn miniblock_levels_reject_non_wrapped_count_mismatch() {
11879 let frame = encoded_u16_frame(&[1_u16; 8], RunLengthWidth::U8);
11880 let decompressor = RleDecompressor::new(16);
11881
11882 let error = DecodeMiniBlockTask::resolve_num_levels(
11883 Some(&decompressor),
11884 Some(&frame),
11885 None,
11886 None,
11887 7,
11888 )
11889 .unwrap_err();
11890 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
11891 assert!(error.to_string().contains("not congruent modulo 65536"));
11892 }
11893
11894 #[test]
11895 fn miniblock_levels_reject_cross_stream_count_disagreement() {
11896 let rep_levels = vec![1_u16; usize::from(u16::MAX) + 8];
11897 let def_levels = vec![1_u16; rep_levels.len() + usize::from(u16::MAX) + 1];
11898 let rep_frame = encoded_u16_frame(&rep_levels, RunLengthWidth::U8);
11899 let def_frame = encoded_u16_frame(&def_levels, RunLengthWidth::U8);
11900 let decompressor = RleDecompressor::new(16);
11901
11902 let error = DecodeMiniBlockTask::resolve_num_levels(
11903 Some(&decompressor),
11904 Some(&rep_frame),
11905 Some(&decompressor),
11906 Some(&def_frame),
11907 7,
11908 )
11909 .unwrap_err();
11910 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
11911 assert!(
11912 error
11913 .to_string()
11914 .contains("structural streams disagree on the level count")
11915 );
11916 }
11917
11918 fn physical_levels(levels: &[u16]) -> LazyLevels {
11919 LazyLevels::Runs(Arc::new(RunStorage::Physical(
11920 encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(),
11921 )))
11922 }
11923
11924 fn coalesced_levels(levels: &[u16]) -> LazyLevels {
11925 let mut values = Vec::new();
11926 let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len());
11927 for (index, &value) in levels.iter().enumerate() {
11928 if values.last() == Some(&value) {
11929 ends.set_last(index + 1).unwrap();
11930 } else {
11931 values.push(value);
11932 ends.push(index + 1).unwrap();
11933 }
11934 }
11935 LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
11936 values: values.into_boxed_slice(),
11937 ends: ends.finish(),
11938 }))
11939 }
11940
11941 #[test]
11942 fn lazy_levels_runs_match_dense() {
11943 let expanded: Vec<u16> = vec![3, 3, 1, 3, 3, 3, 0, 0];
11945 let coalesced = coalesced_levels(&expanded);
11946 let physical = physical_levels(&expanded);
11947 let dense = LazyLevels::Dense(ScalarBuffer::<u16>::from(expanded.clone()));
11948 let n = expanded.len();
11949
11950 assert_eq!(coalesced.len(), n);
11951 assert_eq!(physical.len(), n);
11952 assert_eq!(dense.len(), n);
11953
11954 let max_rep = 3u16;
11956 let row_starts: Vec<usize> = (0..n).filter(|&i| expanded[i] == max_rep).collect();
11957 for target in 0..=row_starts.len() as u64 {
11958 let want = row_starts.get(target as usize).copied().unwrap_or(n);
11959 for runs in [&coalesced, &physical] {
11960 let mut cursor = LevelCursor::default();
11961 assert_eq!(
11962 runs.seek_row_start(&mut cursor, target, max_rep).unwrap(),
11963 want,
11964 "seek_row_start({target})"
11965 );
11966 }
11967 let mut c_dense = LevelCursor::default();
11968 assert_eq!(
11969 dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(),
11970 want
11971 );
11972 }
11973
11974 for start in 0..=n {
11977 for end in start..=n {
11978 for max in [0u16, 1, 2, 3] {
11979 let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64;
11980 for runs in [&coalesced, &physical] {
11981 let mut cursor = RunPosition::default();
11982 assert_eq!(
11983 runs.count_le_cursor(&mut cursor, start..end, max).0,
11984 want,
11985 "count_le_cursor({start}..{end}, {max})"
11986 );
11987 }
11988 let mut d_cur = RunPosition::default();
11989 assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want);
11990 }
11991 for runs in [&coalesced, &physical] {
11992 let mut got = Vec::new();
11993 runs.extend_into(start..end, RunPosition::default(), &mut got);
11994 assert_eq!(
11995 got,
11996 expanded[start..end].to_vec(),
11997 "extend_into({start}..{end})"
11998 );
11999 }
12000 let mut got_dense = Vec::new();
12001 dense.extend_into(start..end, RunPosition::default(), &mut got_dense);
12002 assert_eq!(got_dense, expanded[start..end].to_vec());
12003 }
12004 }
12005 }
12006
12007 #[test]
12008 fn physical_run_hints_support_deferred_materialization() {
12009 let expanded: Vec<u16> = vec![3, 3, 1, 1, 2, 2, 0, 0];
12010 let physical = physical_levels(&expanded);
12011 let LazyLevels::Runs(runs) = &physical else {
12012 panic!("expected physical runs");
12013 };
12014 let mut first_hint = RunPosition::default();
12015 runs.seek(&mut first_hint, 2);
12016 let mut second_hint = RunPosition::default();
12017 runs.seek(&mut second_hint, 6);
12018
12019 let mut second = Vec::new();
12020 physical.extend_into(6..8, second_hint, &mut second);
12021 let mut first = Vec::new();
12022 physical.extend_into(2..4, first_hint, &mut first);
12023 assert_eq!(second, expanded[6..8]);
12024 assert_eq!(first, expanded[2..4]);
12025 }
12026
12027 mod complex_all_null_drain_parity {
12032 use std::ops::Range;
12033
12034 use arrow_buffer::ScalarBuffer;
12035 use proptest::prelude::*;
12036
12037 use super::super::{LazyLevels, LevelCursor, RunPosition};
12038 use super::{coalesced_levels, physical_levels};
12039 use crate::Result;
12040
12041 #[derive(Debug, Clone)]
12042 struct DrainInput {
12043 max_rep: u16,
12044 max_visible: u16,
12045 rep: Option<Vec<u16>>,
12046 def: Option<Vec<u16>>,
12047 ranges: Vec<Range<u64>>,
12048 }
12049
12050 fn dense_levels(levels: &[u16]) -> LazyLevels {
12051 LazyLevels::Dense(ScalarBuffer::from(levels.to_vec()))
12052 }
12053
12054 fn rle_levels(levels: &[u16]) -> LazyLevels {
12055 coalesced_levels(levels)
12056 }
12057
12058 fn seek(
12059 rep: Option<&LazyLevels>,
12060 cursor: &mut LevelCursor,
12061 row: u64,
12062 max_rep: u16,
12063 ) -> Result<usize> {
12064 match rep {
12065 Some(rep) => rep.seek_row_start(cursor, row, max_rep),
12066 None => {
12067 cursor.row = row;
12068 cursor.level = row as usize;
12069 Ok(row as usize)
12070 }
12071 }
12072 }
12073
12074 fn simulate_drain(
12077 rep: Option<&LazyLevels>,
12078 def: Option<&LazyLevels>,
12079 max_rep: u16,
12080 max_visible: u16,
12081 ranges: &[Range<u64>],
12082 ) -> Result<(Vec<Range<usize>>, u64)> {
12083 let mut rep_cursor = LevelCursor::default();
12084 let mut def_run_cursor = RunPosition::default();
12085 let mut slices: Vec<Range<usize>> = Vec::new();
12086 let mut visible = 0u64;
12087 for range in ranges {
12088 let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?;
12089 let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?;
12090 visible += match def {
12091 Some(def) => {
12092 def.count_le_cursor(
12093 &mut def_run_cursor,
12094 level_start..level_end,
12095 max_visible,
12096 )
12097 .0
12098 }
12099 None => (level_end - level_start) as u64,
12100 };
12101 match slices.last_mut() {
12102 Some(last) if last.end == level_start => last.end = level_end,
12103 _ => slices.push(level_start..level_end),
12104 }
12105 }
12106 Ok((slices, visible))
12107 }
12108
12109 fn reference_drain(
12111 rep: Option<&[u16]>,
12112 def: Option<&[u16]>,
12113 max_rep: u16,
12114 max_visible: u16,
12115 ranges: &[Range<u64>],
12116 ) -> (Vec<Range<usize>>, u64) {
12117 let total_levels = rep
12118 .map(|r| r.len())
12119 .or_else(|| def.map(|d| d.len()))
12120 .unwrap_or(0);
12121 let row_starts: Vec<usize> = match rep {
12123 Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(),
12124 None => (0..total_levels).collect(),
12125 };
12126 let level_of_row = |row: u64| {
12127 row_starts
12128 .get(row as usize)
12129 .copied()
12130 .unwrap_or(total_levels)
12131 };
12132
12133 let mut slices: Vec<Range<usize>> = Vec::new();
12134 let mut visible = 0u64;
12135 for range in ranges {
12136 let ls = level_of_row(range.start);
12137 let le = level_of_row(range.end);
12138 visible += match def {
12139 Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64,
12140 None => (le - ls) as u64,
12141 };
12142 match slices.last_mut() {
12143 Some(last) if last.end == ls => last.end = le,
12144 _ => slices.push(ls..le),
12145 }
12146 }
12147 (slices, visible)
12148 }
12149
12150 fn ranges_strategy(num_rows: u64) -> BoxedStrategy<Vec<Range<u64>>> {
12151 if num_rows == 0 {
12152 return Just(Vec::new()).boxed();
12153 }
12154 proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8)
12157 .prop_map(move |pairs| {
12158 let mut ranges = Vec::new();
12159 let mut pos = 0u64;
12160 for (gap, len) in pairs {
12161 pos = pos.saturating_add(gap);
12162 if pos >= num_rows {
12163 break;
12164 }
12165 let end = (pos + len).min(num_rows);
12166 ranges.push(pos..end);
12167 pos = end;
12168 }
12169 ranges
12170 })
12171 .boxed()
12172 }
12173
12174 fn drain_input() -> impl Strategy<Value = DrainInput> {
12175 (
12176 1u16..=3,
12177 0u16..=3,
12178 any::<bool>(),
12179 any::<bool>(),
12180 1usize..=48,
12181 )
12182 .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| {
12183 let has_def = has_def || !has_rep;
12186 let rep = if has_rep {
12187 proptest::collection::vec(0u16..=max_rep, len)
12188 .prop_map(move |mut v| {
12189 v[0] = max_rep;
12191 Some(v)
12192 })
12193 .boxed()
12194 } else {
12195 Just(None).boxed()
12196 };
12197 let def = if has_def {
12198 proptest::collection::vec(0u16..=(max_visible + 2), len)
12199 .prop_map(Some)
12200 .boxed()
12201 } else {
12202 Just(None).boxed()
12203 };
12204 (Just(max_rep), Just(max_visible), rep, def)
12205 })
12206 .prop_flat_map(|(max_rep, max_visible, rep, def)| {
12207 let num_rows = match &rep {
12208 Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64,
12209 None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0),
12210 };
12211 ranges_strategy(num_rows).prop_map(move |ranges| DrainInput {
12212 max_rep,
12213 max_visible,
12214 rep: rep.clone(),
12215 def: def.clone(),
12216 ranges,
12217 })
12218 })
12219 }
12220
12221 proptest! {
12222 #![proptest_config(ProptestConfig::with_cases(512))]
12223
12224 #[test]
12225 fn drain_matches_reference(input in drain_input()) {
12226 let DrainInput { max_rep, max_visible, rep, def, ranges } = input;
12227
12228 let reference =
12229 reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges);
12230
12231 let rep_dense = rep.as_deref().map(dense_levels);
12232 let def_dense = def.as_deref().map(dense_levels);
12233 let got_dense =
12234 simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges)
12235 .unwrap();
12236 prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference");
12237
12238 let rep_rle = rep.as_deref().map(rle_levels);
12239 let def_rle = def.as_deref().map(rle_levels);
12240 let got_rle =
12241 simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges)
12242 .unwrap();
12243 prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference");
12244
12245 let rep_physical = rep.as_deref().map(physical_levels);
12246 let def_physical = def.as_deref().map(physical_levels);
12247 let got_physical =
12248 simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges)
12249 .unwrap();
12250 prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference");
12251 }
12252 }
12253 }
12254
12255 #[test]
12256 fn lazy_levels_runs_are_compact() {
12257 let single_run = |n: usize| {
12258 let mut ends = RunEndsBuilder::with_capacity(n, 1);
12259 ends.push(n).unwrap();
12260 LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
12261 values: vec![1u16].into_boxed_slice(),
12262 ends: ends.finish(),
12263 }))
12264 };
12265 assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size());
12267 assert!(single_run(10_000_000).deep_size() < 100);
12268 assert_eq!(single_run(10_000_000).len(), 10_000_000);
12269 assert_eq!(
12271 LazyLevels::Dense(ScalarBuffer::<u16>::from(vec![1u16; 1000])).deep_size(),
12272 2000
12273 );
12274 }
12275
12276 #[test]
12277 fn lazy_levels_selects_smallest_representation() {
12278 let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8);
12279 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
12280
12281 let equal_size: Vec<u16> = std::iter::repeat_n(0, 256)
12282 .chain(std::iter::repeat_n(1, 100))
12283 .chain(std::iter::repeat_n(2, 100))
12284 .collect();
12285 let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8);
12286 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
12287
12288 let moderate_runs: Vec<u16> = (0..250)
12289 .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
12290 .collect();
12291 let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8);
12292 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
12293
12294 let split_constant = vec![7u16; 5000];
12295 let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8);
12296 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
12297
12298 let high_density: Vec<u16> = (0..70_000).map(|index| (index % 2) as u16).collect();
12299 let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32);
12300 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
12301 }
12302
12303 #[test]
12304 fn physical_runs_detach_from_large_encoded_frame() {
12305 let levels: Vec<u16> = (0..250)
12306 .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
12307 .collect();
12308 let frame = encoded_u16_frame(&levels, RunLengthWidth::U8);
12309 let frame_offset = 4096;
12310 let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000];
12311 allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref());
12312 let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len());
12313 let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8)
12314 .decode_u16_runs(frame, levels.len() as u64)
12315 .unwrap();
12316
12317 assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
12318 let cached = LazyLevels::from_rle_runs(runs).unwrap();
12319 assert!(
12320 matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_)))
12321 );
12322 assert_eq!(cached.len(), levels.len());
12323 assert!(cached.deep_size() < 4096);
12324 }
12325
12326 #[test]
12327 fn complex_all_null_levels_reject_invalid_values_and_lengths() {
12328 let invalid_levels = vec![0u16, 3];
12329 for levels in [
12330 LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())),
12331 physical_levels(&invalid_levels),
12332 coalesced_levels(&invalid_levels),
12333 ] {
12334 let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err();
12335 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12336 assert!(error.to_string().contains("Invalid definition level 3"));
12337 }
12338
12339 let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2])));
12340 let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16])));
12341 let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err();
12342 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12343 assert!(
12344 error
12345 .to_string()
12346 .contains("repetition has 2, definition has 1")
12347 );
12348 }
12349
12350 #[test]
12351 fn block_levels_reject_malformed_payload_size() {
12352 let block = DataBlock::FixedWidth(FixedWidthDataBlock {
12353 data: LanceBuffer::from(vec![0]),
12354 bits_per_value: 16,
12355 num_values: 1,
12356 block_info: BlockInfo::new(),
12357 });
12358 let error = dense_levels_from_block(block, 1, "definition").unwrap_err();
12359 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12360 assert!(
12361 error
12362 .to_string()
12363 .contains("expected 2 bytes for 1 values, got 1")
12364 );
12365 }
12366
12367 #[test]
12368 fn complex_all_null_level_codec_validates_rle_metadata() {
12369 let encoding = pb21::CompressiveEncoding {
12370 compression: Some(Compression::Rle(Box::new(pb21::Rle {
12371 values: None,
12372 run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))),
12373 }))),
12374 };
12375
12376 let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default())
12377 .unwrap_err();
12378 assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12379 assert!(
12380 error
12381 .to_string()
12382 .contains("RLE compression missing values encoding")
12383 );
12384 }
12385
12386 #[tokio::test]
12388 async fn test_sparse_boolean_list_roundtrip() {
12389 use arrow_array::builder::{BooleanBuilder, ListBuilder};
12390
12391 let mut list_builder = ListBuilder::new(BooleanBuilder::new());
12392 for i in 0..1000i32 {
12393 if i % 64 == 0 {
12394 list_builder.values().append_value(i % 128 == 0);
12396 list_builder.append(true);
12397 } else {
12398 list_builder.append(false);
12399 }
12400 }
12401 let list_array = Arc::new(list_builder.finish());
12402
12403 let test_cases = TestCases::default().with_structural_encodings();
12404 check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
12405 }
12406
12407 fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
12408 use crate::compression::VariablePerValueDecompressor;
12409 use crate::encodings::physical::binary::VariableDecoder;
12410 use crate::repdef::{ControlWordParser, DefinitionInterpretation};
12411 use std::sync::Arc;
12412 Arc::new(super::FullZipDecodeDetails {
12413 value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
12414 VariableDecoder::default(),
12415 )
12416 as Arc<dyn VariablePerValueDecompressor>),
12417 def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
12418 ctrl_word_parser: ControlWordParser::new(0, 0),
12419 max_rep: 0,
12420 max_visible_def: 0,
12421 })
12422 }
12423
12424 fn decode_variable_full_zip(
12425 buf: Vec<u8>,
12426 bits_per_offset: u8,
12427 ) -> lance_core::Result<super::VariableFullZipDecoder> {
12428 use std::collections::VecDeque;
12429 let mut data = VecDeque::new();
12430 data.push_back(crate::buffer::LanceBuffer::from(buf));
12431 super::VariableFullZipDecoder::new(
12432 truncated_tail_details(),
12433 data,
12434 1,
12435 bits_per_offset,
12436 bits_per_offset,
12437 )
12438 }
12439
12440 #[test]
12442 fn variable_full_zip_wellformed_length_prefix() {
12443 assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok());
12444 assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok());
12445 }
12446
12447 #[test]
12456 fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
12457 use lance_core::Error;
12458
12459 for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
12460 let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
12461 .expect_err("a truncated length prefix must not decode");
12462 assert!(
12463 matches!(err, Error::CorruptFile { .. }),
12464 "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
12465 bits,
12466 buf_len,
12467 err
12468 );
12469 let msg = err.to_string();
12470 assert!(
12471 msg.contains("truncated length prefix"),
12472 "error should say what is wrong, got: {msg}"
12473 );
12474 }
12475 }
12476
12477 fn drain_fixed_full_zip(buf: Vec<u8>, bits_rep: u8, num_rows: u64) -> lance_core::Result<()> {
12482 use crate::compression::FixedPerValueDecompressor;
12483 use crate::decoder::StructuralPageDecoder;
12484 use crate::repdef::{ControlWordParser, DefinitionInterpretation};
12485 use std::collections::VecDeque;
12486 use std::sync::Arc;
12487
12488 #[derive(Debug)]
12489 struct UnusedFixedDecompressor;
12490
12491 impl FixedPerValueDecompressor for UnusedFixedDecompressor {
12492 fn decompress(&self, _: FixedWidthDataBlock, _: u64) -> crate::Result<DataBlock> {
12493 unreachable!("these cases fail while slicing, before any value is decoded")
12494 }
12495
12496 fn bits_per_value(&self) -> u64 {
12497 32
12498 }
12499 }
12500
12501 let details = Arc::new(super::FullZipDecodeDetails {
12502 value_decompressor: super::PerValueDecompressor::Fixed(Arc::new(
12503 UnusedFixedDecompressor,
12504 )),
12505 def_meaning: vec![DefinitionInterpretation::AllValidList].into(),
12506 ctrl_word_parser: ControlWordParser::new(bits_rep, 0),
12507 max_rep: 1,
12508 max_visible_def: 0,
12509 });
12510 let bytes_per_word = bits_rep.div_ceil(8) as usize;
12511
12512 let mut data = VecDeque::new();
12513 data.push_back(crate::buffer::LanceBuffer::from(buf));
12514 let mut decoder = super::FixedFullZipDecoder {
12515 details,
12516 data,
12517 offset_in_current: 0,
12518 bytes_per_value: 4,
12519 total_bytes_per_value: 4 + bytes_per_word,
12520 num_rows,
12521 };
12522 decoder.drain(num_rows).map(|_| ())
12523 }
12524
12525 #[test]
12527 fn fixed_full_zip_wellformed_page_slices() {
12528 assert!(drain_fixed_full_zip(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 2).is_ok());
12529 }
12530
12531 #[rstest::rstest]
12540 #[case::truncated_item(vec![1, 0, 0, 0, 0, 1, 0, 0, 0], 1, 2, "truncated item")]
12542 #[case::truncated_control_word(
12544 vec![1, 0, 0, 0, 0, 0, 1],
12545 9,
12546 2,
12547 "truncated control word"
12548 )]
12549 #[case::page_exhausted(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 3, "page data ran out")]
12551 fn fixed_full_zip_malformed_page_is_corrupt_file(
12552 #[case] buf: Vec<u8>,
12553 #[case] bits_rep: u8,
12554 #[case] num_rows: u64,
12555 #[case] expected_msg: &str,
12556 ) {
12557 use lance_core::Error;
12558
12559 let err = drain_fixed_full_zip(buf, bits_rep, num_rows)
12560 .expect_err("a malformed page must not slice");
12561 assert!(
12562 matches!(err, Error::CorruptFile { .. }),
12563 "expected CorruptFile, got: {err:?}"
12564 );
12565 let msg = err.to_string();
12566 assert!(
12567 msg.contains(expected_msg),
12568 "error should say what is wrong, got: {msg}"
12569 );
12570 }
12571}