Skip to main content

lance_encoding/encodings/logical/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    any::Any,
6    collections::{HashMap, VecDeque},
7    env,
8    fmt::Debug,
9    iter,
10    ops::Range,
11    sync::Arc,
12    vec,
13};
14
15use crate::{
16    constants::{
17        STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK,
18    },
19    data::DictionaryDataBlock,
20    encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler},
21    format::{
22        ProtobufUtils21,
23        pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression},
24    },
25};
26use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type};
27use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
28use arrow_schema::{DataType, Field as ArrowField};
29use bytes::Bytes;
30use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
31use itertools::Itertools;
32use lance_arrow::DataTypeExt;
33use lance_arrow::deepcopy::deep_copy_nulls;
34use lance_core::{
35    cache::{CacheKey, Context, DeepSizeOf},
36    error::{Error, LanceOptionExt},
37    utils::bit::pad_bytes,
38};
39use log::trace;
40
41use crate::{
42    compression::{
43        BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
44    },
45    data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
46    utils::bytepack::BytepackedIntegerEncoder,
47};
48use crate::{
49    compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
50    encodings::logical::primitive::fullzip::PerValueDataBlock,
51};
52use crate::{
53    encodings::logical::primitive::miniblock::MiniBlockChunk, utils::bytepack::ByteUnpacker,
54};
55use crate::{
56    encodings::logical::primitive::miniblock::MiniBlockCompressed,
57    statistics::{ComputeStat, GetStat, Stat},
58};
59use crate::{
60    repdef::{
61        CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
62        RepDefSlicer, build_control_word_iterator,
63    },
64    utils::accumulation::AccumulationQueue,
65};
66use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
67
68use crate::constants::{
69    COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
70    DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
71    DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
72    DICT_VALUES_COMPRESSION_META_KEY,
73};
74use crate::version::LanceFileVersion;
75use crate::{
76    EncodingsIo,
77    buffer::LanceBuffer,
78    data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
79    decoder::{
80        ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
81        MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
82        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
83        StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
84    },
85    encoder::{
86        EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
87    },
88    repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
89};
90
91pub mod blob;
92pub mod constant;
93pub mod dict;
94pub mod fullzip;
95pub mod miniblock;
96
97const FILL_BYTE: u8 = 0xFE;
98const DEFAULT_DICT_DIVISOR: u64 = 2;
99const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
100const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
101const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
102
103struct PageLoadTask {
104    decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
105    num_rows: u64,
106}
107
108/// A trait for figuring out how to schedule the data within
109/// a single page.
110trait StructuralPageScheduler: std::fmt::Debug + Send {
111    /// Fetches any metadata required for the page
112    fn initialize<'a>(
113        &'a mut self,
114        io: &Arc<dyn EncodingsIo>,
115    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
116    /// Loads metadata from a previous initialize call
117    fn load(&mut self, data: &Arc<dyn CachedPageData>);
118    /// Schedules the read of the given ranges in the page
119    ///
120    /// The read may be split into multiple "shards" if the page is extremely large.
121    /// Each shard maps to one or more rows and can be decoded independently.
122    ///
123    /// Note: this sharding is for splitting up very large pages into smaller reads to
124    /// avoid buffering too much data in memory.  It is not related to the batch size or
125    /// compute units in any way.
126    fn schedule_ranges(
127        &self,
128        ranges: &[Range<u64>],
129        io: &Arc<dyn EncodingsIo>,
130    ) -> Result<Vec<PageLoadTask>>;
131}
132
133/// Metadata describing the decoded size of a mini-block
134#[derive(Debug)]
135struct ChunkMeta {
136    num_values: u64,
137    chunk_size_bytes: u64,
138    offset_bytes: u64,
139}
140
141/// A mini-block chunk that has been decoded and decompressed
142#[derive(Debug, Clone)]
143struct DecodedMiniBlockChunk {
144    rep: Option<ScalarBuffer<u16>>,
145    def: Option<ScalarBuffer<u16>>,
146    values: DataBlock,
147}
148
149/// A task to decode a one or more mini-blocks of data into an output batch
150///
151/// Note: Two batches might share the same mini-block of data.  When this happens
152/// then each batch gets a copy of the block and each batch decodes the block independently.
153///
154/// This means we have duplicated work but it is necessary to avoid having to synchronize
155/// the decoding of the block. (TODO: test this theory)
156#[derive(Debug)]
157struct DecodeMiniBlockTask {
158    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
159    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
160    value_decompressor: Arc<dyn MiniBlockDecompressor>,
161    dictionary_data: Option<Arc<DataBlock>>,
162    def_meaning: Arc<[DefinitionInterpretation]>,
163    num_buffers: u64,
164    max_visible_level: u16,
165    instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
166    has_large_chunk: bool,
167}
168
169impl DecodeMiniBlockTask {
170    fn decode_levels(
171        rep_decompressor: &dyn BlockDecompressor,
172        levels: LanceBuffer,
173        num_levels: u16,
174    ) -> Result<ScalarBuffer<u16>> {
175        let rep = rep_decompressor.decompress(levels, num_levels as u64)?;
176        let rep = rep.as_fixed_width().unwrap();
177        debug_assert_eq!(rep.num_values, num_levels as u64);
178        debug_assert_eq!(rep.bits_per_value, 16);
179        Ok(rep.data.borrow_to_typed_slice::<u16>())
180    }
181
182    // We are building a LevelBuffer (levels) and want to copy into it `total_len`
183    // values from `level_buf` starting at `offset`.
184    //
185    // We need to handle both the case where `levels` is None (no nulls encountered
186    // yet) and the case where `level_buf` is None (the input we are copying from has
187    // no nulls)
188    fn extend_levels(
189        range: Range<u64>,
190        levels: &mut Option<LevelBuffer>,
191        level_buf: &Option<impl AsRef<[u16]>>,
192        dest_offset: usize,
193    ) {
194        if let Some(level_buf) = level_buf {
195            if levels.is_none() {
196                // This is the first non-empty def buf we've hit, fill in the past
197                // with 0 (valid)
198                let mut new_levels_vec =
199                    LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
200                new_levels_vec.extend(iter::repeat_n(0, dest_offset));
201                *levels = Some(new_levels_vec);
202            }
203            levels.as_mut().unwrap().extend(
204                level_buf.as_ref()[range.start as usize..range.end as usize]
205                    .iter()
206                    .copied(),
207            );
208        } else if let Some(levels) = levels {
209            let num_values = (range.end - range.start) as usize;
210            // This is an all-valid level_buf but we had nulls earlier and so we
211            // need to materialize it
212            levels.extend(iter::repeat_n(0, num_values));
213        }
214    }
215
216    /// Maps a range of rows to a range of items and a range of levels
217    ///
218    /// If there is no repetition information this just returns the range as-is.
219    ///
220    /// If there is repetition information then we need to do some work to figure out what
221    /// range of items corresponds to the requested range of rows.
222    ///
223    /// For example, if the data is [[1, 2, 3], [4, 5], [6, 7]] and the range is 1..2 (i.e. just row
224    /// 1) then the user actually wants items 3..5.  In the above case the rep levels would be:
225    ///
226    /// Idx: 0 1 2 3 4 5 6
227    /// Rep: 1 0 0 1 0 1 0
228    ///
229    /// So the start (1) maps to the second 1 (idx=3) and the end (2) maps to the third 1 (idx=5)
230    ///
231    /// If there are invisible items then we don't count them when calculating the range of items we
232    /// are interested in but we do count them when calculating the range of levels we are interested
233    /// in.  As a result we have to return both the item range (first return value) and the level range
234    /// (second return value).
235    ///
236    /// For example, if the data is [[1, 2, 3], [4, 5], NULL, [6, 7, 8]] and the range is 2..4 then the
237    /// user wants items 5..8 but they want levels 5..9.  In the above case the rep/def levels would be:
238    ///
239    /// Idx: 0 1 2 3 4 5 6 7 8
240    /// Rep: 1 0 0 1 0 1 1 0 0
241    /// Def: 0 0 0 0 0 1 0 0 0
242    /// Itm: 1 2 3 4 5 6 7 8
243    ///
244    /// Finally, we have to contend with the fact that chunks may or may not start with a "preamble" of
245    /// trailing values that finish up a list from the previous chunk.  In this case the first item does
246    /// not start at max_rep because it is a continuation of the previous chunk.  For our purposes we do
247    /// not consider this a "row" and so the range 0..1 will refer to the first row AFTER the preamble.
248    ///
249    /// We have a separate parameter (`preamble_action`) to control whether we want the preamble or not.
250    ///
251    /// Note that the "trailer" is considered a "row" and if we want it we should include it in the range.
252    fn map_range(
253        range: Range<u64>,
254        rep: Option<&impl AsRef<[u16]>>,
255        def: Option<&impl AsRef<[u16]>>,
256        max_rep: u16,
257        max_visible_def: u16,
258        // The total number of items (not rows) in the chunk.  This is not quite the same as
259        // rep.len() / def.len() because it doesn't count invisible items
260        total_items: u64,
261        preamble_action: PreambleAction,
262    ) -> (Range<u64>, Range<u64>) {
263        if let Some(rep) = rep {
264            let mut rep = rep.as_ref();
265            // If there is a preamble and we need to skip it then do that first.  The work is the same
266            // whether there is def information or not
267            let mut items_in_preamble = 0_u64;
268            let first_row_start = match preamble_action {
269                PreambleAction::Skip | PreambleAction::Take => {
270                    let first_row_start = if let Some(def) = def.as_ref() {
271                        let mut first_row_start = None;
272                        for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
273                            if *rep == max_rep {
274                                first_row_start = Some(idx as u64);
275                                break;
276                            }
277                            if *def <= max_visible_def {
278                                items_in_preamble += 1;
279                            }
280                        }
281                        first_row_start
282                    } else {
283                        let first_row_start =
284                            rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
285                        items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
286                        first_row_start
287                    };
288                    // It is possible for a chunk to be entirely partial values but if it is then it
289                    // should never show up as a preamble to skip
290                    if first_row_start.is_none() {
291                        assert!(preamble_action == PreambleAction::Take);
292                        return (0..total_items, 0..rep.len() as u64);
293                    }
294                    let first_row_start = first_row_start.unwrap();
295                    rep = &rep[first_row_start as usize..];
296                    first_row_start
297                }
298                PreambleAction::Absent => {
299                    debug_assert!(rep[0] == max_rep);
300                    0
301                }
302            };
303
304            // We hit this case when all we needed was the preamble
305            if range.start == range.end {
306                debug_assert!(preamble_action == PreambleAction::Take);
307                debug_assert!(items_in_preamble <= total_items);
308                return (0..items_in_preamble, 0..first_row_start);
309            }
310            assert!(range.start < range.end);
311
312            let mut rows_seen = 0;
313            let mut new_start = 0;
314            let mut new_levels_start = 0;
315
316            if let Some(def) = def {
317                let def = &def.as_ref()[first_row_start as usize..];
318
319                // range.start == 0 always maps to 0 (even with invis items), otherwise we need to walk
320                let mut lead_invis_seen = 0;
321
322                if range.start > 0 {
323                    if def[0] > max_visible_def {
324                        lead_invis_seen += 1;
325                    }
326                    for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
327                        if *rep == max_rep {
328                            rows_seen += 1;
329                            if rows_seen == range.start {
330                                new_start = idx as u64 + 1 - lead_invis_seen;
331                                new_levels_start = idx as u64 + 1;
332                                break;
333                            }
334                        }
335                        if *def > max_visible_def {
336                            lead_invis_seen += 1;
337                        }
338                    }
339                }
340
341                rows_seen += 1;
342
343                let mut new_end = u64::MAX;
344                let mut new_levels_end = rep.len() as u64;
345                let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
346                let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
347                for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
348                    .iter()
349                    .zip(&def[(new_levels_start + 1) as usize..])
350                    .enumerate()
351                {
352                    if *rep == max_rep {
353                        rows_seen += 1;
354                        if rows_seen == range.end + 1 {
355                            new_end = idx as u64 + new_start + 1 - tail_invis_seen;
356                            new_levels_end = idx as u64 + new_levels_start + 1;
357                            break;
358                        }
359                    }
360                    if *def > max_visible_def {
361                        tail_invis_seen += 1;
362                    }
363                }
364
365                if new_end == u64::MAX {
366                    new_levels_end = rep.len() as u64;
367                    let total_invis_seen = lead_invis_seen + tail_invis_seen;
368                    new_end = rep.len() as u64 - total_invis_seen;
369                }
370
371                assert_ne!(new_end, u64::MAX);
372
373                // Adjust for any skipped preamble
374                if preamble_action == PreambleAction::Skip {
375                    new_start += items_in_preamble;
376                    new_end += items_in_preamble;
377                    new_levels_start += first_row_start;
378                    new_levels_end += first_row_start;
379                } else if preamble_action == PreambleAction::Take {
380                    debug_assert_eq!(new_start, 0);
381                    debug_assert_eq!(new_levels_start, 0);
382                    new_end += items_in_preamble;
383                    new_levels_end += first_row_start;
384                }
385
386                debug_assert!(new_end <= total_items);
387                (new_start..new_end, new_levels_start..new_levels_end)
388            } else {
389                // Easy case, there are no invisible items, so we don't need to check for them
390                // The items range and levels range will be the same.  We do still need to walk
391                // the rep levels to find the row boundaries
392
393                // range.start == 0 always maps to 0, otherwise we need to walk
394                if range.start > 0 {
395                    for (idx, rep) in rep.iter().skip(1).enumerate() {
396                        if *rep == max_rep {
397                            rows_seen += 1;
398                            if rows_seen == range.start {
399                                new_start = idx as u64 + 1;
400                                break;
401                            }
402                        }
403                    }
404                }
405                let mut new_end = rep.len() as u64;
406                // range.end == max_items always maps to rep.len(), otherwise we need to walk
407                if range.end < total_items {
408                    for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
409                        if *rep == max_rep {
410                            rows_seen += 1;
411                            if rows_seen == range.end {
412                                new_end = idx as u64 + new_start + 1;
413                                break;
414                            }
415                        }
416                    }
417                }
418
419                // Adjust for any skipped preamble
420                if preamble_action == PreambleAction::Skip {
421                    new_start += first_row_start;
422                    new_end += first_row_start;
423                } else if preamble_action == PreambleAction::Take {
424                    debug_assert_eq!(new_start, 0);
425                    new_end += first_row_start;
426                }
427
428                debug_assert!(new_end <= total_items);
429                (new_start..new_end, new_start..new_end)
430            }
431        } else {
432            // No repetition info, easy case, just use the range as-is and the item
433            // and level ranges are the same
434            (range.clone(), range)
435        }
436    }
437
438    // read `num_buffers` buffer sizes from `buf` starting at `offset`
439    fn read_buffer_sizes<const LARGE: bool>(
440        buf: &[u8],
441        offset: &mut usize,
442        num_buffers: u64,
443    ) -> Vec<u32> {
444        let read_size = if LARGE { 4 } else { 2 };
445        (0..num_buffers)
446            .map(|_| {
447                let bytes = &buf[*offset..*offset + read_size];
448                let size = if LARGE {
449                    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
450                } else {
451                    // the buffer size is read from u16 but is stored as u32 after decoding for consistency
452                    u16::from_le_bytes([bytes[0], bytes[1]]) as u32
453                };
454                *offset += read_size;
455                size
456            })
457            .collect()
458    }
459
460    // Unserialize a miniblock into a collection of vectors
461    fn decode_miniblock_chunk(
462        &self,
463        buf: &LanceBuffer,
464        items_in_chunk: u64,
465    ) -> Result<DecodedMiniBlockChunk> {
466        let mut offset = 0;
467        let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
468        offset += 2;
469
470        let rep_size = if self.rep_decompressor.is_some() {
471            let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
472            offset += 2;
473            Some(rep_size)
474        } else {
475            None
476        };
477        let def_size = if self.def_decompressor.is_some() {
478            let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
479            offset += 2;
480            Some(def_size)
481        } else {
482            None
483        };
484
485        let buffer_sizes = if self.has_large_chunk {
486            Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
487        } else {
488            Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
489        };
490
491        offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
492
493        let rep = rep_size.map(|rep_size| {
494            let rep = buf.slice_with_length(offset, rep_size as usize);
495            offset += rep_size as usize;
496            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
497            rep
498        });
499
500        let def = def_size.map(|def_size| {
501            let def = buf.slice_with_length(offset, def_size as usize);
502            offset += def_size as usize;
503            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
504            def
505        });
506
507        let buffers = buffer_sizes
508            .into_iter()
509            .map(|buf_size| {
510                let buf = buf.slice_with_length(offset, buf_size as usize);
511                offset += buf_size as usize;
512                offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
513                buf
514            })
515            .collect::<Vec<_>>();
516
517        let values = self
518            .value_decompressor
519            .decompress(buffers, items_in_chunk)?;
520
521        let rep = rep
522            .map(|rep| {
523                Self::decode_levels(
524                    self.rep_decompressor.as_ref().unwrap().as_ref(),
525                    rep,
526                    num_levels,
527                )
528            })
529            .transpose()?;
530        let def = def
531            .map(|def| {
532                Self::decode_levels(
533                    self.def_decompressor.as_ref().unwrap().as_ref(),
534                    def,
535                    num_levels,
536                )
537            })
538            .transpose()?;
539
540        Ok(DecodedMiniBlockChunk { rep, def, values })
541    }
542}
543
544impl DecodePageTask for DecodeMiniBlockTask {
545    fn decode(self: Box<Self>) -> Result<DecodedPage> {
546        // First, we create output buffers for the rep and def and data
547        let mut repbuf: Option<LevelBuffer> = None;
548        let mut defbuf: Option<LevelBuffer> = None;
549
550        let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
551
552        // This is probably an over-estimate but it's quick and easy to calculate
553        let estimated_size_bytes = self
554            .instructions
555            .iter()
556            .map(|(_, chunk)| chunk.data.len())
557            .sum::<usize>()
558            * 2;
559        let mut data_builder =
560            DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
561
562        // We need to keep track of the offset into repbuf/defbuf that we are building up
563        let mut level_offset = 0;
564
565        // Pre-compute caching needs for each chunk by checking if the next chunk is the same
566        let needs_caching: Vec<bool> = self
567            .instructions
568            .windows(2)
569            .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
570            .chain(std::iter::once(false)) // the last one never needs caching
571            .collect();
572
573        // Cache for storing decoded chunks when beneficial
574        let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
575
576        // Now we iterate through each instruction and process it
577        for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
578            let should_cache_this_chunk = needs_caching[idx];
579
580            let decoded_chunk = match &chunk_cache {
581                Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
582                    // Clone only when we have a cache hit (much cheaper than decoding)
583                    cached_chunk.clone()
584                }
585                _ => {
586                    // Cache miss, need to decode
587                    let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
588
589                    // Only update cache if this chunk will benefit the next access
590                    if should_cache_this_chunk {
591                        chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
592                    }
593                    decoded
594                }
595            };
596
597            let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
598
599            // Our instructions tell us which rows we want to take from this chunk
600            let row_range_start =
601                instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
602            let row_range_end = row_range_start + instructions.rows_to_take;
603
604            // We use the rep info to map the row range to an item range / levels range
605            let (item_range, level_range) = Self::map_range(
606                row_range_start..row_range_end,
607                rep.as_ref(),
608                def.as_ref(),
609                max_rep,
610                self.max_visible_level,
611                chunk.items_in_chunk,
612                instructions.preamble_action,
613            );
614            if item_range.end - item_range.start > chunk.items_in_chunk {
615                return Err(lance_core::Error::internal(format!(
616                    "Item range {:?} is greater than chunk items in chunk {:?}",
617                    item_range, chunk.items_in_chunk
618                )));
619            }
620
621            // Now we append the data to the output buffers
622            Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
623            Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
624            level_offset += (level_range.end - level_range.start) as usize;
625            data_builder.append(&values, item_range);
626        }
627
628        let mut data = data_builder.finish();
629
630        let unraveler =
631            RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
632
633        if let Some(dictionary) = &self.dictionary_data {
634            // Don't decode here, that happens later (if needed)
635            let DataBlock::FixedWidth(indices) = data else {
636                return Err(lance_core::Error::internal(format!(
637                    "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
638                    data
639                )));
640            };
641            data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
642                indices,
643                dictionary.as_ref().clone(),
644            ));
645        }
646
647        Ok(DecodedPage {
648            data,
649            repdef: unraveler,
650        })
651    }
652}
653
654/// A chunk that has been loaded by the miniblock scheduler (but not
655/// yet decoded)
656#[derive(Debug)]
657struct LoadedChunk {
658    data: LanceBuffer,
659    items_in_chunk: u64,
660    byte_range: Range<u64>,
661    chunk_idx: usize,
662}
663
664impl Clone for LoadedChunk {
665    fn clone(&self) -> Self {
666        Self {
667            // Safe as we always create borrowed buffers here
668            data: self.data.clone(),
669            items_in_chunk: self.items_in_chunk,
670            byte_range: self.byte_range.clone(),
671            chunk_idx: self.chunk_idx,
672        }
673    }
674}
675
676/// Decodes mini-block formatted data.  See [`PrimitiveStructuralEncoder`] for more
677/// details on the different layouts.
678#[derive(Debug)]
679struct MiniBlockDecoder {
680    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
681    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
682    value_decompressor: Arc<dyn MiniBlockDecompressor>,
683    def_meaning: Arc<[DefinitionInterpretation]>,
684    loaded_chunks: VecDeque<LoadedChunk>,
685    instructions: VecDeque<ChunkInstructions>,
686    offset_in_current_chunk: u64,
687    num_rows: u64,
688    num_buffers: u64,
689    dictionary: Option<Arc<DataBlock>>,
690    has_large_chunk: bool,
691}
692
693/// See [`MiniBlockScheduler`] for more details on the scheduling and decoding
694/// process for miniblock encoded data.
695impl StructuralPageDecoder for MiniBlockDecoder {
696    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
697        let mut items_desired = num_rows;
698        let mut need_preamble = false;
699        let mut skip_in_chunk = self.offset_in_current_chunk;
700        let mut drain_instructions = Vec::new();
701        while items_desired > 0 || need_preamble {
702            let (instructions, consumed) = self
703                .instructions
704                .front()
705                .unwrap()
706                .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
707
708            while self.loaded_chunks.front().unwrap().chunk_idx
709                != instructions.chunk_instructions.chunk_idx
710            {
711                self.loaded_chunks.pop_front();
712            }
713            drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
714            if consumed {
715                self.instructions.pop_front();
716            }
717        }
718        // We can throw away need_preamble here because it must be false.  If it were true it would mean
719        // we were still in the middle of loading rows.  We do need to latch skip_in_chunk though.
720        self.offset_in_current_chunk = skip_in_chunk;
721
722        let max_visible_level = self
723            .def_meaning
724            .iter()
725            .take_while(|l| !l.is_list())
726            .map(|l| l.num_def_levels())
727            .sum::<u16>();
728
729        Ok(Box::new(DecodeMiniBlockTask {
730            instructions: drain_instructions,
731            def_decompressor: self.def_decompressor.clone(),
732            rep_decompressor: self.rep_decompressor.clone(),
733            value_decompressor: self.value_decompressor.clone(),
734            dictionary_data: self.dictionary.clone(),
735            def_meaning: self.def_meaning.clone(),
736            num_buffers: self.num_buffers,
737            max_visible_level,
738            has_large_chunk: self.has_large_chunk,
739        }))
740    }
741
742    fn num_rows(&self) -> u64 {
743        self.num_rows
744    }
745}
746
747#[derive(Debug)]
748struct CachedComplexAllNullState {
749    rep: Option<ScalarBuffer<u16>>,
750    def: Option<ScalarBuffer<u16>>,
751}
752
753impl DeepSizeOf for CachedComplexAllNullState {
754    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
755        self.rep.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
756            + self.def.as_ref().map(|buf| buf.len() * 2).unwrap_or(0)
757    }
758}
759
760impl CachedPageData for CachedComplexAllNullState {
761    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
762        self
763    }
764}
765
766/// A scheduler for all-null data that has repetition and definition levels
767///
768/// We still need to do some I/O in this case because we need to figure out what kind of null we
769/// are dealing with (null list, null struct, what level null struct, etc.)
770///
771/// TODO: Right now we just load the entire rep/def at initialization time and cache it.  This is a touch
772/// RAM aggressive and maybe we want something more lazy in the future.  On the other hand, it's simple
773/// and fast so...maybe not :)
774#[derive(Debug)]
775pub struct ComplexAllNullScheduler {
776    // Set from protobuf
777    buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
778    def_meaning: Arc<[DefinitionInterpretation]>,
779    repdef: Option<Arc<CachedComplexAllNullState>>,
780    max_rep: u16,
781    max_visible_level: u16,
782    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
783    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
784    num_rep_values: u64,
785    num_def_values: u64,
786}
787
788impl ComplexAllNullScheduler {
789    pub fn new(
790        buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
791        def_meaning: Arc<[DefinitionInterpretation]>,
792        rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
793        def_decompressor: Option<Arc<dyn BlockDecompressor>>,
794        num_rep_values: u64,
795        num_def_values: u64,
796    ) -> Self {
797        let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
798        let max_visible_level = def_meaning
799            .iter()
800            .take_while(|l| !l.is_list())
801            .map(|l| l.num_def_levels())
802            .sum::<u16>();
803        Self {
804            buffer_offsets_and_sizes,
805            def_meaning,
806            repdef: None,
807            max_rep,
808            max_visible_level,
809            rep_decompressor,
810            def_decompressor,
811            num_rep_values,
812            num_def_values,
813        }
814    }
815}
816
817impl StructuralPageScheduler for ComplexAllNullScheduler {
818    fn initialize<'a>(
819        &'a mut self,
820        io: &Arc<dyn EncodingsIo>,
821    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
822        // Fully load the rep & def buffers, as needed
823        let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
824        let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
825        let has_rep = rep_size > 0;
826        let has_def = def_size > 0;
827
828        let mut reads = Vec::with_capacity(2);
829        if has_rep {
830            reads.push(rep_pos..rep_pos + rep_size);
831        }
832        if has_def {
833            reads.push(def_pos..def_pos + def_size);
834        }
835
836        let data = io.submit_request(reads, 0);
837        let rep_decompressor = self.rep_decompressor.clone();
838        let def_decompressor = self.def_decompressor.clone();
839        let num_rep_values = self.num_rep_values;
840        let num_def_values = self.num_def_values;
841
842        async move {
843            let data = data.await?;
844            let mut data_iter = data.into_iter();
845
846            let decompress_levels = |compressed_bytes: Bytes,
847                                     decompressor: &Arc<dyn BlockDecompressor>,
848                                     num_values: u64,
849                                     level_type: &str|
850             -> Result<ScalarBuffer<u16>> {
851                let compressed_buffer = LanceBuffer::from_bytes(compressed_bytes, 1);
852                let decompressed = decompressor.decompress(compressed_buffer, num_values)?;
853                match decompressed {
854                    DataBlock::FixedWidth(block) => {
855                        if block.num_values != num_values {
856                            return Err(Error::invalid_input_source(format!(
857                                "Unexpected {} level count after decompression: expected {}, got {}",
858                                level_type, num_values, block.num_values
859                            )
860                            .into()));
861                        }
862                        if block.bits_per_value != 16 {
863                            return Err(Error::invalid_input_source(format!(
864                                "Unexpected {} level bit width after decompression: expected 16, got {}",
865                                level_type, block.bits_per_value
866                            )
867                            .into()));
868                        }
869                        Ok(block.data.borrow_to_typed_slice::<u16>())
870                    }
871                    _ => Err(Error::invalid_input_source(format!(
872                        "Expected fixed-width data block for {} levels",
873                        level_type
874                    )
875                    .into())),
876                }
877            };
878
879            let rep = if has_rep {
880                let rep = data_iter.next().unwrap();
881                if let Some(rep_decompressor) = rep_decompressor.as_ref() {
882                    Some(decompress_levels(
883                        rep,
884                        rep_decompressor,
885                        num_rep_values,
886                        "repetition",
887                    )?)
888                } else {
889                    let rep = LanceBuffer::from_bytes(rep, 2);
890                    let rep = rep.borrow_to_typed_slice::<u16>();
891                    Some(rep)
892                }
893            } else {
894                None
895            };
896
897            let def = if has_def {
898                let def = data_iter.next().unwrap();
899                if let Some(def_decompressor) = def_decompressor.as_ref() {
900                    Some(decompress_levels(
901                        def,
902                        def_decompressor,
903                        num_def_values,
904                        "definition",
905                    )?)
906                } else {
907                    let def = LanceBuffer::from_bytes(def, 2);
908                    let def = def.borrow_to_typed_slice::<u16>();
909                    Some(def)
910                }
911            } else {
912                None
913            };
914
915            let repdef = Arc::new(CachedComplexAllNullState { rep, def });
916
917            self.repdef = Some(repdef.clone());
918
919            Ok(repdef as Arc<dyn CachedPageData>)
920        }
921        .boxed()
922    }
923
924    fn load(&mut self, data: &Arc<dyn CachedPageData>) {
925        self.repdef = Some(
926            data.clone()
927                .as_arc_any()
928                .downcast::<CachedComplexAllNullState>()
929                .unwrap(),
930        );
931    }
932
933    fn schedule_ranges(
934        &self,
935        ranges: &[Range<u64>],
936        _io: &Arc<dyn EncodingsIo>,
937    ) -> Result<Vec<PageLoadTask>> {
938        let ranges = VecDeque::from_iter(ranges.iter().cloned());
939        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
940        let decoder = Box::new(ComplexAllNullPageDecoder {
941            ranges,
942            rep: self.repdef.as_ref().unwrap().rep.clone(),
943            def: self.repdef.as_ref().unwrap().def.clone(),
944            num_rows,
945            def_meaning: self.def_meaning.clone(),
946            max_rep: self.max_rep,
947            max_visible_level: self.max_visible_level,
948            cursor_row: 0,
949            cursor_level: 0,
950        }) as Box<dyn StructuralPageDecoder>;
951        let page_load_task = PageLoadTask {
952            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
953            num_rows,
954        };
955        Ok(vec![page_load_task])
956    }
957}
958
959#[derive(Debug)]
960pub struct ComplexAllNullPageDecoder {
961    ranges: VecDeque<Range<u64>>,
962    rep: Option<ScalarBuffer<u16>>,
963    def: Option<ScalarBuffer<u16>>,
964    num_rows: u64,
965    def_meaning: Arc<[DefinitionInterpretation]>,
966    max_rep: u16,
967    max_visible_level: u16,
968    cursor_row: u64,
969    cursor_level: usize,
970}
971
972impl ComplexAllNullPageDecoder {
973    fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
974        let mut rows_desired = num_rows;
975        let mut ranges = Vec::with_capacity(self.ranges.len());
976        while rows_desired > 0 {
977            let front = self.ranges.front_mut().unwrap();
978            let avail = front.end - front.start;
979            if avail > rows_desired {
980                ranges.push(front.start..front.start + rows_desired);
981                front.start += rows_desired;
982                rows_desired = 0;
983            } else {
984                ranges.push(self.ranges.pop_front().unwrap());
985                rows_desired -= avail;
986            }
987        }
988        ranges
989    }
990
991    fn take_row(&mut self) -> Result<(Range<usize>, u64)> {
992        let start = self.cursor_level;
993        let end = if let Some(rep) = &self.rep {
994            if start >= rep.len() {
995                return Err(Error::internal(
996                    "Invalid complex all-null layout: repetition buffer too short",
997                ));
998            }
999            if rep[start] != self.max_rep {
1000                return Err(Error::internal(
1001                    "Invalid complex all-null layout: row did not start at max repetition level",
1002                ));
1003            }
1004            let mut end = start + 1;
1005            while end < rep.len() && rep[end] != self.max_rep {
1006                end += 1;
1007            }
1008            end
1009        } else {
1010            start + 1
1011        };
1012
1013        let visible = if let Some(def) = &self.def {
1014            if end > def.len() {
1015                return Err(Error::internal(
1016                    "Invalid complex all-null layout: definition buffer too short",
1017                ));
1018            }
1019            def[start..end]
1020                .iter()
1021                .filter(|d| **d <= self.max_visible_level)
1022                .count() as u64
1023        } else {
1024            (end - start) as u64
1025        };
1026
1027        self.cursor_level = end;
1028        self.cursor_row += 1;
1029        Ok((start..end, visible))
1030    }
1031
1032    fn skip_to_row(&mut self, target_row: u64) -> Result<()> {
1033        while self.cursor_row < target_row {
1034            self.take_row()?;
1035        }
1036        Ok(())
1037    }
1038}
1039
1040impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1041    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1042        let drained_ranges = self.drain_ranges(num_rows);
1043        let mut level_slices: Vec<Range<usize>> = Vec::new();
1044        let mut visible_items_total = 0;
1045
1046        for range in drained_ranges {
1047            self.skip_to_row(range.start)?;
1048            for _ in range.start..range.end {
1049                let (level_range, visible) = self.take_row()?;
1050                visible_items_total += visible;
1051                if let Some(last) = level_slices.last_mut()
1052                    && last.end == level_range.start
1053                {
1054                    last.end = level_range.end;
1055                    continue;
1056                }
1057                level_slices.push(level_range);
1058            }
1059        }
1060
1061        Ok(Box::new(DecodeComplexAllNullTask {
1062            level_slices,
1063            visible_items_total,
1064            rep: self.rep.clone(),
1065            def: self.def.clone(),
1066            def_meaning: self.def_meaning.clone(),
1067            max_visible_level: self.max_visible_level,
1068        }))
1069    }
1070
1071    fn num_rows(&self) -> u64 {
1072        self.num_rows
1073    }
1074}
1075
1076/// We use `level_slices` to slice into `rep` and `def` and create rep/def buffers
1077/// for the null data.
1078#[derive(Debug)]
1079pub struct DecodeComplexAllNullTask {
1080    level_slices: Vec<Range<usize>>,
1081    visible_items_total: u64,
1082    rep: Option<ScalarBuffer<u16>>,
1083    def: Option<ScalarBuffer<u16>>,
1084    def_meaning: Arc<[DefinitionInterpretation]>,
1085    max_visible_level: u16,
1086}
1087
1088impl DecodeComplexAllNullTask {
1089    fn decode_level(&self, levels: &Option<ScalarBuffer<u16>>) -> Option<Vec<u16>> {
1090        levels.as_ref().map(|levels| {
1091            let num_levels = self
1092                .level_slices
1093                .iter()
1094                .map(|range| range.end - range.start)
1095                .sum();
1096            let mut referenced_levels = Vec::with_capacity(num_levels);
1097            for range in &self.level_slices {
1098                referenced_levels.extend(levels[range.start..range.end].iter().copied());
1099            }
1100            referenced_levels
1101        })
1102    }
1103}
1104
1105impl DecodePageTask for DecodeComplexAllNullTask {
1106    fn decode(self: Box<Self>) -> Result<DecodedPage> {
1107        let rep = self.decode_level(&self.rep);
1108        let def = self.decode_level(&self.def);
1109
1110        // If there are definition levels there may be empty / null lists which are not visible
1111        // in the items array.  We need to account for that here to figure out how many values
1112        // should be in the items array.
1113        let num_values = if let Some(def) = &def {
1114            def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
1115        } else {
1116            self.visible_items_total
1117        };
1118
1119        let data = DataBlock::AllNull(AllNullDataBlock { num_values });
1120        let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
1121        Ok(DecodedPage {
1122            data,
1123            repdef: unraveler,
1124        })
1125    }
1126}
1127
1128/// A scheduler for simple all-null data
1129///
1130/// "simple" all-null data is data that is all null and only has a single level of definition and
1131/// no repetition.  We don't need to read any data at all in this case.
1132#[derive(Debug, Default)]
1133pub struct SimpleAllNullScheduler {}
1134
1135impl StructuralPageScheduler for SimpleAllNullScheduler {
1136    fn initialize<'a>(
1137        &'a mut self,
1138        _io: &Arc<dyn EncodingsIo>,
1139    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1140        std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
1141    }
1142
1143    fn load(&mut self, _cache: &Arc<dyn CachedPageData>) {}
1144
1145    fn schedule_ranges(
1146        &self,
1147        ranges: &[Range<u64>],
1148        _io: &Arc<dyn EncodingsIo>,
1149    ) -> Result<Vec<PageLoadTask>> {
1150        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1151        let decoder =
1152            Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
1153        let page_load_task = PageLoadTask {
1154            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1155            num_rows,
1156        };
1157        Ok(vec![page_load_task])
1158    }
1159}
1160
1161/// A page decode task for all-null data without any
1162/// repetition and only a single level of definition
1163#[derive(Debug)]
1164struct SimpleAllNullDecodePageTask {
1165    num_values: u64,
1166}
1167impl DecodePageTask for SimpleAllNullDecodePageTask {
1168    fn decode(self: Box<Self>) -> Result<DecodedPage> {
1169        let unraveler = RepDefUnraveler::new(
1170            None,
1171            Some(vec![1; self.num_values as usize]),
1172            Arc::new([DefinitionInterpretation::NullableItem]),
1173            self.num_values,
1174        );
1175        Ok(DecodedPage {
1176            data: DataBlock::AllNull(AllNullDataBlock {
1177                num_values: self.num_values,
1178            }),
1179            repdef: unraveler,
1180        })
1181    }
1182}
1183
1184#[derive(Debug)]
1185pub struct SimpleAllNullPageDecoder {
1186    num_rows: u64,
1187}
1188
1189impl StructuralPageDecoder for SimpleAllNullPageDecoder {
1190    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1191        Ok(Box::new(SimpleAllNullDecodePageTask {
1192            num_values: num_rows,
1193        }))
1194    }
1195
1196    fn num_rows(&self) -> u64 {
1197        self.num_rows
1198    }
1199}
1200
1201#[derive(Debug, Clone)]
1202struct MiniBlockSchedulerDictionary {
1203    // These come from the protobuf
1204    dictionary_decompressor: Arc<dyn BlockDecompressor>,
1205    dictionary_buf_position_and_size: (u64, u64),
1206    dictionary_data_alignment: u64,
1207    num_dictionary_items: u64,
1208}
1209
1210/// Individual block metadata within a MiniBlock repetition index.
1211#[derive(Debug)]
1212struct MiniBlockRepIndexBlock {
1213    // The index of the first row that starts after the beginning of this block.  If the block
1214    // has a preamble this will be the row after the preamble.  If the block is entirely preamble
1215    // then this will be a row that starts in some future block.
1216    first_row: u64,
1217    // The number of rows in the block, including the trailer but not the preamble.
1218    // Can be 0 if the block is entirely preamble
1219    starts_including_trailer: u64,
1220    // Whether the block has a preamble
1221    has_preamble: bool,
1222    // Whether the block has a trailer
1223    has_trailer: bool,
1224}
1225
1226impl DeepSizeOf for MiniBlockRepIndexBlock {
1227    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
1228        0
1229    }
1230}
1231
1232/// Repetition index for MiniBlock encoding.
1233///
1234/// Stores block-level offset information to enable efficient random
1235/// access to nested data structures within mini-blocks.
1236#[derive(Debug)]
1237struct MiniBlockRepIndex {
1238    blocks: Vec<MiniBlockRepIndexBlock>,
1239}
1240
1241impl DeepSizeOf for MiniBlockRepIndex {
1242    fn deep_size_of_children(&self, context: &mut Context) -> usize {
1243        self.blocks.deep_size_of_children(context)
1244    }
1245}
1246
1247impl MiniBlockRepIndex {
1248    /// Decode repetition index from chunk metadata using default values.
1249    ///
1250    /// This creates a repetition index where each chunk has no partial values
1251    /// and no trailers, suitable for simple sequential data layouts.
1252    pub fn default_from_chunks(chunks: &[ChunkMeta]) -> Self {
1253        let mut blocks = Vec::with_capacity(chunks.len());
1254        let mut offset: u64 = 0;
1255
1256        for c in chunks {
1257            blocks.push(MiniBlockRepIndexBlock {
1258                first_row: offset,
1259                starts_including_trailer: c.num_values,
1260                has_preamble: false,
1261                has_trailer: false,
1262            });
1263
1264            offset += c.num_values;
1265        }
1266
1267        Self { blocks }
1268    }
1269
1270    /// Decode repetition index from raw bytes in little-endian format.
1271    ///
1272    /// The bytes should contain u64 values arranged in groups of `stride` elements,
1273    /// where the first two values of each group represent ends_count and partial_count.
1274    /// Returns an empty index if no bytes are provided.
1275    pub fn decode_from_bytes(rep_bytes: &[u8], stride: usize) -> Self {
1276        // Convert bytes to u64 slice, handling alignment automatically
1277        let buffer = crate::buffer::LanceBuffer::from(rep_bytes.to_vec());
1278        let u64_slice = buffer.borrow_to_typed_slice::<u64>();
1279        let n = u64_slice.len() / stride;
1280
1281        let mut blocks = Vec::with_capacity(n);
1282        let mut chunk_has_preamble = false;
1283        let mut offset: u64 = 0;
1284
1285        // Extract first two values from each block: ends_count and partial_count
1286        for i in 0..n {
1287            let base_idx = i * stride;
1288            let ends = u64_slice[base_idx];
1289            let partial = u64_slice[base_idx + 1];
1290
1291            let has_trailer = partial > 0;
1292            // Convert branches to arithmetic for better compiler optimization
1293            let starts_including_trailer =
1294                ends + (has_trailer as u64) - (chunk_has_preamble as u64);
1295
1296            blocks.push(MiniBlockRepIndexBlock {
1297                first_row: offset,
1298                starts_including_trailer,
1299                has_preamble: chunk_has_preamble,
1300                has_trailer,
1301            });
1302
1303            chunk_has_preamble = has_trailer;
1304            offset += starts_including_trailer;
1305        }
1306
1307        Self { blocks }
1308    }
1309}
1310
1311/// State that is loaded once and cached for future lookups
1312#[derive(Debug)]
1313struct MiniBlockCacheableState {
1314    /// Metadata that describes each chunk in the page
1315    chunk_meta: Vec<ChunkMeta>,
1316    /// The decoded repetition index
1317    rep_index: MiniBlockRepIndex,
1318    /// The dictionary for the page, if any
1319    dictionary: Option<Arc<DataBlock>>,
1320}
1321
1322impl DeepSizeOf for MiniBlockCacheableState {
1323    fn deep_size_of_children(&self, context: &mut Context) -> usize {
1324        self.rep_index.deep_size_of_children(context)
1325            + self
1326                .dictionary
1327                .as_ref()
1328                .map(|dict| dict.data_size() as usize)
1329                .unwrap_or(0)
1330    }
1331}
1332
1333impl CachedPageData for MiniBlockCacheableState {
1334    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1335        self
1336    }
1337}
1338
1339/// A scheduler for a page that has been encoded with the mini-block layout
1340///
1341/// Scheduling mini-block encoded data is simple in concept and somewhat complex
1342/// in practice.
1343///
1344/// First, during initialization, we load the chunk metadata, the repetition index,
1345/// and the dictionary (these last two may not be present)
1346///
1347/// Then, during scheduling, we use the user's requested row ranges and the repetition
1348/// index to determine which chunks we need and which rows we need from those chunks.
1349///
1350/// For example, if the repetition index is: [50, 3], [50, 0], [10, 0] and the range
1351/// from the user is 40..60 then we need to:
1352///
1353///  - Read the first chunk and skip the first 40 rows, then read 10 full rows, and
1354///    then read 3 items for the 11th row of our range.
1355///  - Read the second chunk and read the remaining items in our 11th row and then read
1356///    the remaining 9 full rows.
1357///
1358/// Then, if we are going to decode that in batches of 5, we need to make decode tasks.
1359/// The first two decode tasks will just need the first chunk.  The third decode task will
1360/// need the first chunk (for the trailer which has the 11th row in our range) and the second
1361/// chunk.  The final decode task will just need the second chunk.
1362///
1363/// The above prose descriptions are what are represented by `ChunkInstructions` and
1364/// `ChunkDrainInstructions`.
1365#[derive(Debug)]
1366pub struct MiniBlockScheduler {
1367    // These come from the protobuf
1368    buffer_offsets_and_sizes: Vec<(u64, u64)>,
1369    priority: u64,
1370    items_in_page: u64,
1371    repetition_index_depth: u16,
1372    num_buffers: u64,
1373    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
1374    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
1375    value_decompressor: Arc<dyn MiniBlockDecompressor>,
1376    def_meaning: Arc<[DefinitionInterpretation]>,
1377    dictionary: Option<MiniBlockSchedulerDictionary>,
1378    // This is set after initialization
1379    page_meta: Option<Arc<MiniBlockCacheableState>>,
1380    has_large_chunk: bool,
1381}
1382
1383impl MiniBlockScheduler {
1384    fn try_new(
1385        buffer_offsets_and_sizes: &[(u64, u64)],
1386        priority: u64,
1387        items_in_page: u64,
1388        layout: &pb21::MiniBlockLayout,
1389        decompressors: &dyn DecompressionStrategy,
1390    ) -> Result<Self> {
1391        let rep_decompressor = layout
1392            .rep_compression
1393            .as_ref()
1394            .map(|rep_compression| {
1395                decompressors
1396                    .create_block_decompressor(rep_compression)
1397                    .map(Arc::from)
1398            })
1399            .transpose()?;
1400        let def_decompressor = layout
1401            .def_compression
1402            .as_ref()
1403            .map(|def_compression| {
1404                decompressors
1405                    .create_block_decompressor(def_compression)
1406                    .map(Arc::from)
1407            })
1408            .transpose()?;
1409        let def_meaning = layout
1410            .layers
1411            .iter()
1412            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
1413            .collect::<Vec<_>>();
1414        let value_decompressor = decompressors.create_miniblock_decompressor(
1415            layout.value_compression.as_ref().unwrap(),
1416            decompressors,
1417        )?;
1418
1419        let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
1420            let num_dictionary_items = layout.num_dictionary_items;
1421            let dictionary_decompressor = decompressors
1422                .create_block_decompressor(dictionary_encoding)?
1423                .into();
1424            let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
1425            {
1426                Compression::Variable(_) => 4,
1427                Compression::Flat(_) => 16,
1428                Compression::General(_) => 1,
1429                Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
1430                    crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
1431                }
1432                _ => {
1433                    return Err(Error::invalid_input_source(
1434                        format!(
1435                            "Unsupported mini-block dictionary encoding: {:?}",
1436                            dictionary_encoding.compression.as_ref().unwrap()
1437                        )
1438                        .into(),
1439                    ));
1440                }
1441            };
1442            Some(MiniBlockSchedulerDictionary {
1443                dictionary_decompressor,
1444                dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
1445                dictionary_data_alignment,
1446                num_dictionary_items,
1447            })
1448        } else {
1449            None
1450        };
1451
1452        Ok(Self {
1453            buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
1454            rep_decompressor,
1455            def_decompressor,
1456            value_decompressor: value_decompressor.into(),
1457            repetition_index_depth: layout.repetition_index_depth as u16,
1458            num_buffers: layout.num_buffers,
1459            priority,
1460            items_in_page,
1461            dictionary,
1462            def_meaning: def_meaning.into(),
1463            page_meta: None,
1464            has_large_chunk: layout.has_large_chunk,
1465        })
1466    }
1467
1468    fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
1469        let page_meta = self.page_meta.as_ref().unwrap();
1470        chunk_indices
1471            .iter()
1472            .map(|&chunk_idx| {
1473                let chunk_meta = &page_meta.chunk_meta[chunk_idx];
1474                let bytes_start = chunk_meta.offset_bytes;
1475                let bytes_end = bytes_start + chunk_meta.chunk_size_bytes;
1476                LoadedChunk {
1477                    byte_range: bytes_start..bytes_end,
1478                    items_in_chunk: chunk_meta.num_values,
1479                    chunk_idx,
1480                    data: LanceBuffer::empty(),
1481                }
1482            })
1483            .collect()
1484    }
1485}
1486
1487#[derive(Debug, PartialEq, Eq, Clone, Copy)]
1488enum PreambleAction {
1489    Take,
1490    Skip,
1491    Absent,
1492}
1493
1494// When we schedule a chunk we use the repetition index (or, if none exists, just the # of items
1495// in each chunk) to map a user requested range into a set of ChunkInstruction objects which tell
1496// us how exactly to read from the chunk.
1497//
1498// Examples:
1499//
1500// | Chunk 0     | Chunk 1   | Chunk 2   | Chunk 3 |
1501// | xxxxyyyyzzz | zzzzzzzzz | zzzzzzzzz | aaabbcc |
1502//
1503// Full read (0..6)
1504//
1505// Chunk 0: (several rows, ends with trailer)
1506//   preamble: absent
1507//   rows_to_skip: 0
1508//   rows_to_take: 3 (x, y, z)
1509//   take_trailer: true
1510//
1511// Chunk 1: (all preamble, ends with trailer)
1512//   preamble: take
1513//   rows_to_skip: 0
1514//   rows_to_take: 0
1515//   take_trailer: true
1516//
1517// Chunk 2: (all preamble, no trailer)
1518//   preamble: take
1519//   rows_to_skip: 0
1520//   rows_to_take: 0
1521//   take_trailer: false
1522//
1523// Chunk 3: (several rows, no trailer or preamble)
1524//   preamble: absent
1525//   rows_to_skip: 0
1526//   rows_to_take: 3 (a, b, c)
1527//   take_trailer: false
1528#[derive(Clone, Debug, PartialEq, Eq)]
1529struct ChunkInstructions {
1530    // The index of the chunk to read
1531    chunk_idx: usize,
1532    // A "preamble" is when a chunk begins with a continuation of the previous chunk's list.  If there
1533    // is no repetition index there is never a preamble.
1534    //
1535    // It's possible for a chunk to be entirely premable.  For example, if there is a really large list
1536    // that spans several chunks.
1537    preamble: PreambleAction,
1538    // How many complete rows (not including the preamble or trailer) to skip
1539    //
1540    // If this is non-zero then premable must not be Take
1541    rows_to_skip: u64,
1542    // How many rows to take.  If a row splits across chunks then we will count the row in the first
1543    // chunk that contains the row.
1544    rows_to_take: u64,
1545    // A "trailer" is when a chunk ends with a partial list.  If there is no repetition index there is
1546    // never a trailer.
1547    //
1548    // A chunk that is all preamble may or may not have a trailer.
1549    //
1550    // If this is true then we want to include the trailer
1551    take_trailer: bool,
1552}
1553
1554// First, we schedule a bunch of [`ChunkInstructions`] based on the users ranges.  Then we
1555// start decoding them, based on a batch size, which might not align with what we scheduled.
1556//
1557// This results in `ChunkDrainInstructions` which targets a contiguous slice of a `ChunkInstructions`
1558//
1559// So if `ChunkInstructions` is "skip preamble, skip 10, take 50, take trailer" and we are decoding in
1560// batches of size 10 we might have a `ChunkDrainInstructions` that targets that chunk and has its own
1561// skip of 17 and take of 10.  This would mean we decode the chunk, skip the preamble and 27 rows, and
1562// then take 10 rows.
1563//
1564// One very confusing bit is that `rows_to_take` includes the trailer.  So if we have two chunks:
1565//  -no preamble, skip 5, take 10, take trailer
1566//  -take preamble, skip 0, take 50, no trailer
1567//
1568// and we are draining 20 rows then the drain instructions for the first batch will be:
1569//  - no preamble, skip 0 (from chunk 0), take 11 (from chunk 0)
1570//  - take preamble (from chunk 1), skip 0 (from chunk 1), take 9 (from chunk 1)
1571#[derive(Debug, PartialEq, Eq)]
1572struct ChunkDrainInstructions {
1573    chunk_instructions: ChunkInstructions,
1574    rows_to_skip: u64,
1575    rows_to_take: u64,
1576    preamble_action: PreambleAction,
1577}
1578
1579impl ChunkInstructions {
1580    // Given a repetition index and a set of user ranges we need to figure out how to read from the chunks
1581    //
1582    // We assume that `user_ranges` are in sorted order and non-overlapping
1583    //
1584    // The output will be a set of `ChunkInstructions` which tell us how to read from the chunks
1585    fn schedule_instructions(
1586        rep_index: &MiniBlockRepIndex,
1587        user_ranges: &[Range<u64>],
1588    ) -> Vec<Self> {
1589        // This is an in-exact capacity guess but pretty good.  The actual capacity can be
1590        // smaller if instructions are merged.  It can be larger if there are multiple instructions
1591        // per row which can happen with lists.
1592        let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
1593
1594        for user_range in user_ranges {
1595            let mut rows_needed = user_range.end - user_range.start;
1596            let mut need_preamble = false;
1597
1598            // Need to find the first chunk with a first row >= user_range.start.  If there are
1599            // multiple chunks with the same first row we need to take the first one.
1600            let mut block_index = match rep_index
1601                .blocks
1602                .binary_search_by_key(&user_range.start, |block| block.first_row)
1603            {
1604                Ok(idx) => {
1605                    // Slightly tricky case, we may need to walk backwards a bit to make sure we
1606                    // are grabbing first eligible chunk
1607                    let mut idx = idx;
1608                    while idx > 0 && rep_index.blocks[idx - 1].first_row == user_range.start {
1609                        idx -= 1;
1610                    }
1611                    idx
1612                }
1613                // Easy case.  idx is greater, and idx - 1 is smaller, so idx - 1 contains the start
1614                Err(idx) => idx - 1,
1615            };
1616
1617            let mut to_skip = user_range.start - rep_index.blocks[block_index].first_row;
1618
1619            while rows_needed > 0 || need_preamble {
1620                // Check if we've gone past the last block (should not happen)
1621                if block_index >= rep_index.blocks.len() {
1622                    log::warn!(
1623                        "schedule_instructions inconsistency: block_index >= rep_index.blocks.len(), exiting early"
1624                    );
1625                    break;
1626                }
1627
1628                let chunk = &rep_index.blocks[block_index];
1629                let rows_avail = chunk.starts_including_trailer.saturating_sub(to_skip);
1630
1631                // Handle blocks that are entirely preamble (rows_avail = 0)
1632                // These blocks have no rows to take but may have a preamble we need
1633                // We only look for preamble if to_skip == 0 (we're not skipping rows)
1634                if rows_avail == 0 && to_skip == 0 {
1635                    // Only process if this chunk has a preamble we need
1636                    if chunk.has_preamble && need_preamble {
1637                        chunk_instructions.push(Self {
1638                            chunk_idx: block_index,
1639                            preamble: PreambleAction::Take,
1640                            rows_to_skip: 0,
1641                            rows_to_take: 0,
1642                            // We still need to look at has_trailer to distinguish between "all preamble
1643                            // and row ends at end of chunk" and "all preamble and row bleeds into next
1644                            // chunk".  Both cases will have 0 rows available.
1645                            take_trailer: chunk.has_trailer,
1646                        });
1647                        // Only set need_preamble = false if the chunk has at least one row,
1648                        // Or we are reaching the last block,
1649                        // Otherwise, the chunk is entirely preamble and we need the next chunk's preamble too
1650                        if chunk.starts_including_trailer > 0
1651                            || block_index == rep_index.blocks.len() - 1
1652                        {
1653                            need_preamble = false;
1654                        }
1655                    }
1656                    // Move to next block
1657                    block_index += 1;
1658                    continue;
1659                }
1660
1661                // Edge case: if rows_avail == 0 but to_skip > 0
1662                // This theoretically shouldn't happen (binary search should avoid it)
1663                // but handle it for safety
1664                if rows_avail == 0 && to_skip > 0 {
1665                    // This block doesn't have enough rows to skip, move to next block
1666                    // Adjust to_skip by the number of rows in this block
1667                    to_skip -= chunk.starts_including_trailer;
1668                    block_index += 1;
1669                    continue;
1670                }
1671
1672                let rows_to_take = rows_avail.min(rows_needed);
1673                rows_needed -= rows_to_take;
1674
1675                let mut take_trailer = false;
1676                let preamble = if chunk.has_preamble {
1677                    if need_preamble {
1678                        PreambleAction::Take
1679                    } else {
1680                        PreambleAction::Skip
1681                    }
1682                } else {
1683                    PreambleAction::Absent
1684                };
1685
1686                // Are we taking the trailer?  If so, make sure we mark that we need the preamble
1687                if rows_to_take == rows_avail && chunk.has_trailer {
1688                    take_trailer = true;
1689                    need_preamble = true;
1690                } else {
1691                    need_preamble = false;
1692                };
1693
1694                chunk_instructions.push(Self {
1695                    preamble,
1696                    chunk_idx: block_index,
1697                    rows_to_skip: to_skip,
1698                    rows_to_take,
1699                    take_trailer,
1700                });
1701
1702                to_skip = 0;
1703                block_index += 1;
1704            }
1705        }
1706
1707        // If there were multiple ranges we may have multiple instructions for a single chunk.  Merge them now if they
1708        // are _adjacent_ (i.e. don't merge "take first row of chunk 0" and "take third row of chunk 0" into "take 2
1709        // rows of chunk 0 starting at 0")
1710        if user_ranges.len() > 1 {
1711            // TODO: Could probably optimize this allocation away
1712            let mut merged_instructions = Vec::with_capacity(chunk_instructions.len());
1713            let mut instructions_iter = chunk_instructions.into_iter();
1714            merged_instructions.push(instructions_iter.next().unwrap());
1715            for instruction in instructions_iter {
1716                let last = merged_instructions.last_mut().unwrap();
1717                if last.chunk_idx == instruction.chunk_idx
1718                    && last.rows_to_take + last.rows_to_skip == instruction.rows_to_skip
1719                {
1720                    last.rows_to_take += instruction.rows_to_take;
1721                    last.take_trailer |= instruction.take_trailer;
1722                } else {
1723                    merged_instructions.push(instruction);
1724                }
1725            }
1726            merged_instructions
1727        } else {
1728            chunk_instructions
1729        }
1730    }
1731
1732    fn drain_from_instruction(
1733        &self,
1734        rows_desired: &mut u64,
1735        need_preamble: &mut bool,
1736        skip_in_chunk: &mut u64,
1737    ) -> (ChunkDrainInstructions, bool) {
1738        // If we need the premable then we shouldn't be skipping anything
1739        debug_assert!(!*need_preamble || *skip_in_chunk == 0);
1740        let rows_avail = self.rows_to_take - *skip_in_chunk;
1741        let has_preamble = self.preamble != PreambleAction::Absent;
1742        let preamble_action = match (*need_preamble, has_preamble) {
1743            (true, true) => PreambleAction::Take,
1744            (true, false) => panic!("Need preamble but there isn't one"),
1745            (false, true) => PreambleAction::Skip,
1746            (false, false) => PreambleAction::Absent,
1747        };
1748
1749        // How many rows are we actually taking in this take step (including the preamble
1750        // and trailer both as individual rows)
1751        let rows_taking = if *rows_desired >= rows_avail {
1752            // We want all the rows.  If there is a trailer we are grabbing it and will need
1753            // the preamble of the next chunk
1754            // If there is a trailer and we are taking all the rows then we need the preamble
1755            // of the next chunk.
1756            //
1757            // Also, if this chunk is entirely preamble (rows_avail == 0 && !take_trailer) then we
1758            // need the preamble of the next chunk.
1759            *need_preamble = self.take_trailer;
1760            rows_avail
1761        } else {
1762            // We aren't taking all the rows.  Even if there is a trailer we aren't taking
1763            // it so we will not need the preamble
1764            *need_preamble = false;
1765            *rows_desired
1766        };
1767        let rows_skipped = *skip_in_chunk;
1768
1769        // Update the state for the next iteration
1770        let consumed_chunk = if *rows_desired >= rows_avail {
1771            *rows_desired -= rows_avail;
1772            *skip_in_chunk = 0;
1773            true
1774        } else {
1775            *skip_in_chunk += *rows_desired;
1776            *rows_desired = 0;
1777            false
1778        };
1779
1780        (
1781            ChunkDrainInstructions {
1782                chunk_instructions: self.clone(),
1783                rows_to_skip: rows_skipped,
1784                rows_to_take: rows_taking,
1785                preamble_action,
1786            },
1787            consumed_chunk,
1788        )
1789    }
1790}
1791
1792enum Words {
1793    U16(ScalarBuffer<u16>),
1794    U32(ScalarBuffer<u32>),
1795}
1796
1797struct WordsIter<'a> {
1798    iter: Box<dyn Iterator<Item = u32> + 'a>,
1799}
1800
1801impl Words {
1802    pub fn len(&self) -> usize {
1803        match self {
1804            Self::U16(b) => b.len(),
1805            Self::U32(b) => b.len(),
1806        }
1807    }
1808
1809    pub fn iter(&self) -> WordsIter<'_> {
1810        match self {
1811            Self::U16(buf) => WordsIter {
1812                iter: Box::new(buf.iter().map(|&x| x as u32)),
1813            },
1814            Self::U32(buf) => WordsIter {
1815                iter: Box::new(buf.iter().copied()),
1816            },
1817        }
1818    }
1819
1820    pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
1821        let bytes_per_value = if has_large_chunk { 4 } else { 2 };
1822        assert_eq!(bytes.len() % bytes_per_value, 0);
1823        let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
1824        if has_large_chunk {
1825            Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
1826        } else {
1827            Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
1828        }
1829    }
1830}
1831
1832impl<'a> Iterator for WordsIter<'a> {
1833    type Item = u32;
1834
1835    fn next(&mut self) -> Option<Self::Item> {
1836        self.iter.next()
1837    }
1838}
1839
1840impl StructuralPageScheduler for MiniBlockScheduler {
1841    fn initialize<'a>(
1842        &'a mut self,
1843        io: &Arc<dyn EncodingsIo>,
1844    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1845        // We always need to fetch chunk metadata.  We may also need to fetch a dictionary and
1846        // we may also need to fetch the repetition index.  Here, we gather what buffers we
1847        // need.
1848        let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
1849        let value_buf_position = self.buffer_offsets_and_sizes[1].0;
1850        let mut bufs_needed = 1;
1851        if self.dictionary.is_some() {
1852            bufs_needed += 1;
1853        }
1854        if self.repetition_index_depth > 0 {
1855            bufs_needed += 1;
1856        }
1857        let mut required_ranges = Vec::with_capacity(bufs_needed);
1858        required_ranges.push(meta_buf_position..meta_buf_position + meta_buf_size);
1859        if let Some(ref dictionary) = self.dictionary {
1860            required_ranges.push(
1861                dictionary.dictionary_buf_position_and_size.0
1862                    ..dictionary.dictionary_buf_position_and_size.0
1863                        + dictionary.dictionary_buf_position_and_size.1,
1864            );
1865        }
1866        if self.repetition_index_depth > 0 {
1867            let (rep_index_pos, rep_index_size) = self.buffer_offsets_and_sizes.last().unwrap();
1868            required_ranges.push(*rep_index_pos..*rep_index_pos + *rep_index_size);
1869        }
1870        let io_req = io.submit_request(required_ranges, 0);
1871
1872        async move {
1873            let mut buffers = io_req.await?.into_iter().fuse();
1874            let meta_bytes = buffers.next().unwrap();
1875            let dictionary_bytes = self.dictionary.as_ref().and_then(|_| buffers.next());
1876            let rep_index_bytes = buffers.next();
1877
1878            // Parse the metadata and build the chunk meta
1879            let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
1880            let mut chunk_meta = Vec::with_capacity(words.len());
1881
1882            let mut rows_counter = 0;
1883            let mut offset_bytes = value_buf_position;
1884            for (word_idx, word) in words.iter().enumerate() {
1885                let log_num_values = word & 0x0F;
1886                let divided_bytes = word >> 4;
1887                let num_bytes = (divided_bytes as usize + 1) * MINIBLOCK_ALIGNMENT;
1888                debug_assert!(num_bytes > 0);
1889                let num_values = if word_idx < words.len() - 1 {
1890                    debug_assert!(log_num_values > 0);
1891                    1 << log_num_values
1892                } else {
1893                    debug_assert!(
1894                        log_num_values == 0
1895                            || (1 << log_num_values) == (self.items_in_page - rows_counter)
1896                    );
1897                    self.items_in_page - rows_counter
1898                };
1899                rows_counter += num_values;
1900
1901                chunk_meta.push(ChunkMeta {
1902                    num_values,
1903                    chunk_size_bytes: num_bytes as u64,
1904                    offset_bytes,
1905                });
1906                offset_bytes += num_bytes as u64;
1907            }
1908
1909            // Build the repetition index
1910            let rep_index = if let Some(rep_index_data) = rep_index_bytes {
1911                assert!(rep_index_data.len() % 8 == 0);
1912                let stride = self.repetition_index_depth as usize + 1;
1913                MiniBlockRepIndex::decode_from_bytes(&rep_index_data, stride)
1914            } else {
1915                MiniBlockRepIndex::default_from_chunks(&chunk_meta)
1916            };
1917
1918            let mut page_meta = MiniBlockCacheableState {
1919                chunk_meta,
1920                rep_index,
1921                dictionary: None,
1922            };
1923
1924            // decode dictionary
1925            if let Some(ref mut dictionary) = self.dictionary {
1926                let dictionary_data = dictionary_bytes.unwrap();
1927                page_meta.dictionary =
1928                    Some(Arc::new(dictionary.dictionary_decompressor.decompress(
1929                        LanceBuffer::from_bytes(
1930                            dictionary_data,
1931                            dictionary.dictionary_data_alignment,
1932                        ),
1933                        dictionary.num_dictionary_items,
1934                    )?));
1935            };
1936            let page_meta = Arc::new(page_meta);
1937            self.page_meta = Some(page_meta.clone());
1938            Ok(page_meta as Arc<dyn CachedPageData>)
1939        }
1940        .boxed()
1941    }
1942
1943    fn load(&mut self, data: &Arc<dyn CachedPageData>) {
1944        self.page_meta = Some(
1945            data.clone()
1946                .as_arc_any()
1947                .downcast::<MiniBlockCacheableState>()
1948                .unwrap(),
1949        );
1950    }
1951
1952    fn schedule_ranges(
1953        &self,
1954        ranges: &[Range<u64>],
1955        io: &Arc<dyn EncodingsIo>,
1956    ) -> Result<Vec<PageLoadTask>> {
1957        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
1958
1959        let page_meta = self.page_meta.as_ref().unwrap();
1960
1961        let chunk_instructions =
1962            ChunkInstructions::schedule_instructions(&page_meta.rep_index, ranges);
1963
1964        debug_assert_eq!(
1965            num_rows,
1966            chunk_instructions
1967                .iter()
1968                .map(|ci| ci.rows_to_take)
1969                .sum::<u64>()
1970        );
1971
1972        let chunks_needed = chunk_instructions
1973            .iter()
1974            .map(|ci| ci.chunk_idx)
1975            .unique()
1976            .collect::<Vec<_>>();
1977
1978        let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
1979        let chunk_ranges = loaded_chunks
1980            .iter()
1981            .map(|c| c.byte_range.clone())
1982            .collect::<Vec<_>>();
1983        let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
1984
1985        let rep_decompressor = self.rep_decompressor.clone();
1986        let def_decompressor = self.def_decompressor.clone();
1987        let value_decompressor = self.value_decompressor.clone();
1988        let num_buffers = self.num_buffers;
1989        let has_large_chunk = self.has_large_chunk;
1990        let dictionary = page_meta
1991            .dictionary
1992            .as_ref()
1993            .map(|dictionary| dictionary.clone());
1994        let def_meaning = self.def_meaning.clone();
1995
1996        let res = async move {
1997            let loaded_chunk_data = loaded_chunk_data.await?;
1998            for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
1999                loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
2000            }
2001
2002            Ok(Box::new(MiniBlockDecoder {
2003                rep_decompressor,
2004                def_decompressor,
2005                value_decompressor,
2006                def_meaning,
2007                loaded_chunks: VecDeque::from_iter(loaded_chunks),
2008                instructions: VecDeque::from(chunk_instructions),
2009                offset_in_current_chunk: 0,
2010                dictionary,
2011                num_rows,
2012                num_buffers,
2013                has_large_chunk,
2014            }) as Box<dyn StructuralPageDecoder>)
2015        }
2016        .boxed();
2017        let page_load_task = PageLoadTask {
2018            decoder_fut: res,
2019            num_rows,
2020        };
2021        Ok(vec![page_load_task])
2022    }
2023}
2024
2025#[derive(Debug, Clone, Copy)]
2026struct FullZipRepIndexDetails {
2027    buf_position: u64,
2028    bytes_per_value: u64, // Will be 1, 2, 4, or 8
2029}
2030
2031#[derive(Debug)]
2032enum PerValueDecompressor {
2033    Fixed(Arc<dyn FixedPerValueDecompressor>),
2034    Variable(Arc<dyn VariablePerValueDecompressor>),
2035}
2036
2037#[derive(Debug)]
2038struct FullZipDecodeDetails {
2039    value_decompressor: PerValueDecompressor,
2040    def_meaning: Arc<[DefinitionInterpretation]>,
2041    ctrl_word_parser: ControlWordParser,
2042    max_rep: u16,
2043    max_visible_def: u16,
2044}
2045
2046/// Describes where FullZip byte ranges should be read from.
2047///
2048/// FullZip decoding always needs a list of byte ranges, but those bytes can come
2049/// from two different places:
2050/// - Remote I/O (normal path): ranges are fetched from the underlying `EncodingsIo`.
2051/// - A prefetched full page (full scan fast path): the entire page has already been
2052///   loaded once and ranges should be sliced from memory.
2053///
2054/// This abstraction keeps scheduling code focused on "which ranges are needed"
2055/// instead of "how bytes are fetched", and it lets full-page scans avoid the
2056/// two-stage rep-index -> data I/O pipeline.
2057#[derive(Debug, Clone)]
2058enum FullZipReadSource {
2059    /// Fetch ranges from the storage backend through the encoding I/O interface.
2060    Remote(Arc<dyn EncodingsIo>),
2061    /// Slice ranges from an already-loaded FullZip page buffer.
2062    PrefetchedPage { base_offset: u64, data: LanceBuffer },
2063}
2064
2065impl FullZipReadSource {
2066    /// Materialize the requested ranges as decode-ready `LanceBuffer`s.
2067    ///
2068    /// The returned buffers preserve the input range order.
2069    fn fetch(
2070        &self,
2071        ranges: &[Range<u64>],
2072        priority: u64,
2073    ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
2074        match self {
2075            Self::Remote(io) => {
2076                let io = io.clone();
2077                let ranges = ranges.to_vec();
2078                async move {
2079                    let data = io.submit_request(ranges, priority).await?;
2080                    Ok(data
2081                        .into_iter()
2082                        .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2083                        .collect::<VecDeque<_>>())
2084                }
2085                .boxed()
2086            }
2087            Self::PrefetchedPage { base_offset, data } => {
2088                let base_offset = *base_offset;
2089                let data = data.clone();
2090                let page_end = base_offset + data.len() as u64;
2091                std::future::ready(
2092                    ranges
2093                        .iter()
2094                        .map(|range| {
2095                            if range.start > range.end
2096                                || range.start < base_offset
2097                                || range.end > page_end
2098                            {
2099                                return Err(Error::internal(format!(
2100                                    "Requested range {:?} is outside page range {}..{}",
2101                                    range, base_offset, page_end
2102                                )));
2103                            }
2104                            let start = (range.start - base_offset) as usize;
2105                            let len = (range.end - range.start) as usize;
2106                            Ok(data.slice_with_length(start, len))
2107                        })
2108                        .collect::<Result<VecDeque<_>>>(),
2109                )
2110                .boxed()
2111            }
2112        }
2113    }
2114}
2115
2116/// A scheduler for full-zip encoded data
2117///
2118/// When the data type has a fixed-width then we simply need to map from
2119/// row ranges to byte ranges using the fixed-width of the data type.
2120///
2121/// When the data type is variable-width or has any repetition then a
2122/// repetition index is required.
2123#[derive(Debug)]
2124pub struct FullZipScheduler {
2125    data_buf_position: u64,
2126    data_buf_size: u64,
2127    rep_index: Option<FullZipRepIndexDetails>,
2128    priority: u64,
2129    rows_in_page: u64,
2130    bits_per_offset: u8,
2131    details: Arc<FullZipDecodeDetails>,
2132    /// Cached state containing the decoded repetition index
2133    cached_state: Option<Arc<FullZipCacheableState>>,
2134    /// Whether repetition index metadata should be cached during initialize.
2135    enable_cache: bool,
2136}
2137
2138impl FullZipScheduler {
2139    fn try_new(
2140        buffer_offsets_and_sizes: &[(u64, u64)],
2141        priority: u64,
2142        rows_in_page: u64,
2143        layout: &pb21::FullZipLayout,
2144        decompressors: &dyn DecompressionStrategy,
2145    ) -> Result<Self> {
2146        let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
2147        let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
2148            let num_reps = rows_in_page + 1;
2149            let bytes_per_rep = len / num_reps;
2150            debug_assert_eq!(len % num_reps, 0);
2151            debug_assert!(
2152                bytes_per_rep == 1
2153                    || bytes_per_rep == 2
2154                    || bytes_per_rep == 4
2155                    || bytes_per_rep == 8
2156            );
2157            FullZipRepIndexDetails {
2158                buf_position: *pos,
2159                bytes_per_value: bytes_per_rep,
2160            }
2161        });
2162
2163        let value_decompressor = match layout.details {
2164            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
2165                let decompressor = decompressors.create_fixed_per_value_decompressor(
2166                    layout.value_compression.as_ref().unwrap(),
2167                )?;
2168                PerValueDecompressor::Fixed(decompressor.into())
2169            }
2170            Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
2171                let decompressor = decompressors.create_variable_per_value_decompressor(
2172                    layout.value_compression.as_ref().unwrap(),
2173                )?;
2174                PerValueDecompressor::Variable(decompressor.into())
2175            }
2176            None => {
2177                panic!("Full-zip layout must have a `details` field");
2178            }
2179        };
2180        let ctrl_word_parser = ControlWordParser::new(
2181            layout.bits_rep.try_into().unwrap(),
2182            layout.bits_def.try_into().unwrap(),
2183        );
2184        let def_meaning = layout
2185            .layers
2186            .iter()
2187            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2188            .collect::<Vec<_>>();
2189
2190        let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
2191        let max_visible_def = def_meaning
2192            .iter()
2193            .filter(|d| !d.is_list())
2194            .map(|d| d.num_def_levels())
2195            .sum();
2196
2197        let bits_per_offset = match layout.details {
2198            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
2199            Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
2200                bits_per_offset as u8
2201            }
2202            None => panic!("Full-zip layout must have a `details` field"),
2203        };
2204
2205        let details = Arc::new(FullZipDecodeDetails {
2206            value_decompressor,
2207            def_meaning: def_meaning.into(),
2208            ctrl_word_parser,
2209            max_rep,
2210            max_visible_def,
2211        });
2212        Ok(Self {
2213            data_buf_position,
2214            data_buf_size,
2215            rep_index,
2216            details,
2217            priority,
2218            rows_in_page,
2219            bits_per_offset,
2220            cached_state: None,
2221            enable_cache: false,
2222        })
2223    }
2224
2225    fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
2226        if ranges.is_empty() {
2227            return false;
2228        }
2229        let mut expected_start = 0;
2230        for range in ranges {
2231            if range.start != expected_start || range.end > rows_in_page || range.end < range.start
2232            {
2233                return false;
2234            }
2235            expected_start = range.end;
2236        }
2237        expected_start == rows_in_page
2238    }
2239
2240    fn create_page_load_task(
2241        io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
2242        num_rows: u64,
2243        details: Arc<FullZipDecodeDetails>,
2244        bits_per_offset: u8,
2245    ) -> PageLoadTask {
2246        let load_task = async move {
2247            let buffers = io_future.await?;
2248            let data = buffers
2249                .into_iter()
2250                .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
2251                .collect::<VecDeque<_>>();
2252            Self::create_decoder(details, data, num_rows, bits_per_offset)
2253        }
2254        .boxed();
2255        PageLoadTask {
2256            decoder_fut: load_task,
2257            num_rows,
2258        }
2259    }
2260
2261    /// Creates a decoder from the loaded data
2262    fn create_decoder(
2263        details: Arc<FullZipDecodeDetails>,
2264        data: VecDeque<LanceBuffer>,
2265        num_rows: u64,
2266        bits_per_offset: u8,
2267    ) -> Result<Box<dyn StructuralPageDecoder>> {
2268        match &details.value_decompressor {
2269            PerValueDecompressor::Fixed(decompressor) => {
2270                let bits_per_value = decompressor.bits_per_value();
2271                if bits_per_value % 8 != 0 {
2272                    return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
2273                }
2274                let bytes_per_value = bits_per_value / 8;
2275                let total_bytes_per_value =
2276                    bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
2277                if total_bytes_per_value == 0 {
2278                    return Err(lance_core::Error::internal(
2279                        "Invalid encoding: per-row byte width must be greater than 0",
2280                    ));
2281                }
2282                Ok(Box::new(FixedFullZipDecoder {
2283                    details,
2284                    data,
2285                    num_rows,
2286                    offset_in_current: 0,
2287                    bytes_per_value: bytes_per_value as usize,
2288                    total_bytes_per_value,
2289                }) as Box<dyn StructuralPageDecoder>)
2290            }
2291            PerValueDecompressor::Variable(_decompressor) => {
2292                Ok(Box::new(VariableFullZipDecoder::new(
2293                    details,
2294                    data,
2295                    num_rows,
2296                    bits_per_offset,
2297                    bits_per_offset,
2298                )?))
2299            }
2300        }
2301    }
2302
2303    /// Extracts byte ranges from a repetition index buffer
2304    /// The buffer contains pairs of (start, end) values for each range
2305    fn extract_byte_ranges_from_pairs(
2306        buffer: LanceBuffer,
2307        bytes_per_value: u64,
2308        data_buf_position: u64,
2309    ) -> Vec<Range<u64>> {
2310        ByteUnpacker::new(buffer, bytes_per_value as usize)
2311            .chunks(2)
2312            .into_iter()
2313            .map(|mut c| {
2314                let start = c.next().unwrap() + data_buf_position;
2315                let end = c.next().unwrap() + data_buf_position;
2316                start..end
2317            })
2318            .collect::<Vec<_>>()
2319    }
2320
2321    /// Extracts byte ranges from a cached repetition index buffer
2322    /// The buffer contains all values and we need to extract specific ranges
2323    fn extract_byte_ranges_from_cached(
2324        buffer: &LanceBuffer,
2325        ranges: &[Range<u64>],
2326        bytes_per_value: u64,
2327        data_buf_position: u64,
2328    ) -> Vec<Range<u64>> {
2329        ranges
2330            .iter()
2331            .map(|r| {
2332                let start_offset = (r.start * bytes_per_value) as usize;
2333                let end_offset = (r.end * bytes_per_value) as usize;
2334
2335                let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
2336                let start_val =
2337                    ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
2338                        .next()
2339                        .unwrap();
2340
2341                let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
2342                let end_val =
2343                    ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
2344                        .next()
2345                        .unwrap();
2346
2347                (data_buf_position + start_val)..(data_buf_position + end_val)
2348            })
2349            .collect()
2350    }
2351
2352    /// Computes the ranges in the repetition index that need to be loaded
2353    fn compute_rep_index_ranges(
2354        ranges: &[Range<u64>],
2355        rep_index: &FullZipRepIndexDetails,
2356    ) -> Vec<Range<u64>> {
2357        ranges
2358            .iter()
2359            .flat_map(|r| {
2360                let first_val_start =
2361                    rep_index.buf_position + (r.start * rep_index.bytes_per_value);
2362                let first_val_end = first_val_start + rep_index.bytes_per_value;
2363                let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
2364                let last_val_end = last_val_start + rep_index.bytes_per_value;
2365                [first_val_start..first_val_end, last_val_start..last_val_end]
2366            })
2367            .collect()
2368    }
2369
2370    /// Schedules ranges in the presence of a repetition index
2371    fn schedule_ranges_rep(
2372        &self,
2373        ranges: &[Range<u64>],
2374        io: &Arc<dyn EncodingsIo>,
2375        rep_index: FullZipRepIndexDetails,
2376    ) -> Result<Vec<PageLoadTask>> {
2377        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2378        let data_buf_position = self.data_buf_position;
2379        let priority = self.priority;
2380        let details = self.details.clone();
2381        let bits_per_offset = self.bits_per_offset;
2382
2383        if Self::covers_entire_page(ranges, self.rows_in_page) {
2384            let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
2385            let page_data = io.submit_single(full_range.clone(), priority);
2386            let load_task = async move {
2387                let page_data = page_data.await?;
2388                let source = FullZipReadSource::PrefetchedPage {
2389                    base_offset: full_range.start,
2390                    data: LanceBuffer::from_bytes(page_data, 1),
2391                };
2392                let read_ranges = vec![full_range];
2393                let data = source.fetch(&read_ranges, priority).await?;
2394                Self::create_decoder(details, data, num_rows, bits_per_offset)
2395            }
2396            .boxed();
2397            let page_load_task = PageLoadTask {
2398                decoder_fut: load_task,
2399                num_rows,
2400            };
2401            return Ok(vec![page_load_task]);
2402        }
2403
2404        if let Some(cached_state) = &self.cached_state {
2405            let byte_ranges = Self::extract_byte_ranges_from_cached(
2406                &cached_state.rep_index_buffer,
2407                ranges,
2408                rep_index.bytes_per_value,
2409                data_buf_position,
2410            );
2411            let io_future = io.submit_request(byte_ranges, priority);
2412            let page_load_task =
2413                Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
2414            return Ok(vec![page_load_task]);
2415        }
2416
2417        let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
2418        let rep_data = io.submit_request(rep_ranges, priority);
2419        let io_clone = io.clone();
2420        let load_task = async move {
2421            let rep_data = rep_data.await?;
2422            let rep_buffer = LanceBuffer::concat(
2423                &rep_data
2424                    .into_iter()
2425                    .map(|d| LanceBuffer::from_bytes(d, 1))
2426                    .collect::<Vec<_>>(),
2427            );
2428            let byte_ranges = Self::extract_byte_ranges_from_pairs(
2429                rep_buffer,
2430                rep_index.bytes_per_value,
2431                data_buf_position,
2432            );
2433            let source = FullZipReadSource::Remote(io_clone);
2434            let data = source.fetch(&byte_ranges, priority).await?;
2435            Self::create_decoder(details, data, num_rows, bits_per_offset)
2436        }
2437        .boxed();
2438        let page_load_task = PageLoadTask {
2439            decoder_fut: load_task,
2440            num_rows,
2441        };
2442        Ok(vec![page_load_task])
2443    }
2444
2445    // In the simple case there is no repetition and we just have large fixed-width
2446    // rows of data.  We can just map row ranges to byte ranges directly using the
2447    // fixed-width of the data type.
2448    fn schedule_ranges_simple(
2449        &self,
2450        ranges: &[Range<u64>],
2451        io: &Arc<dyn EncodingsIo>,
2452    ) -> Result<Vec<PageLoadTask>> {
2453        // Convert row ranges to item ranges (i.e. multiply by items per row)
2454        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2455
2456        let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
2457            unreachable!()
2458        };
2459
2460        // Convert item ranges to byte ranges (i.e. multiply by bytes per item)
2461        let bits_per_value = decompressor.bits_per_value();
2462        assert_eq!(bits_per_value % 8, 0);
2463        let bytes_per_value = bits_per_value / 8;
2464        let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
2465        let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
2466        let byte_ranges = ranges
2467            .iter()
2468            .map(|r| {
2469                debug_assert!(r.end <= self.rows_in_page);
2470                let start = self.data_buf_position + r.start * total_bytes_per_value;
2471                let end = self.data_buf_position + r.end * total_bytes_per_value;
2472                start..end
2473            })
2474            .collect::<Vec<_>>();
2475
2476        let io_future = io.submit_request(byte_ranges, self.priority);
2477        let page_load_task = Self::create_page_load_task(
2478            io_future,
2479            num_rows,
2480            self.details.clone(),
2481            self.bits_per_offset,
2482        );
2483        Ok(vec![page_load_task])
2484    }
2485}
2486
2487/// Cacheable state for FullZip encoding, storing the decoded repetition index
2488#[derive(Debug)]
2489struct FullZipCacheableState {
2490    /// The raw repetition index buffer for future decoding
2491    rep_index_buffer: LanceBuffer,
2492}
2493
2494impl DeepSizeOf for FullZipCacheableState {
2495    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
2496        self.rep_index_buffer.len()
2497    }
2498}
2499
2500impl CachedPageData for FullZipCacheableState {
2501    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
2502        self
2503    }
2504}
2505
2506impl StructuralPageScheduler for FullZipScheduler {
2507    fn initialize<'a>(
2508        &'a mut self,
2509        io: &Arc<dyn EncodingsIo>,
2510    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2511        if self.enable_cache
2512            && let Some(rep_index) = self.rep_index
2513        {
2514            let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
2515            let rep_index_range = rep_index.buf_position..(rep_index.buf_position + total_size);
2516            let io_clone = io.clone();
2517            return async move {
2518                let rep_index_data = io_clone.submit_request(vec![rep_index_range], 0).await?;
2519                let state = Arc::new(FullZipCacheableState {
2520                    rep_index_buffer: LanceBuffer::from_bytes(rep_index_data[0].clone(), 1),
2521                });
2522                self.cached_state = Some(state.clone());
2523                Ok(state as Arc<dyn CachedPageData>)
2524            }
2525            .boxed();
2526        }
2527        std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
2528    }
2529
2530    /// Loads previously cached repetition index data from the cache system.
2531    /// This method is called when a scheduler instance needs to use cached data
2532    /// that was initialized by another instance or in a previous operation.
2533    fn load(&mut self, cache: &Arc<dyn CachedPageData>) {
2534        // Try to downcast to our specific cache type
2535        if let Ok(cached_state) = cache
2536            .clone()
2537            .as_arc_any()
2538            .downcast::<FullZipCacheableState>()
2539        {
2540            // Store the cached state for use in schedule_ranges
2541            self.cached_state = Some(cached_state);
2542        }
2543    }
2544
2545    fn schedule_ranges(
2546        &self,
2547        ranges: &[Range<u64>],
2548        io: &Arc<dyn EncodingsIo>,
2549    ) -> Result<Vec<PageLoadTask>> {
2550        if let Some(rep_index) = self.rep_index {
2551            self.schedule_ranges_rep(ranges, io, rep_index)
2552        } else {
2553            self.schedule_ranges_simple(ranges, io)
2554        }
2555    }
2556}
2557
2558/// A decoder for full-zip encoded data when the data has a fixed-width
2559///
2560/// Here we need to unzip the control words from the values themselves and
2561/// then decompress the requested values.
2562///
2563/// We use a PerValueDecompressor because we will only be decompressing the
2564/// requested data.  This decoder / scheduler does not do any read amplification.
2565#[derive(Debug)]
2566struct FixedFullZipDecoder {
2567    details: Arc<FullZipDecodeDetails>,
2568    data: VecDeque<LanceBuffer>,
2569    offset_in_current: usize,
2570    bytes_per_value: usize,
2571    total_bytes_per_value: usize,
2572    num_rows: u64,
2573}
2574
2575impl FixedFullZipDecoder {
2576    fn slice_next_task(&mut self, num_rows: u64) -> FullZipDecodeTaskItem {
2577        debug_assert!(num_rows > 0);
2578        let cur_buf = self.data.front_mut().unwrap();
2579        let start = self.offset_in_current;
2580        if self.details.ctrl_word_parser.has_rep() {
2581            // This is a slightly slower path.  In order to figure out where to split we need to
2582            // examine the rep index so we can convert num_lists to num_rows
2583            let mut rows_started = 0;
2584            // We always need at least one value.  Now loop through until we have passed num_rows
2585            // values
2586            let mut num_items = 0;
2587            while self.offset_in_current < cur_buf.len() {
2588                let control = self.details.ctrl_word_parser.parse_desc(
2589                    &cur_buf[self.offset_in_current..],
2590                    self.details.max_rep,
2591                    self.details.max_visible_def,
2592                );
2593                if control.is_new_row {
2594                    if rows_started == num_rows {
2595                        break;
2596                    }
2597                    rows_started += 1;
2598                }
2599                num_items += 1;
2600                if control.is_visible {
2601                    self.offset_in_current += self.total_bytes_per_value;
2602                } else {
2603                    self.offset_in_current += self.details.ctrl_word_parser.bytes_per_word();
2604                }
2605            }
2606
2607            let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
2608            if self.offset_in_current == cur_buf.len() {
2609                self.data.pop_front();
2610                self.offset_in_current = 0;
2611            }
2612
2613            FullZipDecodeTaskItem {
2614                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2615                    data: task_slice,
2616                    bits_per_value: self.bytes_per_value as u64 * 8,
2617                    num_values: num_items,
2618                    block_info: BlockInfo::new(),
2619                }),
2620                rows_in_buf: rows_started,
2621            }
2622        } else {
2623            // If there's no repetition we can calculate the slicing point by just multiplying
2624            // the number of rows by the total bytes per value
2625            let cur_buf = self.data.front_mut().unwrap();
2626            let bytes_avail = cur_buf.len() - self.offset_in_current;
2627            let offset_in_cur = self.offset_in_current;
2628
2629            let bytes_needed = num_rows as usize * self.total_bytes_per_value;
2630            let mut rows_taken = num_rows;
2631            let task_slice = if bytes_needed >= bytes_avail {
2632                self.offset_in_current = 0;
2633                rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
2634                self.data
2635                    .pop_front()
2636                    .unwrap()
2637                    .slice_with_length(offset_in_cur, bytes_avail)
2638            } else {
2639                self.offset_in_current += bytes_needed;
2640                cur_buf.slice_with_length(offset_in_cur, bytes_needed)
2641            };
2642            FullZipDecodeTaskItem {
2643                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
2644                    data: task_slice,
2645                    bits_per_value: self.bytes_per_value as u64 * 8,
2646                    num_values: rows_taken,
2647                    block_info: BlockInfo::new(),
2648                }),
2649                rows_in_buf: rows_taken,
2650            }
2651        }
2652    }
2653}
2654
2655impl StructuralPageDecoder for FixedFullZipDecoder {
2656    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2657        let mut task_data = Vec::with_capacity(self.data.len());
2658        let mut remaining = num_rows;
2659        while remaining > 0 {
2660            let task_item = self.slice_next_task(remaining);
2661            remaining -= task_item.rows_in_buf;
2662            task_data.push(task_item);
2663        }
2664        Ok(Box::new(FixedFullZipDecodeTask {
2665            details: self.details.clone(),
2666            data: task_data,
2667            bytes_per_value: self.bytes_per_value,
2668            num_rows: num_rows as usize,
2669        }))
2670    }
2671
2672    fn num_rows(&self) -> u64 {
2673        self.num_rows
2674    }
2675}
2676
2677/// A decoder for full-zip encoded data when the data has a variable-width
2678///
2679/// Here we need to unzip the control words AND lengths from the values and
2680/// then decompress the requested values.
2681#[derive(Debug)]
2682struct VariableFullZipDecoder {
2683    details: Arc<FullZipDecodeDetails>,
2684    decompressor: Arc<dyn VariablePerValueDecompressor>,
2685    data: LanceBuffer,
2686    offsets: LanceBuffer,
2687    rep: ScalarBuffer<u16>,
2688    def: ScalarBuffer<u16>,
2689    repdef_starts: Vec<usize>,
2690    data_starts: Vec<usize>,
2691    offset_starts: Vec<usize>,
2692    visible_item_counts: Vec<u64>,
2693    bits_per_offset: u8,
2694    current_idx: usize,
2695    num_rows: u64,
2696}
2697
2698fn corrupt_file_named(name: &str, message: impl Into<String>) -> Error {
2699    Error::corrupt_file(name.into(), message)
2700}
2701
2702impl VariableFullZipDecoder {
2703    fn new(
2704        details: Arc<FullZipDecodeDetails>,
2705        data: VecDeque<LanceBuffer>,
2706        num_rows: u64,
2707        in_bits_per_length: u8,
2708        out_bits_per_offset: u8,
2709    ) -> Result<Self> {
2710        let decompressor = match details.value_decompressor {
2711            PerValueDecompressor::Variable(ref d) => d.clone(),
2712            _ => unreachable!(),
2713        };
2714
2715        assert_eq!(in_bits_per_length % 8, 0);
2716        assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
2717
2718        let mut decoder = Self {
2719            details,
2720            decompressor,
2721            data: LanceBuffer::empty(),
2722            offsets: LanceBuffer::empty(),
2723            rep: LanceBuffer::empty().borrow_to_typed_slice(),
2724            def: LanceBuffer::empty().borrow_to_typed_slice(),
2725            bits_per_offset: out_bits_per_offset,
2726            repdef_starts: Vec::with_capacity(num_rows as usize + 1),
2727            data_starts: Vec::with_capacity(num_rows as usize + 1),
2728            offset_starts: Vec::with_capacity(num_rows as usize + 1),
2729            visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
2730            current_idx: 0,
2731            num_rows,
2732        };
2733
2734        // There's no great time to do this and this is the least worst time.  If we don't unzip then
2735        // we can't slice the data during the decode phase.  This is because we need the offsets to be
2736        // unpacked to know where the values start and end.
2737        //
2738        // We don't want to unzip on the decode thread because that is a single-threaded path
2739        // We don't want to unzip on the scheduling thread because that is a single-threaded path
2740        //
2741        // Fortunately, we know variable length data will always be read indirectly and so we can do it
2742        // here, which should be on the indirect thread.  The primary disadvantage to doing it here is that
2743        // we load all the data into memory and then throw it away only to load it all into memory again during
2744        // the decode.
2745        //
2746        // There are some alternatives to investigate:
2747        //   - Instead of just reading the beginning and end of the rep index we could read the entire
2748        //     range in between.  This will give us the break points that we need for slicing and won't increase
2749        //     the number of IOPs but it will mean we are doing more total I/O and we need to load the rep index
2750        //     even when doing a full scan.
2751        //   - We could force each decode task to do a full unzip of all the data.  Each decode task now
2752        //     has to do more work but the work is all fused.
2753        //   - We could just try doing this work on the decode thread and see if it is a problem.
2754        decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
2755
2756        Ok(decoder)
2757    }
2758
2759    fn slice_batch_data_and_rebase_offsets_typed<T>(
2760        data: &LanceBuffer,
2761        offsets: &LanceBuffer,
2762    ) -> Result<(LanceBuffer, LanceBuffer)>
2763    where
2764        T: arrow_buffer::ArrowNativeType
2765            + Copy
2766            + PartialOrd
2767            + std::ops::Sub<Output = T>
2768            + std::fmt::Display
2769            + TryInto<usize>,
2770    {
2771        let offsets_slice = offsets.borrow_to_typed_slice::<T>();
2772        let offsets_slice = offsets_slice.as_ref();
2773        if offsets_slice.is_empty() {
2774            return Err(Error::internal(
2775                "Variable offsets cannot be empty".to_string(),
2776            ));
2777        }
2778
2779        let base = offsets_slice[0];
2780        let end = *offsets_slice.last().unwrap();
2781        if end < base {
2782            return Err(Error::internal(format!(
2783                "Invalid variable offsets: end ({end}) is less than base ({base})"
2784            )));
2785        }
2786
2787        let data_start = base.try_into().map_err(|_| {
2788            Error::internal(format!("Variable offset ({base}) does not fit into usize"))
2789        })?;
2790        let data_end = end.try_into().map_err(|_| {
2791            Error::internal(format!("Variable offset ({end}) does not fit into usize"))
2792        })?;
2793        if data_end > data.len() {
2794            return Err(Error::internal(format!(
2795                "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
2796                data.len()
2797            )));
2798        }
2799
2800        let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
2801        for &offset in offsets_slice {
2802            if offset < base {
2803                return Err(Error::internal(format!(
2804                    "Invalid variable offsets: offset ({offset}) is less than base ({base})"
2805                )));
2806            }
2807            rebased_offsets.push(offset - base);
2808        }
2809
2810        let sliced_data = data.slice_with_length(data_start, data_end - data_start);
2811        // Copy into a compact buffer so each output batch owns only what it references.
2812        let sliced_data = LanceBuffer::copy_slice(&sliced_data);
2813        let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
2814        Ok((sliced_data, rebased_offsets))
2815    }
2816
2817    fn slice_batch_data_and_rebase_offsets(
2818        data: &LanceBuffer,
2819        offsets: &LanceBuffer,
2820        bits_per_offset: u8,
2821    ) -> Result<(LanceBuffer, LanceBuffer)> {
2822        match bits_per_offset {
2823            32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
2824            64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
2825            _ => Err(Error::internal(format!(
2826                "Unsupported bits_per_offset={bits_per_offset}"
2827            ))),
2828        }
2829    }
2830
2831    /// Reads a single length prefix from the front of `data`.
2832    ///
2833    /// The bytes come from the file. A page whose item walk ends with a partial
2834    /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this
2835    /// is bounds checked and reports a corrupt file rather than reading past the
2836    /// end of the buffer.
2837    fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
2838        let width = bits_per_offset as usize / 8;
2839        if data.len() < width {
2840            return Err(corrupt_file_named(
2841                "variable_full_zip",
2842                format!(
2843                    "truncated length prefix: {} byte(s) remain in the page buffer but a \
2844                     {}-bit length prefix requires {}",
2845                    data.len(),
2846                    bits_per_offset,
2847                    width
2848                ),
2849            ));
2850        }
2851        Ok(match bits_per_offset {
2852            8 => data[0] as u64,
2853            16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
2854            32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
2855            64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
2856            _ => unreachable!(),
2857        })
2858    }
2859
2860    fn unzip(
2861        &mut self,
2862        data: VecDeque<LanceBuffer>,
2863        in_bits_per_length: u8,
2864        out_bits_per_offset: u8,
2865        num_rows: u64,
2866    ) -> Result<()> {
2867        // This undercounts if there are lists but, at this point, we don't really know how many items we have
2868        let mut rep = Vec::with_capacity(num_rows as usize);
2869        let mut def = Vec::with_capacity(num_rows as usize);
2870        let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
2871
2872        // This undercounts if there are lists
2873        // It can also overcount if there are invisible items
2874        let bytes_per_offset = out_bits_per_offset as usize / 8;
2875        let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
2876        let mut offsets_data = Vec::with_capacity(bytes_offsets);
2877
2878        let bytes_per_length = in_bits_per_length as usize / 8;
2879        let bytes_lengths = bytes_per_length * num_rows as usize;
2880
2881        let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
2882        // This overcounts since bytes_lengths and bytes_cw are undercounts
2883        // It can also undercount if there are invisible items (hence the saturating_sub)
2884        let mut unzipped_data =
2885            Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
2886
2887        let mut current_offset = 0_u64;
2888        let mut visible_item_count = 0_u64;
2889        for databuf in data.into_iter() {
2890            let mut databuf = databuf.as_ref();
2891            while !databuf.is_empty() {
2892                let data_start = unzipped_data.len();
2893                let offset_start = offsets_data.len();
2894                // We might have only-rep or only-def, neither, or both.  They move at the same
2895                // speed though so we only need one index into it
2896                let repdef_start = rep.len().max(def.len());
2897                // TODO: Kind of inefficient we parse the control word twice here
2898                let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
2899                    databuf,
2900                    self.details.max_rep,
2901                    self.details.max_visible_def,
2902                );
2903                self.details
2904                    .ctrl_word_parser
2905                    .parse(databuf, &mut rep, &mut def);
2906                databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
2907
2908                if ctrl_desc.is_new_row {
2909                    self.repdef_starts.push(repdef_start);
2910                    self.data_starts.push(data_start);
2911                    self.offset_starts.push(offset_start);
2912                    self.visible_item_counts.push(visible_item_count);
2913                }
2914                if ctrl_desc.is_visible {
2915                    visible_item_count += 1;
2916                    if ctrl_desc.is_valid_item {
2917                        let length = Self::parse_length(databuf, in_bits_per_length)?;
2918                        match out_bits_per_offset {
2919                            32 => offsets_data
2920                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2921                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2922                            _ => unreachable!(),
2923                        };
2924                        databuf = &databuf[bytes_per_offset..];
2925                        unzipped_data.extend_from_slice(&databuf[..length as usize]);
2926                        databuf = &databuf[length as usize..];
2927                        current_offset += length;
2928                    } else {
2929                        // Null items still get an offset
2930                        match out_bits_per_offset {
2931                            32 => offsets_data
2932                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
2933                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2934                            _ => unreachable!(),
2935                        }
2936                    }
2937                }
2938            }
2939        }
2940        self.repdef_starts.push(rep.len().max(def.len()));
2941        self.data_starts.push(unzipped_data.len());
2942        self.offset_starts.push(offsets_data.len());
2943        self.visible_item_counts.push(visible_item_count);
2944        match out_bits_per_offset {
2945            32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
2946            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
2947            _ => unreachable!(),
2948        };
2949        self.rep = ScalarBuffer::from(rep);
2950        self.def = ScalarBuffer::from(def);
2951        self.data = LanceBuffer::from(unzipped_data);
2952        self.offsets = LanceBuffer::from(offsets_data);
2953        Ok(())
2954    }
2955}
2956
2957impl StructuralPageDecoder for VariableFullZipDecoder {
2958    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2959        let start = self.current_idx;
2960        let end = start + num_rows as usize;
2961
2962        let offset_start = self.offset_starts[start];
2963        let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
2964        let offsets = self
2965            .offsets
2966            .slice_with_length(offset_start, offset_end - offset_start);
2967        // Keep each batch's variable data buffer bounded to the selected rows.
2968        let (data, offsets) =
2969            Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
2970
2971        let repdef_start = self.repdef_starts[start];
2972        let repdef_end = self.repdef_starts[end];
2973        let rep = if self.rep.is_empty() {
2974            self.rep.clone()
2975        } else {
2976            self.rep.slice(repdef_start, repdef_end - repdef_start)
2977        };
2978        let def = if self.def.is_empty() {
2979            self.def.clone()
2980        } else {
2981            self.def.slice(repdef_start, repdef_end - repdef_start)
2982        };
2983
2984        let visible_item_counts_start = self.visible_item_counts[start];
2985        let visible_item_counts_end = self.visible_item_counts[end];
2986        let num_visible_items = visible_item_counts_end - visible_item_counts_start;
2987
2988        self.current_idx += num_rows as usize;
2989
2990        Ok(Box::new(VariableFullZipDecodeTask {
2991            details: self.details.clone(),
2992            decompressor: self.decompressor.clone(),
2993            data,
2994            offsets,
2995            bits_per_offset: self.bits_per_offset,
2996            num_visible_items,
2997            rep,
2998            def,
2999        }))
3000    }
3001
3002    fn num_rows(&self) -> u64 {
3003        self.num_rows
3004    }
3005}
3006
3007#[derive(Debug)]
3008struct VariableFullZipDecodeTask {
3009    details: Arc<FullZipDecodeDetails>,
3010    decompressor: Arc<dyn VariablePerValueDecompressor>,
3011    data: LanceBuffer,
3012    offsets: LanceBuffer,
3013    bits_per_offset: u8,
3014    num_visible_items: u64,
3015    rep: ScalarBuffer<u16>,
3016    def: ScalarBuffer<u16>,
3017}
3018
3019impl DecodePageTask for VariableFullZipDecodeTask {
3020    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3021        let block = VariableWidthBlock {
3022            data: self.data,
3023            offsets: self.offsets,
3024            bits_per_offset: self.bits_per_offset,
3025            num_values: self.num_visible_items,
3026            block_info: BlockInfo::new(),
3027        };
3028        let decomopressed = self.decompressor.decompress(block)?;
3029        let rep = if self.rep.is_empty() {
3030            None
3031        } else {
3032            Some(self.rep.to_vec())
3033        };
3034        let def = if self.def.is_empty() {
3035            None
3036        } else {
3037            Some(self.def.to_vec())
3038        };
3039        let unraveler = RepDefUnraveler::new(
3040            rep,
3041            def,
3042            self.details.def_meaning.clone(),
3043            self.num_visible_items,
3044        );
3045        Ok(DecodedPage {
3046            data: decomopressed,
3047            repdef: unraveler,
3048        })
3049    }
3050}
3051
3052#[derive(Debug)]
3053struct FullZipDecodeTaskItem {
3054    data: PerValueDataBlock,
3055    rows_in_buf: u64,
3056}
3057
3058/// A task to unzip and decompress full-zip encoded data when that data
3059/// has a fixed-width.
3060#[derive(Debug)]
3061struct FixedFullZipDecodeTask {
3062    details: Arc<FullZipDecodeDetails>,
3063    data: Vec<FullZipDecodeTaskItem>,
3064    num_rows: usize,
3065    bytes_per_value: usize,
3066}
3067
3068impl DecodePageTask for FixedFullZipDecodeTask {
3069    fn decode(self: Box<Self>) -> Result<DecodedPage> {
3070        // Multiply by 2 to make a stab at the size of the output buffer (which will be decompressed and thus bigger)
3071        let estimated_size_bytes = self
3072            .data
3073            .iter()
3074            .map(|task_item| task_item.data.data_size() as usize)
3075            .sum::<usize>()
3076            * 2;
3077        let mut data_builder =
3078            DataBlockBuilder::with_capacity_estimate(estimated_size_bytes as u64);
3079
3080        if self.details.ctrl_word_parser.bytes_per_word() == 0 {
3081            // Fast path, no need to unzip because there is no rep/def
3082            //
3083            // We decompress each buffer and add it to our output buffer
3084            for task_item in self.data.into_iter() {
3085                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3086                    unreachable!()
3087                };
3088                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3089                else {
3090                    unreachable!()
3091                };
3092                debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
3093                let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
3094                data_builder.append(&decompressed, 0..task_item.rows_in_buf);
3095            }
3096
3097            let unraveler = RepDefUnraveler::new(
3098                None,
3099                None,
3100                self.details.def_meaning.clone(),
3101                self.num_rows as u64,
3102            );
3103
3104            Ok(DecodedPage {
3105                data: data_builder.finish(),
3106                repdef: unraveler,
3107            })
3108        } else {
3109            // Slow path, unzipping needed
3110            let mut rep = Vec::with_capacity(self.num_rows);
3111            let mut def = Vec::with_capacity(self.num_rows);
3112
3113            for task_item in self.data.into_iter() {
3114                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
3115                    unreachable!()
3116                };
3117                let mut buf_slice = fixed_data.data.as_ref();
3118                let num_values = fixed_data.num_values as usize;
3119                // We will be unzipping repdef in to `rep` and `def` and the
3120                // values into `values` (which contains the compressed values)
3121                let mut values = Vec::with_capacity(
3122                    fixed_data.data.len()
3123                        - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
3124                );
3125                let mut visible_items = 0;
3126                for _ in 0..num_values {
3127                    // Extract rep/def
3128                    self.details
3129                        .ctrl_word_parser
3130                        .parse(buf_slice, &mut rep, &mut def);
3131                    buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
3132
3133                    let is_visible = def
3134                        .last()
3135                        .map(|d| *d <= self.details.max_visible_def)
3136                        .unwrap_or(true);
3137                    if is_visible {
3138                        // Extract value
3139                        values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
3140                        buf_slice = &buf_slice[self.bytes_per_value..];
3141                        visible_items += 1;
3142                    }
3143                }
3144
3145                // Finally, we decompress the values and add them to our output buffer
3146                let values_buf = LanceBuffer::from(values);
3147                let fixed_data = FixedWidthDataBlock {
3148                    bits_per_value: self.bytes_per_value as u64 * 8,
3149                    block_info: BlockInfo::new(),
3150                    data: values_buf,
3151                    num_values: visible_items,
3152                };
3153                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
3154                else {
3155                    unreachable!()
3156                };
3157                let decompressed = decompressor.decompress(fixed_data, visible_items)?;
3158                data_builder.append(&decompressed, 0..visible_items);
3159            }
3160
3161            let repetition = if rep.is_empty() { None } else { Some(rep) };
3162            let definition = if def.is_empty() { None } else { Some(def) };
3163
3164            let unraveler = RepDefUnraveler::new(
3165                repetition,
3166                definition,
3167                self.details.def_meaning.clone(),
3168                self.num_rows as u64,
3169            );
3170            let data = data_builder.finish();
3171
3172            Ok(DecodedPage {
3173                data,
3174                repdef: unraveler,
3175            })
3176        }
3177    }
3178}
3179
3180#[derive(Debug)]
3181struct StructuralPrimitiveFieldSchedulingJob<'a> {
3182    scheduler: &'a StructuralPrimitiveFieldScheduler,
3183    ranges: Vec<Range<u64>>,
3184    page_idx: usize,
3185    range_idx: usize,
3186    global_row_offset: u64,
3187}
3188
3189impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
3190    pub fn new(scheduler: &'a StructuralPrimitiveFieldScheduler, ranges: Vec<Range<u64>>) -> Self {
3191        Self {
3192            scheduler,
3193            ranges,
3194            page_idx: 0,
3195            range_idx: 0,
3196            global_row_offset: 0,
3197        }
3198    }
3199}
3200
3201impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
3202    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
3203        if self.range_idx >= self.ranges.len() {
3204            return Ok(Vec::new());
3205        }
3206        // Get our current range
3207        let mut range = self.ranges[self.range_idx].clone();
3208        let priority = range.start;
3209
3210        let mut cur_page = &self.scheduler.page_schedulers[self.page_idx];
3211        trace!(
3212            "Current range is {:?} and current page has {} rows",
3213            range, cur_page.num_rows
3214        );
3215        // Skip entire pages until we have some overlap with our next range
3216        while cur_page.num_rows + self.global_row_offset <= range.start {
3217            self.global_row_offset += cur_page.num_rows;
3218            self.page_idx += 1;
3219            trace!("Skipping entire page of {} rows", cur_page.num_rows);
3220            cur_page = &self.scheduler.page_schedulers[self.page_idx];
3221        }
3222
3223        // Now the cur_page has overlap with range.  Continue looping through ranges
3224        // until we find a range that exceeds the current page
3225
3226        let mut ranges_in_page = Vec::new();
3227        while cur_page.num_rows + self.global_row_offset > range.start {
3228            range.start = range.start.max(self.global_row_offset);
3229            let start_in_page = range.start - self.global_row_offset;
3230            let end_in_page = start_in_page + (range.end - range.start);
3231            let end_in_page = end_in_page.min(cur_page.num_rows);
3232            let last_in_range = (end_in_page + self.global_row_offset) >= range.end;
3233
3234            ranges_in_page.push(start_in_page..end_in_page);
3235            if last_in_range {
3236                self.range_idx += 1;
3237                if self.range_idx == self.ranges.len() {
3238                    break;
3239                }
3240                range = self.ranges[self.range_idx].clone();
3241            } else {
3242                break;
3243            }
3244        }
3245
3246        trace!(
3247            "Scheduling {} rows across {} ranges from page with {} rows (priority={}, column_index={}, page_index={})",
3248            ranges_in_page.iter().map(|r| r.end - r.start).sum::<u64>(),
3249            ranges_in_page.len(),
3250            cur_page.num_rows,
3251            priority,
3252            self.scheduler.column_index,
3253            cur_page.page_index,
3254        );
3255
3256        self.global_row_offset += cur_page.num_rows;
3257        self.page_idx += 1;
3258
3259        let page_decoders = cur_page
3260            .scheduler
3261            .schedule_ranges(&ranges_in_page, context.io())?;
3262
3263        let cur_path = context.current_path();
3264        page_decoders
3265            .into_iter()
3266            .map(|page_load_task| {
3267                let cur_path = cur_path.clone();
3268                let page_decoder = page_load_task.decoder_fut;
3269                let unloaded_page = async move {
3270                    let page_decoder = page_decoder.await?;
3271                    Ok(LoadedPageShard {
3272                        decoder: page_decoder,
3273                        path: cur_path,
3274                    })
3275                }
3276                .boxed();
3277                Ok(ScheduledScanLine {
3278                    decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
3279                    rows_scheduled: page_load_task.num_rows,
3280                })
3281            })
3282            .collect::<Result<Vec<_>>>()
3283    }
3284}
3285
3286#[derive(Debug)]
3287struct PageInfoAndScheduler {
3288    page_index: usize,
3289    num_rows: u64,
3290    scheduler: Box<dyn StructuralPageScheduler>,
3291}
3292
3293/// A scheduler for a leaf node
3294///
3295/// Here we look at the layout of the various pages and delegate scheduling to a scheduler
3296/// appropriate for the layout of the page.
3297#[derive(Debug)]
3298pub struct StructuralPrimitiveFieldScheduler {
3299    page_schedulers: Vec<PageInfoAndScheduler>,
3300    column_index: u32,
3301    // Identifies the requested decode shape (e.g. blob descriptor struct vs
3302    // raw bytes). Blob columns can produce multiple page scheduler variants
3303    // for the same physical column depending on the target field's data type,
3304    // and the cached page state types differ per variant. The view tag is
3305    // mixed into the cache key so different variants do not collide.
3306    view_tag: String,
3307}
3308
3309impl StructuralPrimitiveFieldScheduler {
3310    pub fn try_new(
3311        column_info: &ColumnInfo,
3312        decompressors: &dyn DecompressionStrategy,
3313        cache_repetition_index: bool,
3314        target_field: &Field,
3315    ) -> Result<Self> {
3316        let page_schedulers = column_info
3317            .page_infos
3318            .iter()
3319            .enumerate()
3320            .map(|(page_index, page_info)| {
3321                Self::page_info_to_scheduler(
3322                    page_info,
3323                    page_index,
3324                    decompressors,
3325                    cache_repetition_index,
3326                    target_field,
3327                )
3328            })
3329            .collect::<Result<Vec<_>>>()?;
3330        Ok(Self {
3331            page_schedulers,
3332            column_index: column_info.index,
3333            view_tag: format!("{:?}", target_field.data_type()),
3334        })
3335    }
3336
3337    fn page_layout_to_scheduler(
3338        page_info: &PageInfo,
3339        page_layout: &PageLayout,
3340        decompressors: &dyn DecompressionStrategy,
3341        cache_repetition_index: bool,
3342        target_field: &Field,
3343    ) -> Result<Box<dyn StructuralPageScheduler>> {
3344        use pb21::page_layout::Layout;
3345        Ok(match page_layout.layout.as_ref().expect_ok()? {
3346            Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
3347                &page_info.buffer_offsets_and_sizes,
3348                page_info.priority,
3349                mini_block.num_items,
3350                mini_block,
3351                decompressors,
3352            )?),
3353            Layout::FullZipLayout(full_zip) => {
3354                let mut scheduler = FullZipScheduler::try_new(
3355                    &page_info.buffer_offsets_and_sizes,
3356                    page_info.priority,
3357                    page_info.num_rows,
3358                    full_zip,
3359                    decompressors,
3360                )?;
3361                scheduler.enable_cache = cache_repetition_index;
3362                Box::new(scheduler)
3363            }
3364            Layout::ConstantLayout(constant_layout) => {
3365                let def_meaning = constant_layout
3366                    .layers
3367                    .iter()
3368                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3369                    .collect::<Vec<_>>();
3370                let has_scalar_value = constant_layout.inline_value.is_some()
3371                    || page_info.buffer_offsets_and_sizes.len() == 1
3372                    || page_info.buffer_offsets_and_sizes.len() == 3;
3373                if has_scalar_value {
3374                    Box::new(constant::ConstantPageScheduler::try_new(
3375                        page_info.buffer_offsets_and_sizes.clone(),
3376                        constant_layout.inline_value.clone(),
3377                        target_field.data_type(),
3378                        def_meaning.into(),
3379                    )?) as Box<dyn StructuralPageScheduler>
3380                } else if def_meaning.len() == 1
3381                    && def_meaning[0] == DefinitionInterpretation::NullableItem
3382                {
3383                    Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
3384                } else {
3385                    let rep_decompressor = constant_layout
3386                        .rep_compression
3387                        .as_ref()
3388                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3389                        .transpose()?
3390                        .map(Arc::from);
3391
3392                    let def_decompressor = constant_layout
3393                        .def_compression
3394                        .as_ref()
3395                        .map(|encoding| decompressors.create_block_decompressor(encoding))
3396                        .transpose()?
3397                        .map(Arc::from);
3398
3399                    Box::new(ComplexAllNullScheduler::new(
3400                        page_info.buffer_offsets_and_sizes.clone(),
3401                        def_meaning.into(),
3402                        rep_decompressor,
3403                        def_decompressor,
3404                        constant_layout.num_rep_values,
3405                        constant_layout.num_def_values,
3406                    )) as Box<dyn StructuralPageScheduler>
3407                }
3408            }
3409            Layout::BlobLayout(blob) => {
3410                let inner_scheduler = Self::page_layout_to_scheduler(
3411                    page_info,
3412                    blob.inner_layout.as_ref().expect_ok()?.as_ref(),
3413                    decompressors,
3414                    cache_repetition_index,
3415                    target_field,
3416                )?;
3417                let def_meaning = blob
3418                    .layers
3419                    .iter()
3420                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3421                    .collect::<Vec<_>>();
3422                if matches!(target_field.data_type(), DataType::Struct(_)) {
3423                    // User wants to decode blob into struct
3424                    Box::new(BlobDescriptionPageScheduler::new(
3425                        inner_scheduler,
3426                        def_meaning.into(),
3427                    ))
3428                } else {
3429                    // User wants to decode blob into binary data
3430                    Box::new(BlobPageScheduler::new(
3431                        inner_scheduler,
3432                        page_info.priority,
3433                        page_info.num_rows,
3434                        def_meaning.into(),
3435                    ))
3436                }
3437            }
3438        })
3439    }
3440
3441    fn page_info_to_scheduler(
3442        page_info: &PageInfo,
3443        page_index: usize,
3444        decompressors: &dyn DecompressionStrategy,
3445        cache_repetition_index: bool,
3446        target_field: &Field,
3447    ) -> Result<PageInfoAndScheduler> {
3448        let page_layout = page_info.encoding.as_structural();
3449        let scheduler = Self::page_layout_to_scheduler(
3450            page_info,
3451            page_layout,
3452            decompressors,
3453            cache_repetition_index,
3454            target_field,
3455        )?;
3456        Ok(PageInfoAndScheduler {
3457            page_index,
3458            num_rows: page_info.num_rows,
3459            scheduler,
3460        })
3461    }
3462}
3463
3464pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
3465    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
3466}
3467
3468pub struct NoCachedPageData;
3469
3470impl DeepSizeOf for NoCachedPageData {
3471    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
3472        0
3473    }
3474}
3475impl CachedPageData for NoCachedPageData {
3476    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3477        self
3478    }
3479}
3480
3481pub struct CachedFieldData {
3482    pages: Vec<Arc<dyn CachedPageData>>,
3483}
3484
3485impl DeepSizeOf for CachedFieldData {
3486    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
3487        self.pages.deep_size_of_children(ctx)
3488    }
3489}
3490
3491// Cache key for field data
3492//
3493// Both `column_index` and `view_tag` are part of the key because a single
3494// physical column can be decoded under more than one shape — a blob column,
3495// for instance, materializes as a `Struct<position, size>` descriptor in one
3496// scheduler variant and as the raw `LargeBinary` bytes in another. Each
3497// variant builds different `CachedPageData` types per page, so two readers
3498// that hit the same `column_index` with different shapes used to collide and
3499// crash with a downcast failure when loading cached state.
3500#[derive(Debug, Clone)]
3501pub struct FieldDataCacheKey {
3502    pub column_index: u32,
3503    pub view_tag: String,
3504}
3505
3506impl CacheKey for FieldDataCacheKey {
3507    type ValueType = CachedFieldData;
3508
3509    fn key(&self) -> std::borrow::Cow<'_, str> {
3510        format!("{}:{}", self.column_index, self.view_tag).into()
3511    }
3512
3513    fn type_name() -> &'static str {
3514        "FieldData"
3515    }
3516}
3517
3518impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
3519    fn initialize<'a>(
3520        &'a mut self,
3521        _filter: &'a FilterExpression,
3522        context: &'a SchedulerContext,
3523    ) -> BoxFuture<'a, Result<()>> {
3524        let cache_key = FieldDataCacheKey {
3525            column_index: self.column_index,
3526            view_tag: self.view_tag.clone(),
3527        };
3528        let cache = context.cache().clone();
3529
3530        async move {
3531            if let Some(cached_data) = cache.get_with_key(&cache_key).await {
3532                self.page_schedulers
3533                    .iter_mut()
3534                    .zip(cached_data.pages.iter())
3535                    .for_each(|(page_scheduler, cached_data)| {
3536                        page_scheduler.scheduler.load(cached_data);
3537                    });
3538                return Ok(());
3539            }
3540
3541            let page_data = self
3542                .page_schedulers
3543                .iter_mut()
3544                .map(|s| s.scheduler.initialize(context.io()))
3545                .collect::<FuturesOrdered<_>>();
3546
3547            let page_data = page_data.try_collect::<Vec<_>>().await?;
3548            let cached_data = Arc::new(CachedFieldData { pages: page_data });
3549            cache.insert_with_key(&cache_key, cached_data).await;
3550            Ok(())
3551        }
3552        .boxed()
3553    }
3554
3555    fn schedule_ranges<'a>(
3556        &'a self,
3557        ranges: &[Range<u64>],
3558        _filter: &FilterExpression,
3559    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
3560        let ranges = ranges.to_vec();
3561        Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::new(
3562            self, ranges,
3563        )))
3564    }
3565}
3566
3567/// Takes the output from several pages decoders and
3568/// concatenates them.
3569#[derive(Debug)]
3570pub struct StructuralCompositeDecodeArrayTask {
3571    tasks: Vec<Box<dyn DecodePageTask>>,
3572    should_validate: bool,
3573    data_type: DataType,
3574}
3575
3576impl StructuralCompositeDecodeArrayTask {
3577    fn restore_validity(
3578        array: Arc<dyn Array>,
3579        unraveler: &mut CompositeRepDefUnraveler,
3580    ) -> Arc<dyn Array> {
3581        let validity = unraveler.unravel_validity(array.len());
3582        let Some(validity) = validity else {
3583            return array;
3584        };
3585        if array.data_type() == &DataType::Null {
3586            // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it
3587            return array;
3588        }
3589        assert_eq!(validity.len(), array.len());
3590        // SAFETY: We've should have already asserted the buffers are all valid, we are just
3591        // adding null buffers to the array here
3592        make_array(unsafe {
3593            array
3594                .to_data()
3595                .into_builder()
3596                .nulls(Some(validity))
3597                .build_unchecked()
3598        })
3599    }
3600}
3601
3602impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
3603    fn decode(self: Box<Self>) -> Result<DecodedArray> {
3604        let mut arrays = Vec::with_capacity(self.tasks.len());
3605        let mut unravelers = Vec::with_capacity(self.tasks.len());
3606        let mut data_size = 0u64;
3607        for task in self.tasks {
3608            let decoded = task.decode()?;
3609            data_size += decoded.data.data_size();
3610            unravelers.push(decoded.repdef);
3611
3612            let array = make_array(
3613                decoded
3614                    .data
3615                    .into_arrow(self.data_type.clone(), self.should_validate)?,
3616            );
3617
3618            arrays.push(array);
3619        }
3620        let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
3621        let array = arrow_select::concat::concat(&array_refs)?;
3622        let mut repdef = CompositeRepDefUnraveler::new(unravelers);
3623
3624        let array = Self::restore_validity(array, &mut repdef);
3625
3626        Ok(DecodedArray {
3627            array,
3628            repdef,
3629            data_size,
3630        })
3631    }
3632}
3633
3634#[derive(Debug)]
3635pub struct StructuralPrimitiveFieldDecoder {
3636    field: Arc<ArrowField>,
3637    page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
3638    should_validate: bool,
3639    rows_drained_in_current: u64,
3640}
3641
3642impl StructuralPrimitiveFieldDecoder {
3643    pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
3644        Self {
3645            field: field.clone(),
3646            page_decoders: VecDeque::new(),
3647            should_validate,
3648            rows_drained_in_current: 0,
3649        }
3650    }
3651}
3652
3653impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
3654    fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
3655        assert!(child.path.is_empty());
3656        self.page_decoders.push_back(child.decoder);
3657        Ok(())
3658    }
3659
3660    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
3661        let mut remaining = num_rows;
3662        let mut tasks = Vec::new();
3663        while remaining > 0 {
3664            let cur_page = self.page_decoders.front_mut().unwrap();
3665            let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
3666            let to_take = num_in_page.min(remaining);
3667
3668            let task = cur_page.drain(to_take)?;
3669            tasks.push(task);
3670
3671            if to_take == num_in_page {
3672                self.page_decoders.pop_front();
3673                self.rows_drained_in_current = 0;
3674            } else {
3675                self.rows_drained_in_current += to_take;
3676            }
3677
3678            remaining -= to_take;
3679        }
3680        Ok(Box::new(StructuralCompositeDecodeArrayTask {
3681            tasks,
3682            should_validate: self.should_validate,
3683            data_type: self.field.data_type().clone(),
3684        }))
3685    }
3686
3687    fn data_type(&self) -> &DataType {
3688        self.field.data_type()
3689    }
3690}
3691
3692/// The serialized representation of full-zip data
3693struct SerializedFullZip {
3694    /// The zipped values buffer
3695    values: LanceBuffer,
3696    /// The repetition index (only present if there is repetition)
3697    repetition_index: Option<LanceBuffer>,
3698}
3699
3700// We align and pad mini-blocks to 8 byte boundaries for two reasons.  First,
3701// to allow us to store a chunk size in 12 bits.
3702//
3703// If we directly record the size in bytes with 12 bits we would be limited to
3704// 4KiB which is too small.  Since we know each mini-block consists of 8 byte
3705// words we can store the # of words instead which gives us 32KiB.  We want
3706// at least 24KiB so we can handle even the worst case of
3707// - 4Ki values compressed into an 8186 byte buffer
3708// - 4 bytes to describe rep & def lengths
3709// - 16KiB of rep & def buffer (this will almost never happen but life is easier if we
3710//   plan for it)
3711//
3712// Second, each chunk in a mini-block is aligned to 8 bytes.  This allows multi-byte
3713// values like offsets to be stored in a mini-block and safely read back out.  It also
3714// helps ensure zero-copy reads in cases where zero-copy is possible (e.g. no decoding
3715// needed).
3716//
3717// Note: by "aligned to 8 bytes" we mean BOTH "aligned to 8 bytes from the start of
3718// the page" and "aligned to 8 bytes from the start of the file."
3719const MINIBLOCK_ALIGNMENT: usize = 8;
3720
3721/// An encoder for primitive (leaf) arrays
3722///
3723/// This encoder is fairly complicated and follows a number of paths depending
3724/// on the data.
3725///
3726/// First, we convert the validity & offsets information into repetition and
3727/// definition levels.  Then we compress the data itself into a single buffer.
3728///
3729/// If the data is narrow then we encode the data in small chunks (each chunk
3730/// should be a few disk sectors and contains a buffer of repetition, a buffer
3731/// of definition, and a buffer of value data).  This approach is called
3732/// "mini-block".  These mini-blocks are stored into a single data buffer.
3733///
3734/// If the data is wide then we zip together the repetition and definition value
3735/// with the value data into a single buffer.  This approach is called "zipped".
3736///
3737/// If there is any repetition information then we create a repetition index
3738///
3739/// In addition, the compression process may create zero or more metadata buffers.
3740/// For example, a dictionary compression will create dictionary metadata.  Any
3741/// mini-block approach has a metadata buffer of block sizes.  This metadata is
3742/// stored in a separate buffer on disk and read at initialization time.
3743///
3744/// TODO: We should concatenate metadata buffers from all pages into a single buffer
3745/// at (roughly) the end of the file so there is, at most, one read per column of
3746/// metadata per file.
3747pub struct PrimitiveStructuralEncoder {
3748    // Accumulates arrays until we have enough data to justify a disk page
3749    accumulation_queue: AccumulationQueue,
3750
3751    keep_original_array: bool,
3752    support_large_chunk: bool,
3753    accumulated_repdefs: Vec<RepDefBuilder>,
3754    // The compression strategy we will use to compress the data
3755    compression_strategy: Arc<dyn CompressionStrategy>,
3756    column_index: u32,
3757    field: Field,
3758    encoding_metadata: Arc<HashMap<String, String>>,
3759    version: LanceFileVersion,
3760}
3761
3762struct CompressedLevelsChunk {
3763    data: LanceBuffer,
3764    num_levels: u16,
3765}
3766
3767struct CompressedLevels {
3768    data: Vec<CompressedLevelsChunk>,
3769    compression: CompressiveEncoding,
3770    rep_index: Option<LanceBuffer>,
3771}
3772
3773struct SerializedMiniBlockPage {
3774    num_buffers: u64,
3775    data: LanceBuffer,
3776    metadata: LanceBuffer,
3777}
3778
3779#[derive(Debug, Clone, Copy)]
3780struct DictEncodingBudget {
3781    max_dict_entries: u32,
3782    max_encoded_size: usize,
3783}
3784
3785impl PrimitiveStructuralEncoder {
3786    pub fn try_new(
3787        options: &EncodingOptions,
3788        compression_strategy: Arc<dyn CompressionStrategy>,
3789        column_index: u32,
3790        field: Field,
3791        encoding_metadata: Arc<HashMap<String, String>>,
3792    ) -> Result<Self> {
3793        Ok(Self {
3794            accumulation_queue: AccumulationQueue::new(
3795                options.cache_bytes_per_column,
3796                column_index,
3797                options.keep_original_array,
3798            ),
3799            support_large_chunk: options.support_large_chunk(),
3800            keep_original_array: options.keep_original_array,
3801            accumulated_repdefs: Vec::new(),
3802            column_index,
3803            compression_strategy,
3804            field,
3805            encoding_metadata,
3806            version: options.version,
3807        })
3808    }
3809
3810    // TODO: This is a heuristic we may need to tune at some point
3811    //
3812    // As data gets narrow then the "zipping" process gets too expensive
3813    //   and we prefer mini-block
3814    // As data gets wide then the # of values per block shrinks (very wide)
3815    //   data doesn't even fit in a mini-block and the block overhead gets
3816    //   too large and we prefer zipped.
3817    fn is_narrow(data_block: &DataBlock) -> bool {
3818        const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
3819
3820        if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
3821            let max_len_array = max_len_array
3822                .as_any()
3823                .downcast_ref::<PrimitiveArray<UInt64Type>>()
3824                .unwrap();
3825            if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
3826                return true;
3827            }
3828        }
3829        false
3830    }
3831
3832    fn prefers_miniblock(
3833        data_block: &DataBlock,
3834        encoding_metadata: &HashMap<String, String>,
3835    ) -> bool {
3836        // If the user specifically requested miniblock then use it
3837        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3838            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
3839        }
3840        // Otherwise only use miniblock if it is narrow
3841        Self::is_narrow(data_block)
3842    }
3843
3844    /// Checks if the rep/def levels are too sparse for miniblock encoding.
3845    ///
3846    /// Two distinct conditions both force a fallback to fullzip:
3847    ///
3848    /// 1. **Compressed rep+def bytes overflow the per-chunk budget**: each
3849    ///    miniblock chunk only has ~16KiB for rep+def buffers combined, and
3850    ///    if the levels-to-values ratio is high enough that a chunk's
3851    ///    bit-packed levels would exceed that budget, we can't use miniblock.
3852    ///    This is a *global* ratio check that captures uniformly-sparse data
3853    ///    (e.g., 2.5M rows with 100 non-empty lists scattered across).
3854    ///
3855    /// 2. **A single chunk's `num_levels` overflows the `u16` header field**
3856    ///    on disk. v2.2 is a stable format whose per-chunk metadata stores
3857    ///    `num_levels` as a `u16`, capping each chunk at 65 535 levels. The
3858    ///    *global* ratio above can look healthy (HNSW: ~1.19 levels per value)
3859    ///    while one chunk — typically the final chunk consumed via
3860    ///    `slice_rest`, which absorbs every trailing empty list — packs far
3861    ///    more than 65 535 levels. We catch this with a single linear pass
3862    ///    over the def levels that simulates the encoder's chunking
3863    ///    (`MAX_MINIBLOCK_VALUES` visible values per non-last chunk, last
3864    ///    chunk takes the rest) and checks the worst-case per-chunk level
3865    ///    count against `u16::MAX`.
3866    ///
3867    /// Cheap on the dense path: when there are no rep/def levels (no nesting,
3868    /// no nullability above the leaf) we early-exit without scanning anything.
3869    fn repdef_too_sparse_for_miniblock(
3870        repdef: &crate::repdef::SerializedRepDefs,
3871        num_values: u64,
3872    ) -> bool {
3873        if num_values == 0 {
3874            return false;
3875        }
3876        let num_levels = repdef
3877            .repetition_levels
3878            .as_ref()
3879            .map(|r| r.len() as u64)
3880            .max(repdef.definition_levels.as_ref().map(|d| d.len() as u64))
3881            .unwrap_or(0);
3882        if num_levels == 0 {
3883            return false;
3884        }
3885
3886        // Compute bits needed per level for each buffer (ceil of log2(max+1))
3887        let bits_per_rep = repdef
3888            .repetition_levels
3889            .as_ref()
3890            .and_then(|r| r.iter().max().copied())
3891            .map(|max_val| u16::BITS - max_val.leading_zeros())
3892            .unwrap_or(0) as u64;
3893        let bits_per_def = repdef
3894            .definition_levels
3895            .as_ref()
3896            .and_then(|d| d.iter().max().copied())
3897            .map(|max_val| u16::BITS - max_val.leading_zeros())
3898            .unwrap_or(0) as u64;
3899
3900        let bits_per_level = bits_per_rep + bits_per_def;
3901        if bits_per_level == 0 {
3902            return false;
3903        }
3904
3905        // (1) Global rep+def budget check.
3906        const REPDEF_BUDGET_BITS: u64 = 16 * 1024 * 8;
3907        let max_levels_per_chunk = REPDEF_BUDGET_BITS / bits_per_level;
3908        let levels_per_chunk =
3909            (num_levels as f64 / num_values as f64) * *miniblock::MAX_MINIBLOCK_VALUES as f64;
3910        if levels_per_chunk > max_levels_per_chunk as f64 {
3911            return true;
3912        }
3913
3914        // (2) Per-chunk num_levels u16-overflow check. The HNSW persistence
3915        // shape (~40k dense level-0 lists followed by ~6x as many empty
3916        // higher-level rows) has a low global ratio (~1.19) so (1) passes,
3917        // but the trailing empties cluster into the final chunk via
3918        // `slice_rest` and overflow.
3919        Self::any_chunk_levels_overflow_u16(repdef, num_values)
3920    }
3921
3922    /// Simulates the miniblock chunking the encoder is about to perform and
3923    /// returns `true` if any chunk would carry more levels than the on-disk
3924    /// `u16 num_levels` header can express (`> u16::MAX`).
3925    ///
3926    /// Mirrors [`RepDefSlicer::slice_next`] semantics: each non-last chunk
3927    /// consumes levels until it has accumulated `MAX_MINIBLOCK_VALUES`
3928    /// *visible* values (def level ≤ max_visible), leaving any boundary
3929    /// invisibles after the chunk's last visible value for the next chunk; the
3930    /// final chunk is via `slice_rest`, which scoops up every remaining level
3931    /// (including any trailing all-empty rows that no value chunk would
3932    /// otherwise have claimed).
3933    fn any_chunk_levels_overflow_u16(
3934        repdef: &crate::repdef::SerializedRepDefs,
3935        num_values: u64,
3936    ) -> bool {
3937        // Without rep levels the encoder uses a 1:1 levels-to-values mapping,
3938        // so each chunk has at most MAX_MINIBLOCK_VALUES levels — well under
3939        // u16::MAX. Nothing to check on the dense / nullable-leaf-only path.
3940        let Some(max_visible) = repdef.max_visible_level else {
3941            return false;
3942        };
3943        let Some(def_levels) = repdef.definition_levels.as_ref() else {
3944            return false;
3945        };
3946        // If even the *total* level count fits in u16, no chunk can overflow,
3947        // so the linear scan would always return false. Cheap O(1) shortcut on
3948        // small/medium pages.
3949        let u16_max = u16::MAX as u64;
3950        if (def_levels.len() as u64) <= u16_max {
3951            return false;
3952        }
3953        let max_visible_per_chunk = *miniblock::MAX_MINIBLOCK_VALUES;
3954
3955        let mut visible_in_chunk: u64 = 0;
3956        let mut levels_in_current_chunk: u64 = 0;
3957        let mut visible_seen_total: u64 = 0;
3958        for &def_lvl in def_levels.iter() {
3959            levels_in_current_chunk += 1;
3960            // Short-circuit: if the current chunk already exceeds u16, no
3961            // further scanning is needed.
3962            if levels_in_current_chunk > u16_max {
3963                return true;
3964            }
3965            if def_lvl <= max_visible {
3966                visible_in_chunk += 1;
3967                visible_seen_total += 1;
3968                let is_last_value = visible_seen_total == num_values;
3969                if visible_in_chunk == max_visible_per_chunk && !is_last_value {
3970                    // Encoder would call slice_next here. Boundary invisibles
3971                    // (none yet — we just finished a visible value) carry to
3972                    // the next chunk.
3973                    visible_in_chunk = 0;
3974                    levels_in_current_chunk = 0;
3975                }
3976            }
3977        }
3978        // Final chunk (slice_rest) takes everything left in the current
3979        // accumulator, including any trailing empties after the last value.
3980        levels_in_current_chunk > u16_max
3981    }
3982
3983    fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
3984        // Fullzip is the backup option so the only reason we wouldn't use it is if the
3985        // user specifically requested not to use it (in which case we're probably going
3986        // to emit an error)
3987        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
3988            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
3989        }
3990        true
3991    }
3992
3993    // Converts value data, repetition levels, and definition levels into a single
3994    // buffer of mini-blocks.  In addition, creates a buffer of mini-block metadata
3995    // which tells us the size of each block.  Finally, if repetition is present then
3996    // we also create a buffer for the repetition index.
3997    //
3998    // Each chunk is serialized as:
3999    // | num_bufs (1 byte) | buf_lens (2 bytes per buffer) | P | buf0 | P | buf1 | ... | bufN | P |
4000    //
4001    // P - Padding inserted to ensure each buffer is 8-byte aligned and the buffer size is a multiple
4002    //     of 8 bytes (so that the next chunk is 8-byte aligned).
4003    //
4004    // Each block has a u16 word of metadata.  The upper 12 bits contain the
4005    // # of 8-byte words in the block (if the block does not fill the final word
4006    // then up to 7 bytes of padding are added).  The lower 4 bits describe the log_2
4007    // number of values (e.g. if there are 1024 then the lower 4 bits will be
4008    // 0xA)  All blocks except the last must have power-of-two number of values.
4009    // This not only makes metadata smaller but it makes decoding easier since
4010    // batch sizes are typically a power of 2.  4 bits would allow us to express
4011    // up to 16Ki values but we restrict this further to 4Ki values.
4012    //
4013    // This means blocks can have 1 to 4Ki values and 8 - 32Ki bytes.
4014    //
4015    // All metadata words are serialized (as little endian) into a single buffer
4016    // of metadata values.
4017    //
4018    // If there is repetition then we also create a repetition index.  This is a
4019    // single buffer of integer vectors (stored in row major order).  There is one
4020    // entry for each chunk.  The size of the vector is based on the depth of random
4021    // access we want to support.
4022    //
4023    // A vector of size 2 is the minimum and will support row-based random access (e.g.
4024    // "take the 57th row").  A vector of size 3 will support 1 level of nested access
4025    // (e.g. "take the 3rd item in the 57th row").  A vector of size 4 will support 2
4026    // levels of nested access and so on.
4027    //
4028    // The first number in the vector is the number of top-level rows that complete in
4029    // the chunk.  The second number is the number of second-level rows that complete
4030    // after the final top-level row completed (or beginning of the chunk if no top-level
4031    // row completes in the chunk).  And so on.  The final number in the vector is always
4032    // the number of leftover items not covered by earlier entries in the vector.
4033    //
4034    // Currently we are limited to 0 levels of nested access but that will change in the
4035    // future.
4036    //
4037    // The repetition index and the chunk metadata are read at initialization time and
4038    // cached in memory.
4039    fn serialize_miniblocks(
4040        miniblocks: MiniBlockCompressed,
4041        rep: Option<Vec<CompressedLevelsChunk>>,
4042        def: Option<Vec<CompressedLevelsChunk>>,
4043        support_large_chunk: bool,
4044    ) -> Result<SerializedMiniBlockPage> {
4045        let bytes_rep = rep
4046            .as_ref()
4047            .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
4048            .unwrap_or(0);
4049        let bytes_def = def
4050            .as_ref()
4051            .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
4052            .unwrap_or(0);
4053        let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
4054        let mut num_buffers = miniblocks.data.len();
4055        if rep.is_some() {
4056            num_buffers += 1;
4057        }
4058        if def.is_some() {
4059            num_buffers += 1;
4060        }
4061        // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer
4062        let max_extra = 9 * num_buffers;
4063        let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
4064        let chunk_size_bytes = if support_large_chunk { 4 } else { 2 };
4065        let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
4066
4067        let mut rep_iter = rep.map(|r| r.into_iter());
4068        let mut def_iter = def.map(|d| d.into_iter());
4069
4070        let mut buffer_offsets = vec![0; miniblocks.data.len()];
4071        for chunk in miniblocks.chunks {
4072            let start_pos = data_buffer.len();
4073            // Start of chunk should be aligned
4074            debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
4075
4076            let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
4077            let def = def_iter.as_mut().map(|d| d.next().unwrap());
4078
4079            // Write the number of levels, or 0 if there is no rep/def
4080            let num_levels = rep
4081                .as_ref()
4082                .map(|r| r.num_levels)
4083                .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
4084            data_buffer.extend_from_slice(&num_levels.to_le_bytes());
4085
4086            // Write the buffer lengths
4087            if let Some(rep) = rep.as_ref() {
4088                let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
4089                    Error::internal(format!(
4090                        "Repetition buffer size ({} bytes) too large",
4091                        rep.data.len()
4092                    ))
4093                })?;
4094                data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
4095            }
4096            if let Some(def) = def.as_ref() {
4097                let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
4098                    Error::internal(format!(
4099                        "Definition buffer size ({} bytes) too large",
4100                        def.data.len()
4101                    ))
4102                })?;
4103                data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
4104            }
4105
4106            if support_large_chunk {
4107                for &buffer_size in &chunk.buffer_sizes {
4108                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
4109                }
4110            } else {
4111                for &buffer_size in &chunk.buffer_sizes {
4112                    data_buffer.extend_from_slice(&(buffer_size as u16).to_le_bytes());
4113                }
4114            }
4115
4116            // Pad
4117            let add_padding = |data_buffer: &mut Vec<u8>| {
4118                let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
4119                data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
4120            };
4121            add_padding(&mut data_buffer);
4122
4123            // Write the buffers themselves
4124            if let Some(rep) = rep.as_ref() {
4125                data_buffer.extend_from_slice(&rep.data);
4126                add_padding(&mut data_buffer);
4127            }
4128            if let Some(def) = def.as_ref() {
4129                data_buffer.extend_from_slice(&def.data);
4130                add_padding(&mut data_buffer);
4131            }
4132            for (buffer_size, (buffer, buffer_offset)) in chunk
4133                .buffer_sizes
4134                .iter()
4135                .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
4136            {
4137                let start = *buffer_offset;
4138                let end = start + *buffer_size as usize;
4139                *buffer_offset += *buffer_size as usize;
4140                data_buffer.extend_from_slice(&buffer[start..end]);
4141                add_padding(&mut data_buffer);
4142            }
4143
4144            let chunk_bytes = data_buffer.len() - start_pos;
4145            let max_chunk_size = if support_large_chunk {
4146                4 * 1024 * 1024 * 1024 // 4GB limit with u32 metadata
4147            } else {
4148                32 * 1024 // 32KiB limit with u16 metadata
4149            };
4150            assert!(chunk_bytes <= max_chunk_size);
4151            assert!(chunk_bytes > 0);
4152            assert_eq!(chunk_bytes % 8, 0);
4153            // 4Ki values max
4154            assert!(chunk.log_num_values <= 12);
4155            // We subtract 1 here from chunk_bytes because we want to be able to express
4156            // a size of 32KiB and not (32Ki - 8)B which is what we'd get otherwise with
4157            // 0xFFF
4158            let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
4159            let divided_bytes_minus_one = (divided_bytes - 1) as u64;
4160
4161            let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
4162            if support_large_chunk {
4163                meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
4164            } else {
4165                meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
4166            }
4167        }
4168
4169        let data_buffer = LanceBuffer::from(data_buffer);
4170        let metadata_buffer = LanceBuffer::from(meta_buffer);
4171
4172        Ok(SerializedMiniBlockPage {
4173            num_buffers: miniblocks.data.len() as u64,
4174            data: data_buffer,
4175            metadata: metadata_buffer,
4176        })
4177    }
4178
4179    /// Compresses a buffer of levels into chunks
4180    ///
4181    /// If these are repetition levels then we also calculate the repetition index here (that
4182    /// is the third return value)
4183    fn compress_levels(
4184        mut levels: RepDefSlicer<'_>,
4185        num_elements: u64,
4186        compression_strategy: &dyn CompressionStrategy,
4187        chunks: &[MiniBlockChunk],
4188        // This will be 0 if we are compressing def levels
4189        max_rep: u16,
4190    ) -> Result<CompressedLevels> {
4191        let mut rep_index = if max_rep > 0 {
4192            Vec::with_capacity(chunks.len())
4193        } else {
4194            vec![]
4195        };
4196        // Make the levels into a FixedWidth data block
4197        let num_levels = levels.num_levels() as u64;
4198        let levels_buf = levels.all_levels().clone();
4199
4200        let mut fixed_width_block = FixedWidthDataBlock {
4201            data: levels_buf,
4202            bits_per_value: 16,
4203            num_values: num_levels,
4204            block_info: BlockInfo::new(),
4205        };
4206        // Compute statistics to enable optimal compression for rep/def levels
4207        fixed_width_block.compute_stat();
4208
4209        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4210        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4211        // Pick a block compressor
4212        let (compressor, compressor_desc) =
4213            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4214        // Compress blocks of levels (sized according to the chunks)
4215        let mut level_chunks = Vec::with_capacity(chunks.len());
4216        let mut values_counter = 0;
4217        for (chunk_idx, chunk) in chunks.iter().enumerate() {
4218            let chunk_num_values = chunk.num_values(values_counter, num_elements);
4219            debug_assert!(chunk_num_values > 0);
4220            values_counter += chunk_num_values;
4221            let chunk_levels = if chunk_idx < chunks.len() - 1 {
4222                levels.slice_next(chunk_num_values as usize)
4223            } else {
4224                levels.slice_rest()
4225            };
4226            let num_chunk_levels = (chunk_levels.len() / 2) as u64;
4227            if max_rep > 0 {
4228                // If max_rep > 0 then we are working with rep levels and we need
4229                // to calculate the repetition index.  The repetition index for a
4230                // chunk is currently 2 values (in the future it may be more).
4231                //
4232                // The first value is the number of rows that _finish_ in the
4233                // chunk.
4234                //
4235                // The second value is the number of "leftovers" after the last
4236                // finished row in the chunk.
4237                let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
4238                let rep_values = rep_values.as_ref();
4239
4240                // We skip 1 here because a max_rep at spot 0 doesn't count as a finished list (we
4241                // will count it in the previous chunk)
4242                let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
4243                let num_leftovers = if chunk_idx < chunks.len() - 1 {
4244                    rep_values
4245                        .iter()
4246                        .rev()
4247                        .position(|v| *v == max_rep)
4248                        // # of leftovers includes the max_rep spot
4249                        .map(|pos| pos + 1)
4250                        .unwrap_or(rep_values.len())
4251                } else {
4252                    // Last chunk can't have leftovers
4253                    0
4254                };
4255
4256                if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
4257                    // This chunk starts with a new row and so, if we thought we had leftovers
4258                    // in the previous chunk, we were mistaken
4259                    // TODO: Can use unchecked here
4260                    let rep_len = rep_index.len();
4261                    if rep_index[rep_len - 1] != 0 {
4262                        // We thought we had leftovers but that was actually a full row
4263                        rep_index[rep_len - 2] += 1;
4264                        rep_index[rep_len - 1] = 0;
4265                    }
4266                }
4267
4268                if chunk_idx == chunks.len() - 1 {
4269                    // The final list
4270                    num_rows += 1;
4271                }
4272                rep_index.push(num_rows as u64);
4273                rep_index.push(num_leftovers as u64);
4274            }
4275            let mut chunk_fixed_width = FixedWidthDataBlock {
4276                data: chunk_levels,
4277                bits_per_value: 16,
4278                num_values: num_chunk_levels,
4279                block_info: BlockInfo::new(),
4280            };
4281            chunk_fixed_width.compute_stat();
4282            let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
4283            let compressed_levels = compressor.compress(chunk_levels_block)?;
4284            // `num_levels` is a `u16` in the v2.2 on-disk per-chunk header
4285            // (`primitive.rs:467` reads it as `u16::from_le_bytes`). Silently
4286            // truncating here used to corrupt the round-trip when chunking
4287            // packed >65 535 levels into one chunk (HNSW persistence shape;
4288            // see `repdef_too_sparse_for_miniblock`). The heuristic should
4289            // route such shapes to fullzip before reaching this point; if it
4290            // doesn't (e.g. a future shape sneaks past the heuristic), surface
4291            // a clear input-shape error rather than a silent corruption.
4292            let num_levels_u16 = u16::try_from(num_chunk_levels).map_err(|_| {
4293                Error::invalid_input_source(
4294                    format!(
4295                        "miniblock chunk has {} levels but per-chunk header is u16 \
4296                         (max {}); this shape needs fullzip encoding",
4297                        num_chunk_levels,
4298                        u16::MAX,
4299                    )
4300                    .into(),
4301                )
4302            })?;
4303            level_chunks.push(CompressedLevelsChunk {
4304                data: compressed_levels,
4305                num_levels: num_levels_u16,
4306            });
4307        }
4308        debug_assert_eq!(levels.num_levels_remaining(), 0);
4309        let rep_index = if rep_index.is_empty() {
4310            None
4311        } else {
4312            Some(LanceBuffer::reinterpret_vec(rep_index))
4313        };
4314        Ok(CompressedLevels {
4315            data: level_chunks,
4316            compression: compressor_desc,
4317            rep_index,
4318        })
4319    }
4320
4321    fn encode_simple_all_null(
4322        column_idx: u32,
4323        num_rows: u64,
4324        row_number: u64,
4325    ) -> Result<EncodedPage> {
4326        let description =
4327            ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
4328        Ok(EncodedPage {
4329            column_idx,
4330            data: vec![],
4331            description: PageEncoding::Structural(description),
4332            num_rows,
4333            row_number,
4334        })
4335    }
4336
4337    fn encode_complex_all_null_vals(
4338        data: &Arc<[u16]>,
4339        compression_strategy: &dyn CompressionStrategy,
4340    ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
4341        let buffer = LanceBuffer::reinterpret_slice(data.clone());
4342        let mut fixed_width_block = FixedWidthDataBlock {
4343            data: buffer,
4344            bits_per_value: 16,
4345            num_values: data.len() as u64,
4346            block_info: BlockInfo::new(),
4347        };
4348        fixed_width_block.compute_stat();
4349
4350        let levels_block = DataBlock::FixedWidth(fixed_width_block);
4351        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
4352        let (compressor, encoding) =
4353            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
4354        let compressed_buffer = compressor.compress(levels_block)?;
4355        Ok((compressed_buffer, encoding))
4356    }
4357
4358    // Encodes a page where all values are null but we have rep/def
4359    // information that we need to store (e.g. to distinguish between
4360    // different kinds of null)
4361    fn encode_complex_all_null(
4362        column_idx: u32,
4363        repdef: crate::repdef::SerializedRepDefs,
4364        row_number: u64,
4365        num_rows: u64,
4366        version: LanceFileVersion,
4367        compression_strategy: &dyn CompressionStrategy,
4368    ) -> Result<EncodedPage> {
4369        if version.resolve() < LanceFileVersion::V2_2 {
4370            let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
4371                LanceBuffer::reinterpret_slice(rep.clone())
4372            } else {
4373                LanceBuffer::empty()
4374            };
4375
4376            let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
4377                LanceBuffer::reinterpret_slice(def.clone())
4378            } else {
4379                LanceBuffer::empty()
4380            };
4381
4382            let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
4383            return Ok(EncodedPage {
4384                column_idx,
4385                data: vec![rep_bytes, def_bytes],
4386                description: PageEncoding::Structural(description),
4387                num_rows,
4388                row_number,
4389            });
4390        }
4391
4392        let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
4393            repdef.repetition_levels.as_ref()
4394        {
4395            let num_values = rep.len() as u64;
4396            let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
4397            (buffer, Some(encoding), num_values)
4398        } else {
4399            (LanceBuffer::empty(), None, 0)
4400        };
4401
4402        let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
4403            repdef.definition_levels.as_ref()
4404        {
4405            let num_values = def.len() as u64;
4406            let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
4407            (buffer, Some(encoding), num_values)
4408        } else {
4409            (LanceBuffer::empty(), None, 0)
4410        };
4411
4412        let description = ProtobufUtils21::compressed_all_null_constant_layout(
4413            &repdef.def_meaning,
4414            rep_encoding,
4415            def_encoding,
4416            num_rep_values,
4417            num_def_values,
4418        );
4419        Ok(EncodedPage {
4420            column_idx,
4421            data: vec![rep_bytes, def_bytes],
4422            description: PageEncoding::Structural(description),
4423            num_rows,
4424            row_number,
4425        })
4426    }
4427
4428    fn leaf_validity(
4429        repdef: &crate::repdef::SerializedRepDefs,
4430        num_values: usize,
4431    ) -> Result<Option<BooleanBuffer>> {
4432        let rep = repdef
4433            .repetition_levels
4434            .as_ref()
4435            .map(|rep| rep.as_ref().to_vec());
4436        let def = repdef
4437            .definition_levels
4438            .as_ref()
4439            .map(|def| def.as_ref().to_vec());
4440        let mut unraveler = RepDefUnraveler::new(
4441            rep,
4442            def,
4443            repdef.def_meaning.clone().into(),
4444            num_values as u64,
4445        );
4446        if unraveler.is_all_valid() {
4447            return Ok(None);
4448        }
4449        let mut validity = BooleanBufferBuilder::new(num_values);
4450        unraveler.unravel_validity(&mut validity);
4451        Ok(Some(validity.finish()))
4452    }
4453
4454    fn is_constant_values(
4455        arrays: &[ArrayRef],
4456        scalar: &ArrayRef,
4457        validity: Option<&BooleanBuffer>,
4458    ) -> Result<bool> {
4459        debug_assert_eq!(scalar.len(), 1);
4460        debug_assert_eq!(scalar.null_count(), 0);
4461
4462        match scalar.data_type() {
4463            DataType::Boolean => {
4464                let mut global_idx = 0usize;
4465                let scalar_val = scalar.as_boolean().value(0);
4466                for arr in arrays {
4467                    let bool_arr = arr.as_boolean();
4468                    for i in 0..arr.len() {
4469                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4470                        global_idx += 1;
4471                        if !is_valid {
4472                            continue;
4473                        }
4474                        if bool_arr.value(i) != scalar_val {
4475                            return Ok(false);
4476                        }
4477                    }
4478                }
4479                Ok(true)
4480            }
4481            DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
4482            DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
4483            DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
4484            DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
4485            data_type => {
4486                let mut global_idx = 0usize;
4487                let Some(byte_width) = data_type.byte_width_opt() else {
4488                    return Ok(false);
4489                };
4490                let scalar_data = scalar.to_data();
4491                if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
4492                    return Ok(false);
4493                }
4494                let scalar_bytes = scalar_data.buffers()[0].as_slice();
4495                if scalar_bytes.len() != byte_width {
4496                    return Ok(false);
4497                }
4498
4499                for arr in arrays {
4500                    let data = arr.to_data();
4501                    if data.buffers().is_empty() {
4502                        return Ok(false);
4503                    }
4504                    let buf = data.buffers()[0].as_slice();
4505                    let base = data.offset();
4506                    for i in 0..arr.len() {
4507                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4508                        global_idx += 1;
4509                        if !is_valid {
4510                            continue;
4511                        }
4512                        let start = (base + i) * byte_width;
4513                        if buf[start..start + byte_width] != scalar_bytes[..] {
4514                            return Ok(false);
4515                        }
4516                    }
4517                }
4518                Ok(true)
4519            }
4520        }
4521    }
4522
4523    fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
4524        arrays: &[ArrayRef],
4525        scalar: &ArrayRef,
4526        validity: Option<&BooleanBuffer>,
4527    ) -> Result<bool> {
4528        debug_assert_eq!(scalar.len(), 1);
4529        let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
4530        let mut global_idx = 0usize;
4531        for arr in arrays {
4532            let str_arr = arr.as_string::<O>();
4533            for i in 0..arr.len() {
4534                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4535                global_idx += 1;
4536                if !is_valid {
4537                    continue;
4538                }
4539                if str_arr.value(i).as_bytes() != scalar_val {
4540                    return Ok(false);
4541                }
4542            }
4543        }
4544        Ok(true)
4545    }
4546
4547    fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
4548        arrays: &[ArrayRef],
4549        scalar: &ArrayRef,
4550        validity: Option<&BooleanBuffer>,
4551    ) -> Result<bool> {
4552        debug_assert_eq!(scalar.len(), 1);
4553        let scalar_val = scalar.as_binary::<O>().value(0);
4554        let mut global_idx = 0usize;
4555        for arr in arrays {
4556            let bin_arr = arr.as_binary::<O>();
4557            for i in 0..arr.len() {
4558                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
4559                global_idx += 1;
4560                if !is_valid {
4561                    continue;
4562                }
4563                if bin_arr.value(i) != scalar_val {
4564                    return Ok(false);
4565                }
4566            }
4567        }
4568        Ok(true)
4569    }
4570
4571    fn find_constant_scalar(
4572        arrays: &[ArrayRef],
4573        validity: Option<&BooleanBuffer>,
4574    ) -> Result<Option<ArrayRef>> {
4575        if arrays.is_empty() {
4576            return Ok(None);
4577        }
4578
4579        let global_scalar_idx = if let Some(validity) = validity {
4580            let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
4581                return Ok(None);
4582            };
4583            idx
4584        } else {
4585            0
4586        };
4587
4588        let mut idx_remaining = global_scalar_idx;
4589        let mut scalar_arr_idx = 0usize;
4590        while scalar_arr_idx < arrays.len() {
4591            let len = arrays[scalar_arr_idx].len();
4592            if idx_remaining < len {
4593                break;
4594            }
4595            idx_remaining -= len;
4596            scalar_arr_idx += 1;
4597        }
4598
4599        if scalar_arr_idx >= arrays.len() {
4600            return Ok(None);
4601        }
4602
4603        let scalar =
4604            lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
4605        if scalar.null_count() != 0 {
4606            return Ok(None);
4607        }
4608        if !Self::is_constant_values(arrays, &scalar, validity)? {
4609            return Ok(None);
4610        }
4611        Ok(Some(scalar))
4612    }
4613
4614    fn resolve_dict_values_compression_metadata(
4615        field_metadata: &HashMap<String, String>,
4616        env_compression: Option<String>,
4617        env_compression_level: Option<String>,
4618    ) -> HashMap<String, String> {
4619        let mut metadata = HashMap::new();
4620
4621        let compression = field_metadata
4622            .get(DICT_VALUES_COMPRESSION_META_KEY)
4623            .cloned()
4624            .or(env_compression)
4625            .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
4626        metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
4627
4628        if let Some(compression_level) = field_metadata
4629            .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
4630            .cloned()
4631            .or(env_compression_level)
4632        {
4633            metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
4634        }
4635
4636        metadata
4637    }
4638
4639    fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
4640        // This is an internal synthetic field used only to feed metadata into
4641        // `create_block_compressor` for dictionary values. The concrete type/name here
4642        // are not semantically meaningful; we rely on explicit metadata below to control
4643        // general compression selection for dictionary values.
4644        let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
4645        dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
4646            &field.metadata,
4647            env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
4648            env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
4649        );
4650        Ok(dict_values_field)
4651    }
4652
4653    #[allow(clippy::too_many_arguments)]
4654    fn encode_miniblock(
4655        column_idx: u32,
4656        field: &Field,
4657        compression_strategy: &dyn CompressionStrategy,
4658        data: DataBlock,
4659        repdef: crate::repdef::SerializedRepDefs,
4660        row_number: u64,
4661        dictionary_data: Option<DataBlock>,
4662        num_rows: u64,
4663        support_large_chunk: bool,
4664    ) -> Result<EncodedPage> {
4665        if let DataBlock::AllNull(_null_block) = data {
4666            // We should not be using mini-block for all-null.  There are other structural
4667            // encodings for that.
4668            unreachable!()
4669        }
4670
4671        let num_items = data.num_values();
4672
4673        let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
4674        let (compressed_data, value_encoding) = compressor.compress(data)?;
4675
4676        let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
4677
4678        let mut compressed_rep = repdef
4679            .rep_slicer()
4680            .map(|rep_slicer| {
4681                Self::compress_levels(
4682                    rep_slicer,
4683                    num_items,
4684                    compression_strategy,
4685                    &compressed_data.chunks,
4686                    max_rep,
4687                )
4688            })
4689            .transpose()?;
4690
4691        let (rep_index, rep_index_depth) =
4692            match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
4693                Some(rep_index) => (Some(rep_index.clone()), 1),
4694                None => (None, 0),
4695            };
4696
4697        let mut compressed_def = repdef
4698            .def_slicer()
4699            .map(|def_slicer| {
4700                Self::compress_levels(
4701                    def_slicer,
4702                    num_items,
4703                    compression_strategy,
4704                    &compressed_data.chunks,
4705                    /*max_rep=*/ 0,
4706                )
4707            })
4708            .transpose()?;
4709
4710        // TODO: Parquet sparsely encodes values here.  We could do the same but
4711        // then we won't have log2 values per chunk.  This means more metadata
4712        // and potentially more decoder asymmetry.  However, it may be worth
4713        // investigating at some point
4714
4715        let rep_data = compressed_rep
4716            .as_mut()
4717            .map(|cr| std::mem::take(&mut cr.data));
4718        let def_data = compressed_def
4719            .as_mut()
4720            .map(|cd| std::mem::take(&mut cd.data));
4721
4722        let serialized =
4723            Self::serialize_miniblocks(compressed_data, rep_data, def_data, support_large_chunk)?;
4724
4725        // Metadata, Data, Dictionary, (maybe) Repetition Index
4726        let mut data = Vec::with_capacity(4);
4727        data.push(serialized.metadata);
4728        data.push(serialized.data);
4729
4730        if let Some(dictionary_data) = dictionary_data {
4731            let num_dictionary_items = dictionary_data.num_values();
4732            let dict_values_field = Self::build_dict_values_compressor_field(field)?;
4733
4734            let (compressor, dictionary_encoding) = compression_strategy
4735                .create_block_compressor(&dict_values_field, &dictionary_data)?;
4736            let dictionary_buffer = compressor.compress(dictionary_data)?;
4737
4738            data.push(dictionary_buffer);
4739            if let Some(rep_index) = rep_index {
4740                data.push(rep_index);
4741            }
4742
4743            let description = ProtobufUtils21::miniblock_layout(
4744                compressed_rep.map(|cr| cr.compression),
4745                compressed_def.map(|cd| cd.compression),
4746                value_encoding,
4747                rep_index_depth,
4748                serialized.num_buffers,
4749                Some((dictionary_encoding, num_dictionary_items)),
4750                &repdef.def_meaning,
4751                num_items,
4752                support_large_chunk,
4753            );
4754            Ok(EncodedPage {
4755                num_rows,
4756                column_idx,
4757                data,
4758                description: PageEncoding::Structural(description),
4759                row_number,
4760            })
4761        } else {
4762            let description = ProtobufUtils21::miniblock_layout(
4763                compressed_rep.map(|cr| cr.compression),
4764                compressed_def.map(|cd| cd.compression),
4765                value_encoding,
4766                rep_index_depth,
4767                serialized.num_buffers,
4768                None,
4769                &repdef.def_meaning,
4770                num_items,
4771                support_large_chunk,
4772            );
4773
4774            if let Some(rep_index) = rep_index {
4775                let view = rep_index.borrow_to_typed_slice::<u64>();
4776                let total = view.chunks_exact(2).map(|c| c[0]).sum::<u64>();
4777                debug_assert_eq!(total, num_rows);
4778
4779                data.push(rep_index);
4780            }
4781
4782            Ok(EncodedPage {
4783                num_rows,
4784                column_idx,
4785                data,
4786                description: PageEncoding::Structural(description),
4787                row_number,
4788            })
4789        }
4790    }
4791
4792    // For fixed-size data we encode < control word | data > for each value
4793    fn serialize_full_zip_fixed(
4794        fixed: FixedWidthDataBlock,
4795        mut repdef: ControlWordIterator,
4796        num_values: u64,
4797    ) -> SerializedFullZip {
4798        let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
4799        let mut zipped_data = Vec::with_capacity(len);
4800
4801        let max_rep_index_val = if repdef.has_repetition() {
4802            len as u64
4803        } else {
4804            // Setting this to 0 means we won't write a repetition index
4805            0
4806        };
4807        let mut rep_index_builder =
4808            BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
4809
4810        // I suppose we can just pad to the nearest byte but I'm not sure we need to worry about this anytime soon
4811        // because it is unlikely compression of large values is going to yield a result that is not byte aligned
4812        assert_eq!(
4813            fixed.bits_per_value % 8,
4814            0,
4815            "Non-byte aligned full-zip compression not yet supported"
4816        );
4817
4818        let bytes_per_value = fixed.bits_per_value as usize / 8;
4819        let mut offset = 0;
4820
4821        if bytes_per_value == 0 {
4822            // No data, just dump the repdef into the buffer
4823            while let Some(control) = repdef.append_next(&mut zipped_data) {
4824                if control.is_new_row {
4825                    // We have finished a row
4826                    debug_assert!(offset <= len);
4827                    // SAFETY: We know that `start <= len`
4828                    unsafe { rep_index_builder.append(offset as u64) };
4829                }
4830                offset = zipped_data.len();
4831            }
4832        } else {
4833            // We have data, zip it with the repdef
4834            let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
4835            while let Some(control) = repdef.append_next(&mut zipped_data) {
4836                if control.is_new_row {
4837                    // We have finished a row
4838                    debug_assert!(offset <= len);
4839                    // SAFETY: We know that `start <= len`
4840                    unsafe { rep_index_builder.append(offset as u64) };
4841                }
4842                if control.is_visible {
4843                    let value = data_iter.next().unwrap();
4844                    zipped_data.extend_from_slice(value);
4845                }
4846                offset = zipped_data.len();
4847            }
4848        }
4849
4850        debug_assert_eq!(zipped_data.len(), len);
4851        // Put the final value in the rep index
4852        // SAFETY: `zipped_data.len() == len`
4853        unsafe {
4854            rep_index_builder.append(zipped_data.len() as u64);
4855        }
4856
4857        let zipped_data = LanceBuffer::from(zipped_data);
4858        let rep_index = rep_index_builder.into_data();
4859        let rep_index = if rep_index.is_empty() {
4860            None
4861        } else {
4862            Some(LanceBuffer::from(rep_index))
4863        };
4864        SerializedFullZip {
4865            values: zipped_data,
4866            repetition_index: rep_index,
4867        }
4868    }
4869
4870    // For variable-size data we encode < control word | length | data > for each value
4871    //
4872    // In addition, we create a second buffer, the repetition index
4873    fn serialize_full_zip_variable(
4874        variable: VariableWidthBlock,
4875        mut repdef: ControlWordIterator,
4876        num_items: u64,
4877    ) -> SerializedFullZip {
4878        let bytes_per_offset = variable.bits_per_offset as usize / 8;
4879        assert_eq!(
4880            variable.bits_per_offset % 8,
4881            0,
4882            "Only byte-aligned offsets supported"
4883        );
4884        let len = variable.data.len()
4885            + repdef.bytes_per_word() * num_items as usize
4886            + bytes_per_offset * variable.num_values as usize;
4887        let mut buf = Vec::with_capacity(len);
4888
4889        let max_rep_index_val = len as u64;
4890        let mut rep_index_builder =
4891            BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
4892
4893        // TODO: byte pack the item lengths with varint encoding
4894        match bytes_per_offset {
4895            4 => {
4896                let offs = variable.offsets.borrow_to_typed_slice::<u32>();
4897                let mut rep_offset = 0;
4898                let mut windows_iter = offs.as_ref().windows(2);
4899                while let Some(control) = repdef.append_next(&mut buf) {
4900                    if control.is_new_row {
4901                        // We have finished a row
4902                        debug_assert!(rep_offset <= len);
4903                        // SAFETY: We know that `buf.len() <= len`
4904                        unsafe { rep_index_builder.append(rep_offset as u64) };
4905                    }
4906                    if control.is_visible {
4907                        let window = windows_iter.next().unwrap();
4908                        if control.is_valid_item {
4909                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4910                            buf.extend_from_slice(
4911                                &variable.data[window[0] as usize..window[1] as usize],
4912                            );
4913                        }
4914                    }
4915                    rep_offset = buf.len();
4916                }
4917            }
4918            8 => {
4919                let offs = variable.offsets.borrow_to_typed_slice::<u64>();
4920                let mut rep_offset = 0;
4921                let mut windows_iter = offs.as_ref().windows(2);
4922                while let Some(control) = repdef.append_next(&mut buf) {
4923                    if control.is_new_row {
4924                        // We have finished a row
4925                        debug_assert!(rep_offset <= len);
4926                        // SAFETY: We know that `buf.len() <= len`
4927                        unsafe { rep_index_builder.append(rep_offset as u64) };
4928                    }
4929                    if control.is_visible {
4930                        let window = windows_iter.next().unwrap();
4931                        if control.is_valid_item {
4932                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
4933                            buf.extend_from_slice(
4934                                &variable.data[window[0] as usize..window[1] as usize],
4935                            );
4936                        }
4937                    }
4938                    rep_offset = buf.len();
4939                }
4940            }
4941            _ => panic!("Unsupported offset size"),
4942        }
4943
4944        // We might have saved a few bytes by not copying lengths when the length was zero.  However,
4945        // if we are over `len` then we have a bug.
4946        debug_assert!(buf.len() <= len);
4947        // Put the final value in the rep index
4948        // SAFETY: `zipped_data.len() == len`
4949        unsafe {
4950            rep_index_builder.append(buf.len() as u64);
4951        }
4952
4953        let zipped_data = LanceBuffer::from(buf);
4954        let rep_index = rep_index_builder.into_data();
4955        debug_assert!(!rep_index.is_empty());
4956        let rep_index = Some(LanceBuffer::from(rep_index));
4957        SerializedFullZip {
4958            values: zipped_data,
4959            repetition_index: rep_index,
4960        }
4961    }
4962
4963    /// Serializes data into a single buffer according to the full-zip format which zips
4964    /// together the repetition, definition, and value data into a single buffer.
4965    fn serialize_full_zip(
4966        compressed_data: PerValueDataBlock,
4967        repdef: ControlWordIterator,
4968        num_items: u64,
4969    ) -> SerializedFullZip {
4970        match compressed_data {
4971            PerValueDataBlock::Fixed(fixed) => {
4972                Self::serialize_full_zip_fixed(fixed, repdef, num_items)
4973            }
4974            PerValueDataBlock::Variable(var) => {
4975                Self::serialize_full_zip_variable(var, repdef, num_items)
4976            }
4977        }
4978    }
4979
4980    fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
4981        debug_assert_eq!(fixed.bits_per_value, 1);
4982        let num_values = fixed.num_values as usize;
4983        let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
4984        let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
4985        FixedWidthDataBlock {
4986            data: LanceBuffer::from(expanded),
4987            bits_per_value: 8,
4988            num_values: fixed.num_values,
4989            block_info: BlockInfo::new(),
4990        }
4991    }
4992
4993    fn encode_full_zip(
4994        column_idx: u32,
4995        field: &Field,
4996        compression_strategy: &dyn CompressionStrategy,
4997        data: DataBlock,
4998        repdef: crate::repdef::SerializedRepDefs,
4999        row_number: u64,
5000        num_lists: u64,
5001    ) -> Result<EncodedPage> {
5002        let max_rep = repdef
5003            .repetition_levels
5004            .as_ref()
5005            .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
5006        let max_def = repdef
5007            .definition_levels
5008            .as_ref()
5009            .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
5010
5011        // To handle FSL we just flatten
5012        // let data = data.flatten();
5013
5014        let (num_items, num_visible_items) =
5015            if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
5016                // If there are rep levels there may be "invisible" items and we need to encode
5017                // rep_levels.len() things which might be larger than data.num_values()
5018                (rep_levels.len() as u64, data.num_values())
5019            } else {
5020                // If there are no rep levels then we encode data.num_values() things
5021                (data.num_values(), data.num_values())
5022            };
5023
5024        let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
5025
5026        let repdef_iter = build_control_word_iterator(
5027            repdef.repetition_levels.as_deref(),
5028            max_rep,
5029            repdef.definition_levels.as_deref(),
5030            max_def,
5031            max_visible_def,
5032            num_items as usize,
5033        );
5034        let bits_rep = repdef_iter.bits_rep();
5035        let bits_def = repdef_iter.bits_def();
5036
5037        // Full-zip requires byte-aligned values; expand 1-bit booleans to 1 byte each.
5038        let data = match data {
5039            DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
5040                DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
5041            }
5042            other => other,
5043        };
5044
5045        let compressor = compression_strategy.create_per_value(field, &data)?;
5046        let (compressed_data, value_encoding) = compressor.compress(data)?;
5047
5048        let description = match &compressed_data {
5049            PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
5050                bits_rep,
5051                bits_def,
5052                fixed.bits_per_value as u32,
5053                value_encoding,
5054                &repdef.def_meaning,
5055                num_items as u32,
5056                num_visible_items as u32,
5057            ),
5058            PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
5059                bits_rep,
5060                bits_def,
5061                variable.bits_per_offset as u32,
5062                value_encoding,
5063                &repdef.def_meaning,
5064                num_items as u32,
5065                num_visible_items as u32,
5066            ),
5067        };
5068
5069        let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items);
5070
5071        let data = if let Some(repindex) = zipped.repetition_index {
5072            vec![zipped.values, repindex]
5073        } else {
5074            vec![zipped.values]
5075        };
5076
5077        Ok(EncodedPage {
5078            num_rows: num_lists,
5079            column_idx,
5080            data,
5081            description: PageEncoding::Structural(description),
5082            row_number,
5083        })
5084    }
5085
5086    fn should_dictionary_encode(
5087        data_block: &DataBlock,
5088        field: &Field,
5089        version: LanceFileVersion,
5090    ) -> Option<DictEncodingBudget> {
5091        const DEFAULT_SAMPLE_SIZE: usize = 4096;
5092        const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
5093
5094        // Since we only dictionary encode FixedWidth and VariableWidth blocks for now, we skip
5095        // estimating the size for other types.
5096        match data_block {
5097            DataBlock::FixedWidth(fixed) => {
5098                if fixed.bits_per_value == 64 && version < LanceFileVersion::V2_2 {
5099                    return None;
5100                }
5101                if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
5102                    return None;
5103                }
5104                if fixed.bits_per_value % 8 != 0 {
5105                    return None;
5106                }
5107            }
5108            DataBlock::VariableWidth(var) => {
5109                if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
5110                    return None;
5111                }
5112            }
5113            _ => return None,
5114        }
5115
5116        // Don't dictionary encode tiny arrays.
5117        let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
5118            .ok()
5119            .and_then(|val| val.parse().ok())
5120            .unwrap_or(100);
5121        if data_block.num_values() < too_small {
5122            return None;
5123        }
5124
5125        let num_values = data_block.num_values();
5126
5127        // Apply divisor threshold and cap. This is intentionally conservative: the goal is to
5128        // avoid spending too much CPU trying to estimate very high cardinalities.
5129        let divisor: u64 = field
5130            .metadata
5131            .get(DICT_DIVISOR_META_KEY)
5132            .and_then(|val| val.parse().ok())
5133            .or_else(|| {
5134                env::var("LANCE_ENCODING_DICT_DIVISOR")
5135                    .ok()
5136                    .and_then(|val| val.parse().ok())
5137            })
5138            .unwrap_or(DEFAULT_DICT_DIVISOR);
5139
5140        let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
5141            .ok()
5142            .and_then(|val| val.parse().ok())
5143            .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
5144
5145        let threshold_cardinality = num_values
5146            .checked_div(divisor.max(1))
5147            .unwrap_or(0)
5148            .min(max_cardinality);
5149        if threshold_cardinality == 0 {
5150            return None;
5151        }
5152
5153        // Get size ratio from metadata or env var.
5154        let threshold_ratio = field
5155            .metadata
5156            .get(DICT_SIZE_RATIO_META_KEY)
5157            .and_then(|val| val.parse::<f64>().ok())
5158            .or_else(|| {
5159                env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
5160                    .ok()
5161                    .and_then(|val| val.parse().ok())
5162            })
5163            .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
5164
5165        if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
5166            panic!(
5167                "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
5168                threshold_ratio
5169            );
5170        }
5171
5172        let data_size = data_block.data_size();
5173        if data_size == 0 {
5174            return None;
5175        }
5176
5177        let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
5178        let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
5179
5180        // Avoid probing dictionary encoding on data that appears to be near-unique.
5181        if Self::sample_is_near_unique(
5182            data_block,
5183            DEFAULT_SAMPLE_SIZE,
5184            DEFAULT_SAMPLE_UNIQUE_RATIO,
5185        )? {
5186            return None;
5187        }
5188
5189        let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
5190        Some(DictEncodingBudget {
5191            max_dict_entries,
5192            max_encoded_size,
5193        })
5194    }
5195
5196    /// Probe whether a page looks near-unique before attempting dictionary encoding.
5197    ///
5198    /// The probe uses deterministic stride sampling (not RNG sampling), which keeps
5199    /// the check cheap and reproducible across runs. The result is only a gate for
5200    /// whether we try dictionary encoding, not a cardinality statistic.
5201    fn sample_is_near_unique(
5202        data_block: &DataBlock,
5203        max_samples: usize,
5204        unique_ratio_threshold: f64,
5205    ) -> Option<bool> {
5206        use std::collections::HashSet;
5207
5208        if unique_ratio_threshold <= 0.0 || unique_ratio_threshold > 1.0 {
5209            return None;
5210        }
5211
5212        let num_values = usize::try_from(data_block.num_values()).ok()?;
5213        if num_values == 0 {
5214            return Some(false);
5215        }
5216
5217        let sample_count = num_values.min(max_samples).max(1);
5218        // Uniform stride sampling across the page.
5219        let step = (num_values / sample_count).max(1);
5220
5221        match data_block {
5222            DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
5223                64 => {
5224                    let values = fixed.data.borrow_to_typed_slice::<u64>();
5225                    let values = values.as_ref();
5226                    let mut unique: HashSet<u64> = HashSet::with_capacity(sample_count.min(1024));
5227                    for idx in (0..num_values).step_by(step).take(sample_count) {
5228                        unique.insert(values.get(idx).copied()?);
5229                    }
5230                    let ratio = unique.len() as f64 / sample_count as f64;
5231                    // Avoid overreacting to tiny pages with too few samples.
5232                    Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5233                }
5234                128 => {
5235                    let values = fixed.data.borrow_to_typed_slice::<u128>();
5236                    let values = values.as_ref();
5237                    let mut unique: HashSet<u128> = HashSet::with_capacity(sample_count.min(1024));
5238                    for idx in (0..num_values).step_by(step).take(sample_count) {
5239                        unique.insert(values.get(idx).copied()?);
5240                    }
5241                    let ratio = unique.len() as f64 / sample_count as f64;
5242                    Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5243                }
5244                _ => Some(false),
5245            },
5246            DataBlock::VariableWidth(var) => {
5247                use xxhash_rust::xxh3::xxh3_64;
5248
5249                // Hash variable-width slices instead of storing borrowed slice keys.
5250                let mut unique: HashSet<u64> = HashSet::with_capacity(sample_count.min(1024));
5251                match var.bits_per_offset {
5252                    32 => {
5253                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
5254                        let offsets: &[u32] = offsets_ref.as_ref();
5255                        for i in (0..num_values).step_by(step).take(sample_count) {
5256                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5257                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5258                            if start > end || end > var.data.len() {
5259                                return None;
5260                            }
5261                            unique.insert(xxh3_64(&var.data[start..end]));
5262                        }
5263                    }
5264                    64 => {
5265                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
5266                        let offsets: &[u64] = offsets_ref.as_ref();
5267                        for i in (0..num_values).step_by(step).take(sample_count) {
5268                            let start = usize::try_from(*offsets.get(i)?).ok()?;
5269                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
5270                            if start > end || end > var.data.len() {
5271                                return None;
5272                            }
5273                            unique.insert(xxh3_64(&var.data[start..end]));
5274                        }
5275                    }
5276                    _ => return Some(false),
5277                }
5278                let ratio = unique.len() as f64 / sample_count as f64;
5279                Some(sample_count >= 1024 && ratio >= unique_ratio_threshold)
5280            }
5281            _ => Some(false),
5282        }
5283    }
5284
5285    // Creates an encode task, consuming all buffered data
5286    fn do_flush(
5287        &mut self,
5288        arrays: Vec<ArrayRef>,
5289        repdefs: Vec<RepDefBuilder>,
5290        row_number: u64,
5291        num_rows: u64,
5292    ) -> Result<Vec<EncodeTask>> {
5293        let column_idx = self.column_index;
5294        let compression_strategy = self.compression_strategy.clone();
5295        let field = self.field.clone();
5296        let encoding_metadata = self.encoding_metadata.clone();
5297        let support_large_chunk = self.support_large_chunk;
5298        let version = self.version;
5299        let task = spawn_cpu(move || {
5300            let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
5301            let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
5302            let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
5303            let repdef = RepDefBuilder::serialize(repdefs);
5304
5305            if num_values == 0 {
5306                // We should not encode empty arrays.  So if we get here that should mean that we
5307                // either have all empty lists or all null lists (or a mix).  We still need to encode
5308                // the rep/def information but we can skip the data encoding.
5309                log::debug!("Encoding column {} with {} items ({} rows) using complex-null layout", column_idx, num_values, num_rows);
5310                return Self::encode_complex_all_null(
5311                    column_idx,
5312                    repdef,
5313                    row_number,
5314                    num_rows,
5315                    version,
5316                    compression_strategy.as_ref(),
5317                );
5318            }
5319
5320            let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
5321            let all_null = leaf_validity
5322                .as_ref()
5323                .map(|validity| validity.count_set_bits() == 0)
5324                .unwrap_or(false);
5325
5326            if all_null {
5327                return if is_simple_validity {
5328                    log::debug!(
5329                        "Encoding column {} with {} items ({} rows) using simple-null layout",
5330                        column_idx,
5331                        num_values,
5332                        num_rows
5333                    );
5334                    Self::encode_simple_all_null(column_idx, num_values, row_number)
5335                } else {
5336                    log::debug!(
5337                        "Encoding column {} with {} items ({} rows) using complex-null layout",
5338                        column_idx,
5339                        num_values,
5340                        num_rows
5341                    );
5342                    Self::encode_complex_all_null(
5343                        column_idx,
5344                        repdef,
5345                        row_number,
5346                        num_rows,
5347                        version,
5348                        compression_strategy.as_ref(),
5349                    )
5350                };
5351            }
5352
5353            if let DataType::Struct(fields) = &field.data_type()
5354                && fields.is_empty()
5355            {
5356                if has_repdef_info {
5357                    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()));
5358                }
5359                // This is maybe a little confusing but the reader should never look at this anyways and it
5360                // seems like overkill to invent a new layout just for "empty structs".
5361                return Self::encode_simple_all_null(column_idx, num_values, row_number);
5362            }
5363
5364            let data_block = DataBlock::from_arrays(&arrays, num_values);
5365
5366            if version.resolve() >= LanceFileVersion::V2_2
5367                && let Some(scalar) = Self::find_constant_scalar(&arrays, leaf_validity.as_ref())?
5368            {
5369                log::debug!(
5370                    "Encoding column {} with {} items ({} rows) using constant layout",
5371                    column_idx,
5372                    num_values,
5373                    num_rows
5374                );
5375                return constant::encode_constant_page(
5376                    column_idx,
5377                    scalar,
5378                    repdef,
5379                    row_number,
5380                    num_rows,
5381                );
5382            }
5383
5384            let requires_full_zip_packed_struct =
5385                if let DataBlock::Struct(ref struct_data_block) = data_block {
5386                    struct_data_block.has_variable_width_child()
5387                } else {
5388                    false
5389                };
5390
5391            if requires_full_zip_packed_struct {
5392                log::debug!(
5393                    "Encoding column {} with {} items using full-zip packed struct layout",
5394                    column_idx,
5395                    num_values
5396                );
5397                return Self::encode_full_zip(
5398                    column_idx,
5399                    &field,
5400                    compression_strategy.as_ref(),
5401                    data_block,
5402                    repdef,
5403                    row_number,
5404                    num_rows,
5405                );
5406            }
5407
5408            // If the rep/def levels are too sparse for miniblock (e.g. many empty
5409            // lists with very few values), fall back to fullzip to avoid exceeding
5410            // the u16 per-chunk rep/def buffer size limit.
5411            let too_sparse = Self::repdef_too_sparse_for_miniblock(&repdef, num_values);
5412
5413            if !too_sparse {
5414                if let DataBlock::Dictionary(dict) = data_block {
5415                    log::debug!("Encoding column {} with {} items using dictionary encoding (already dictionary encoded)", column_idx, num_values);
5416                    let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
5417                    // TODO: https://github.com/lancedb/lance/issues/4809
5418                    // If we compute stats on dictionary_data_block => panic.
5419                    // If we don't compute stats on indices_data_block => panic.
5420                    // This is messy.  Don't make me call compute_stat ever.
5421                    indices_data_block.compute_stat();
5422                    return Self::encode_miniblock(
5423                        column_idx,
5424                        &field,
5425                        compression_strategy.as_ref(),
5426                        indices_data_block,
5427                        repdef,
5428                        row_number,
5429                        Some(dictionary_data_block),
5430                        num_rows,
5431                        support_large_chunk,
5432                    );
5433                }
5434            } else {
5435                log::debug!(
5436                    "Encoding column {} with {} items using full-zip layout \
5437                     (rep/def too sparse for mini-block)",
5438                    column_idx,
5439                    num_values
5440                );
5441            }
5442
5443            {
5444                // Try dictionary encoding first if applicable. If encoding aborts, fall back to the
5445                // preferred structural encoding.
5446                let dict_result = if too_sparse {
5447                    None
5448                } else {
5449                    Self::should_dictionary_encode(&data_block, &field, version)
5450                        .and_then(|budget| {
5451                            log::debug!(
5452                                "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
5453                                column_idx,
5454                                num_values
5455                            );
5456                            dict::dictionary_encode(
5457                                &data_block,
5458                                budget.max_dict_entries,
5459                                budget.max_encoded_size,
5460                            )
5461                        })
5462                };
5463
5464                if let Some((indices_data_block, dictionary_data_block)) = dict_result {
5465                    Self::encode_miniblock(
5466                        column_idx,
5467                        &field,
5468                        compression_strategy.as_ref(),
5469                        indices_data_block,
5470                        repdef,
5471                        row_number,
5472                        Some(dictionary_data_block),
5473                        num_rows,
5474                        support_large_chunk,
5475                    )
5476                } else if !too_sparse && Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
5477                    log::debug!(
5478                        "Encoding column {} with {} items using mini-block layout",
5479                        column_idx,
5480                        num_values
5481                    );
5482                    Self::encode_miniblock(
5483                        column_idx,
5484                        &field,
5485                        compression_strategy.as_ref(),
5486                        data_block,
5487                        repdef,
5488                        row_number,
5489                        None,
5490                        num_rows,
5491                        support_large_chunk,
5492                    )
5493                } else if too_sparse || Self::prefers_fullzip(encoding_metadata.as_ref()) {
5494                    // If the user explicitly asked for miniblock but the
5495                    // heuristic vetoed it, warn so the override is observable
5496                    // rather than silent.
5497                    let user_forced_miniblock = encoding_metadata
5498                        .get(STRUCTURAL_ENCODING_META_KEY)
5499                        .map(|v| v.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK)
5500                        .unwrap_or(false);
5501                    if too_sparse && user_forced_miniblock {
5502                        log::warn!(
5503                            "Field '{}' requested {}={} but the rep/def shape would \
5504                             overflow miniblock per-chunk limits; encoding as fullzip \
5505                             instead to avoid corruption.",
5506                            field.name,
5507                            STRUCTURAL_ENCODING_META_KEY,
5508                            STRUCTURAL_ENCODING_MINIBLOCK,
5509                        );
5510                    }
5511                    log::debug!(
5512                        "Encoding column {} with {} items using full-zip layout",
5513                        column_idx,
5514                        num_values
5515                    );
5516                    Self::encode_full_zip(
5517                        column_idx,
5518                        &field,
5519                        compression_strategy.as_ref(),
5520                        data_block,
5521                        repdef,
5522                        row_number,
5523                        num_rows,
5524                    )
5525                } else {
5526                    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()))
5527                }
5528            }
5529        })
5530        .boxed();
5531        Ok(vec![task])
5532    }
5533
5534    fn extract_validity_buf(
5535        array: Arc<dyn Array>,
5536        repdef: &mut RepDefBuilder,
5537        keep_original_array: bool,
5538    ) -> Result<Arc<dyn Array>> {
5539        if let Some(validity) = array.nulls() {
5540            if keep_original_array {
5541                repdef.add_validity_bitmap(validity.clone());
5542            } else {
5543                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
5544            }
5545            let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
5546            Ok(make_array(data_no_nulls))
5547        } else {
5548            repdef.add_no_null(array.len());
5549            Ok(array)
5550        }
5551    }
5552
5553    fn extract_validity(
5554        mut array: Arc<dyn Array>,
5555        repdef: &mut RepDefBuilder,
5556        keep_original_array: bool,
5557    ) -> Result<Arc<dyn Array>> {
5558        match array.data_type() {
5559            DataType::Null => {
5560                repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
5561                Ok(array)
5562            }
5563            DataType::Dictionary(_, _) => {
5564                array = dict::normalize_dict_nulls(array)?;
5565                Self::extract_validity_buf(array, repdef, keep_original_array)
5566            }
5567            // Extract our validity buf but NOT any child validity bufs. (they will be encoded in
5568            // as part of the values).  Note: for FSL we do not use repdef.add_fsl because we do
5569            // NOT want to increase the repdef depth.
5570            //
5571            // This would be quite catasrophic for something like vector embeddings.  Imagine we
5572            // had thousands of vectors and some were null but no vector contained null items.  If
5573            // we treated the vectors (primitive FSL) like we treat structural FSL we would end up
5574            // with a rep/def value for every single item in the vector.
5575            _ => Self::extract_validity_buf(array, repdef, keep_original_array),
5576        }
5577    }
5578}
5579
5580impl FieldEncoder for PrimitiveStructuralEncoder {
5581    // Buffers data, if there is enough to write a page then we create an encode task
5582    fn maybe_encode(
5583        &mut self,
5584        array: ArrayRef,
5585        _external_buffers: &mut OutOfLineBuffers,
5586        mut repdef: RepDefBuilder,
5587        row_number: u64,
5588        num_rows: u64,
5589    ) -> Result<Vec<EncodeTask>> {
5590        let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
5591        self.accumulated_repdefs.push(repdef);
5592
5593        if let Some((arrays, row_number, num_rows)) =
5594            self.accumulation_queue.insert(array, row_number, num_rows)
5595        {
5596            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5597            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5598        } else {
5599            Ok(vec![])
5600        }
5601    }
5602
5603    // If there is any data left in the buffer then create an encode task from it
5604    fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
5605        if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
5606            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
5607            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
5608        } else {
5609            Ok(vec![])
5610        }
5611    }
5612
5613    fn num_columns(&self) -> u32 {
5614        1
5615    }
5616
5617    fn finish(
5618        &mut self,
5619        _external_buffers: &mut OutOfLineBuffers,
5620    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
5621        std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
5622    }
5623}
5624
5625#[cfg(test)]
5626#[allow(clippy::single_range_in_vec_init)]
5627mod tests {
5628    use super::{
5629        ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor,
5630        FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource,
5631        FullZipRepIndexDetails, FullZipScheduler, MiniBlockRepIndex, PerValueDecompressor,
5632        PreambleAction, StructuralPageScheduler, VariableFullZipDecoder,
5633    };
5634    use crate::buffer::LanceBuffer;
5635    use crate::compression::DefaultDecompressionStrategy;
5636    use crate::constants::{
5637        COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
5638        DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
5639        STRUCTURAL_ENCODING_MINIBLOCK,
5640    };
5641    use crate::data::BlockInfo;
5642    use crate::decoder::PageEncoding;
5643    use crate::encodings::logical::primitive::{
5644        ChunkDrainInstructions, PrimitiveStructuralEncoder,
5645    };
5646    use crate::format::ProtobufUtils21;
5647    use crate::format::pb21;
5648    use crate::format::pb21::compressive_encoding::Compression;
5649    use crate::testing::{TestCases, check_round_trip_encoding_of_data};
5650    use crate::version::LanceFileVersion;
5651    use arrow_array::{ArrayRef, Int8Array, StringArray};
5652    use arrow_schema::DataType;
5653    use std::collections::HashMap;
5654    use std::{collections::VecDeque, sync::Arc};
5655
5656    #[test]
5657    fn test_is_narrow() {
5658        let int8_array = Int8Array::from(vec![1, 2, 3]);
5659        let array_ref: ArrayRef = Arc::new(int8_array);
5660        let block = DataBlock::from_array(array_ref);
5661
5662        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5663
5664        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
5665        let block = DataBlock::from_array(string_array);
5666        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
5667
5668        let string_array = StringArray::from(vec![
5669            Some("hello world".repeat(100)),
5670            Some("world".to_string()),
5671        ]);
5672        let block = DataBlock::from_array(string_array);
5673        assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
5674    }
5675
5676    #[test]
5677    fn test_map_range() {
5678        // Null in the middle
5679        // [[A, B, C], [D, E], NULL, [F, G, H]]
5680        let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
5681        let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
5682        let max_visible_def = 0;
5683        let total_items = 8;
5684        let max_rep = 1;
5685
5686        let check = |range, expected_item_range, expected_level_range| {
5687            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5688                range,
5689                rep.as_ref(),
5690                def.as_ref(),
5691                max_rep,
5692                max_visible_def,
5693                total_items,
5694                PreambleAction::Absent,
5695            );
5696            assert_eq!(item_range, expected_item_range);
5697            assert_eq!(level_range, expected_level_range);
5698        };
5699
5700        check(0..1, 0..3, 0..3);
5701        check(1..2, 3..5, 3..5);
5702        check(2..3, 5..5, 5..6);
5703        check(3..4, 5..8, 6..9);
5704        check(0..2, 0..5, 0..5);
5705        check(1..3, 3..5, 3..6);
5706        check(2..4, 5..8, 5..9);
5707        check(0..3, 0..5, 0..6);
5708        check(1..4, 3..8, 3..9);
5709        check(0..4, 0..8, 0..9);
5710
5711        // Null at start
5712        // [NULL, [A, B], [C]]
5713        let rep = Some(vec![1, 1, 0, 1]);
5714        let def = Some(vec![1, 0, 0, 0]);
5715        let max_visible_def = 0;
5716        let total_items = 3;
5717
5718        let check = |range, expected_item_range, expected_level_range| {
5719            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5720                range,
5721                rep.as_ref(),
5722                def.as_ref(),
5723                max_rep,
5724                max_visible_def,
5725                total_items,
5726                PreambleAction::Absent,
5727            );
5728            assert_eq!(item_range, expected_item_range);
5729            assert_eq!(level_range, expected_level_range);
5730        };
5731
5732        check(0..1, 0..0, 0..1);
5733        check(1..2, 0..2, 1..3);
5734        check(2..3, 2..3, 3..4);
5735        check(0..2, 0..2, 0..3);
5736        check(1..3, 0..3, 1..4);
5737        check(0..3, 0..3, 0..4);
5738
5739        // Null at end
5740        // [[A], [B, C], NULL]
5741        let rep = Some(vec![1, 1, 0, 1]);
5742        let def = Some(vec![0, 0, 0, 1]);
5743        let max_visible_def = 0;
5744        let total_items = 3;
5745
5746        let check = |range, expected_item_range, expected_level_range| {
5747            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5748                range,
5749                rep.as_ref(),
5750                def.as_ref(),
5751                max_rep,
5752                max_visible_def,
5753                total_items,
5754                PreambleAction::Absent,
5755            );
5756            assert_eq!(item_range, expected_item_range);
5757            assert_eq!(level_range, expected_level_range);
5758        };
5759
5760        check(0..1, 0..1, 0..1);
5761        check(1..2, 1..3, 1..3);
5762        check(2..3, 3..3, 3..4);
5763        check(0..2, 0..3, 0..3);
5764        check(1..3, 1..3, 1..4);
5765        check(0..3, 0..3, 0..4);
5766
5767        // No nulls, with repetition
5768        // [[A, B], [C, D], [E, F]]
5769        let rep = Some(vec![1, 0, 1, 0, 1, 0]);
5770        let def: Option<&[u16]> = None;
5771        let max_visible_def = 0;
5772        let total_items = 6;
5773
5774        let check = |range, expected_item_range, expected_level_range| {
5775            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5776                range,
5777                rep.as_ref(),
5778                def.as_ref(),
5779                max_rep,
5780                max_visible_def,
5781                total_items,
5782                PreambleAction::Absent,
5783            );
5784            assert_eq!(item_range, expected_item_range);
5785            assert_eq!(level_range, expected_level_range);
5786        };
5787
5788        check(0..1, 0..2, 0..2);
5789        check(1..2, 2..4, 2..4);
5790        check(2..3, 4..6, 4..6);
5791        check(0..2, 0..4, 0..4);
5792        check(1..3, 2..6, 2..6);
5793        check(0..3, 0..6, 0..6);
5794
5795        // No repetition, with nulls (this case is trivial)
5796        // [A, B, NULL, C]
5797        let rep: Option<&[u16]> = None;
5798        let def = Some(vec![0, 0, 1, 0]);
5799        let max_visible_def = 1;
5800        let total_items = 4;
5801
5802        let check = |range, expected_item_range, expected_level_range| {
5803            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5804                range,
5805                rep.as_ref(),
5806                def.as_ref(),
5807                max_rep,
5808                max_visible_def,
5809                total_items,
5810                PreambleAction::Absent,
5811            );
5812            assert_eq!(item_range, expected_item_range);
5813            assert_eq!(level_range, expected_level_range);
5814        };
5815
5816        check(0..1, 0..1, 0..1);
5817        check(1..2, 1..2, 1..2);
5818        check(2..3, 2..3, 2..3);
5819        check(0..2, 0..2, 0..2);
5820        check(1..3, 1..3, 1..3);
5821        check(0..3, 0..3, 0..3);
5822
5823        // Tricky case, this chunk is a continuation and starts with a rep-index = 0
5824        // [[..., A] [B, C], NULL]
5825        //
5826        // What we do will depend on the preamble action
5827        let rep = Some(vec![0, 1, 0, 1]);
5828        let def = Some(vec![0, 0, 0, 1]);
5829        let max_visible_def = 0;
5830        let total_items = 3;
5831
5832        let check = |range, expected_item_range, expected_level_range| {
5833            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5834                range,
5835                rep.as_ref(),
5836                def.as_ref(),
5837                max_rep,
5838                max_visible_def,
5839                total_items,
5840                PreambleAction::Take,
5841            );
5842            assert_eq!(item_range, expected_item_range);
5843            assert_eq!(level_range, expected_level_range);
5844        };
5845
5846        // If we are taking the preamble then the range must start at 0
5847        check(0..1, 0..3, 0..3);
5848        check(0..2, 0..3, 0..4);
5849
5850        let check = |range, expected_item_range, expected_level_range| {
5851            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5852                range,
5853                rep.as_ref(),
5854                def.as_ref(),
5855                max_rep,
5856                max_visible_def,
5857                total_items,
5858                PreambleAction::Skip,
5859            );
5860            assert_eq!(item_range, expected_item_range);
5861            assert_eq!(level_range, expected_level_range);
5862        };
5863
5864        check(0..1, 1..3, 1..3);
5865        check(1..2, 3..3, 3..4);
5866        check(0..2, 1..3, 1..4);
5867
5868        // Another preamble case but now it doesn't end with a new list
5869        // [[..., A], NULL, [D, E]]
5870        //
5871        // What we do will depend on the preamble action
5872        let rep = Some(vec![0, 1, 1, 0]);
5873        let def = Some(vec![0, 1, 0, 0]);
5874        let max_visible_def = 0;
5875        let total_items = 4;
5876
5877        let check = |range, expected_item_range, expected_level_range| {
5878            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5879                range,
5880                rep.as_ref(),
5881                def.as_ref(),
5882                max_rep,
5883                max_visible_def,
5884                total_items,
5885                PreambleAction::Take,
5886            );
5887            assert_eq!(item_range, expected_item_range);
5888            assert_eq!(level_range, expected_level_range);
5889        };
5890
5891        // If we are taking the preamble then the range must start at 0
5892        check(0..1, 0..1, 0..2);
5893        check(0..2, 0..3, 0..4);
5894
5895        let check = |range, expected_item_range, expected_level_range| {
5896            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5897                range,
5898                rep.as_ref(),
5899                def.as_ref(),
5900                max_rep,
5901                max_visible_def,
5902                total_items,
5903                PreambleAction::Skip,
5904            );
5905            assert_eq!(item_range, expected_item_range);
5906            assert_eq!(level_range, expected_level_range);
5907        };
5908
5909        // If we are taking the preamble then the range must start at 0
5910        check(0..1, 1..1, 1..2);
5911        check(1..2, 1..3, 2..4);
5912        check(0..2, 1..3, 1..4);
5913
5914        // Now a preamble case without any definition levels
5915        // [[..., A] [B, C], [D]]
5916        let rep = Some(vec![0, 1, 0, 1]);
5917        let def: Option<Vec<u16>> = None;
5918        let max_visible_def = 0;
5919        let total_items = 4;
5920
5921        let check = |range, expected_item_range, expected_level_range| {
5922            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5923                range,
5924                rep.as_ref(),
5925                def.as_ref(),
5926                max_rep,
5927                max_visible_def,
5928                total_items,
5929                PreambleAction::Take,
5930            );
5931            assert_eq!(item_range, expected_item_range);
5932            assert_eq!(level_range, expected_level_range);
5933        };
5934
5935        // If we are taking the preamble then the range must start at 0
5936        check(0..1, 0..3, 0..3);
5937        check(0..2, 0..4, 0..4);
5938
5939        let check = |range, expected_item_range, expected_level_range| {
5940            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5941                range,
5942                rep.as_ref(),
5943                def.as_ref(),
5944                max_rep,
5945                max_visible_def,
5946                total_items,
5947                PreambleAction::Skip,
5948            );
5949            assert_eq!(item_range, expected_item_range);
5950            assert_eq!(level_range, expected_level_range);
5951        };
5952
5953        check(0..1, 1..3, 1..3);
5954        check(1..2, 3..4, 3..4);
5955        check(0..2, 1..4, 1..4);
5956
5957        // If we have nested lists then non-top level lists may be empty/null
5958        // and we need to make sure we still handle them as invisible items (we
5959        // failed to do this previously)
5960        let rep = Some(vec![2, 1, 2, 0, 1, 2]);
5961        let def = Some(vec![0, 1, 2, 0, 0, 0]);
5962        let max_rep = 2;
5963        let max_visible_def = 0;
5964        let total_items = 4;
5965
5966        let check = |range, expected_item_range, expected_level_range| {
5967            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5968                range,
5969                rep.as_ref(),
5970                def.as_ref(),
5971                max_rep,
5972                max_visible_def,
5973                total_items,
5974                PreambleAction::Absent,
5975            );
5976            assert_eq!(item_range, expected_item_range);
5977            assert_eq!(level_range, expected_level_range);
5978        };
5979
5980        check(0..3, 0..4, 0..6);
5981        check(0..1, 0..1, 0..2);
5982        check(1..2, 1..3, 2..5);
5983        check(2..3, 3..4, 5..6);
5984
5985        // Invisible items in a preamble that we are taking (regressing a previous failure)
5986        let rep = Some(vec![0, 0, 1, 0, 1, 1]);
5987        let def = Some(vec![0, 1, 0, 0, 0, 0]);
5988        let max_rep = 1;
5989        let max_visible_def = 0;
5990        let total_items = 5;
5991
5992        let check = |range, expected_item_range, expected_level_range| {
5993            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
5994                range,
5995                rep.as_ref(),
5996                def.as_ref(),
5997                max_rep,
5998                max_visible_def,
5999                total_items,
6000                PreambleAction::Take,
6001            );
6002            assert_eq!(item_range, expected_item_range);
6003            assert_eq!(level_range, expected_level_range);
6004        };
6005
6006        check(0..0, 0..1, 0..2);
6007        check(0..1, 0..3, 0..4);
6008        check(0..2, 0..4, 0..5);
6009
6010        // Skip preamble (with invis items) and skip a few rows (with invis items)
6011        // and then take a few rows but not all the rows
6012        let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
6013        let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
6014        let max_rep = 1;
6015        let max_visible_def = 0;
6016        let total_items = 5;
6017
6018        let check = |range, expected_item_range, expected_level_range| {
6019            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
6020                range,
6021                rep.as_ref(),
6022                def.as_ref(),
6023                max_rep,
6024                max_visible_def,
6025                total_items,
6026                PreambleAction::Skip,
6027            );
6028            assert_eq!(item_range, expected_item_range);
6029            assert_eq!(level_range, expected_level_range);
6030        };
6031
6032        check(2..3, 2..4, 5..7);
6033    }
6034
6035    #[test]
6036    fn test_slice_batch_data_and_rebase_offsets_u32() {
6037        let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
6038        let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
6039
6040        let (sliced_data, normalized_offsets) =
6041            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6042                .unwrap();
6043
6044        assert_eq!(sliced_data.as_ref(), b"6789ab");
6045        let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
6046        assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
6047    }
6048
6049    #[test]
6050    fn test_slice_batch_data_and_rebase_offsets_u64() {
6051        let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
6052        let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
6053
6054        let (sliced_data, normalized_offsets) =
6055            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
6056                .unwrap();
6057
6058        assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
6059        let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
6060        assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
6061    }
6062
6063    #[test]
6064    fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
6065        let data = LanceBuffer::copy_slice(b"abcd");
6066        let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
6067
6068        let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
6069            .expect_err("offset end before start should error");
6070        assert!(err.to_string().contains("less than base"));
6071    }
6072
6073    #[test]
6074    fn test_schedule_instructions() {
6075        // Convert repetition index to bytes for testing
6076        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6077        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6078        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6079
6080        let check = |user_ranges, expected_instructions| {
6081            let instructions =
6082                ChunkInstructions::schedule_instructions(&repetition_index, user_ranges);
6083            assert_eq!(instructions, expected_instructions);
6084        };
6085
6086        // The instructions we expect if we're grabbing the whole range
6087        let expected_take_all = vec![
6088            ChunkInstructions {
6089                chunk_idx: 0,
6090                preamble: PreambleAction::Absent,
6091                rows_to_skip: 0,
6092                rows_to_take: 6,
6093                take_trailer: true,
6094            },
6095            ChunkInstructions {
6096                chunk_idx: 1,
6097                preamble: PreambleAction::Take,
6098                rows_to_skip: 0,
6099                rows_to_take: 2,
6100                take_trailer: false,
6101            },
6102            ChunkInstructions {
6103                chunk_idx: 2,
6104                preamble: PreambleAction::Absent,
6105                rows_to_skip: 0,
6106                rows_to_take: 5,
6107                take_trailer: true,
6108            },
6109            ChunkInstructions {
6110                chunk_idx: 3,
6111                preamble: PreambleAction::Take,
6112                rows_to_skip: 0,
6113                rows_to_take: 1,
6114                take_trailer: false,
6115            },
6116        ];
6117
6118        // Take all as 1 range
6119        check(&[0..14], expected_take_all.clone());
6120
6121        // Take all a individual rows
6122        check(
6123            &[
6124                0..1,
6125                1..2,
6126                2..3,
6127                3..4,
6128                4..5,
6129                5..6,
6130                6..7,
6131                7..8,
6132                8..9,
6133                9..10,
6134                10..11,
6135                11..12,
6136                12..13,
6137                13..14,
6138            ],
6139            expected_take_all,
6140        );
6141
6142        // Test some partial takes
6143
6144        // 2 rows in the same chunk but not contiguous
6145        check(
6146            &[0..1, 3..4],
6147            vec![
6148                ChunkInstructions {
6149                    chunk_idx: 0,
6150                    preamble: PreambleAction::Absent,
6151                    rows_to_skip: 0,
6152                    rows_to_take: 1,
6153                    take_trailer: false,
6154                },
6155                ChunkInstructions {
6156                    chunk_idx: 0,
6157                    preamble: PreambleAction::Absent,
6158                    rows_to_skip: 3,
6159                    rows_to_take: 1,
6160                    take_trailer: false,
6161                },
6162            ],
6163        );
6164
6165        // Taking just a trailer/preamble
6166        check(
6167            &[5..6],
6168            vec![
6169                ChunkInstructions {
6170                    chunk_idx: 0,
6171                    preamble: PreambleAction::Absent,
6172                    rows_to_skip: 5,
6173                    rows_to_take: 1,
6174                    take_trailer: true,
6175                },
6176                ChunkInstructions {
6177                    chunk_idx: 1,
6178                    preamble: PreambleAction::Take,
6179                    rows_to_skip: 0,
6180                    rows_to_take: 0,
6181                    take_trailer: false,
6182                },
6183            ],
6184        );
6185
6186        // Skipping an entire chunk
6187        check(
6188            &[7..10],
6189            vec![
6190                ChunkInstructions {
6191                    chunk_idx: 1,
6192                    preamble: PreambleAction::Skip,
6193                    rows_to_skip: 1,
6194                    rows_to_take: 1,
6195                    take_trailer: false,
6196                },
6197                ChunkInstructions {
6198                    chunk_idx: 2,
6199                    preamble: PreambleAction::Absent,
6200                    rows_to_skip: 0,
6201                    rows_to_take: 2,
6202                    take_trailer: false,
6203                },
6204            ],
6205        );
6206    }
6207
6208    #[test]
6209    fn test_drain_instructions() {
6210        fn drain_from_instructions(
6211            instructions: &mut VecDeque<ChunkInstructions>,
6212            mut rows_desired: u64,
6213            need_preamble: &mut bool,
6214            skip_in_chunk: &mut u64,
6215        ) -> Vec<ChunkDrainInstructions> {
6216            // Note: instructions.len() is an upper bound, we typically take much fewer
6217            let mut drain_instructions = Vec::with_capacity(instructions.len());
6218            while rows_desired > 0 || *need_preamble {
6219                let (next_instructions, consumed_chunk) = instructions
6220                    .front()
6221                    .unwrap()
6222                    .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
6223                if consumed_chunk {
6224                    instructions.pop_front();
6225                }
6226                drain_instructions.push(next_instructions);
6227            }
6228            drain_instructions
6229        }
6230
6231        // Convert repetition index to bytes for testing
6232        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
6233        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6234        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6235        let user_ranges = vec![1..7, 10..14];
6236
6237        // First, schedule the ranges
6238        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6239
6240        let mut to_drain = VecDeque::from(scheduled.clone());
6241
6242        // Now we drain in batches of 4
6243
6244        let mut need_preamble = false;
6245        let mut skip_in_chunk = 0;
6246
6247        let next_batch =
6248            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6249
6250        assert!(!need_preamble);
6251        assert_eq!(skip_in_chunk, 4);
6252        assert_eq!(
6253            next_batch,
6254            vec![ChunkDrainInstructions {
6255                chunk_instructions: scheduled[0].clone(),
6256                rows_to_take: 4,
6257                rows_to_skip: 0,
6258                preamble_action: PreambleAction::Absent,
6259            }]
6260        );
6261
6262        let next_batch =
6263            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
6264
6265        assert!(!need_preamble);
6266        assert_eq!(skip_in_chunk, 2);
6267
6268        assert_eq!(
6269            next_batch,
6270            vec![
6271                ChunkDrainInstructions {
6272                    chunk_instructions: scheduled[0].clone(),
6273                    rows_to_take: 1,
6274                    rows_to_skip: 4,
6275                    preamble_action: PreambleAction::Absent,
6276                },
6277                ChunkDrainInstructions {
6278                    chunk_instructions: scheduled[1].clone(),
6279                    rows_to_take: 1,
6280                    rows_to_skip: 0,
6281                    preamble_action: PreambleAction::Take,
6282                },
6283                ChunkDrainInstructions {
6284                    chunk_instructions: scheduled[2].clone(),
6285                    rows_to_take: 2,
6286                    rows_to_skip: 0,
6287                    preamble_action: PreambleAction::Absent,
6288                }
6289            ]
6290        );
6291
6292        let next_batch =
6293            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6294
6295        assert!(!need_preamble);
6296        assert_eq!(skip_in_chunk, 0);
6297
6298        assert_eq!(
6299            next_batch,
6300            vec![
6301                ChunkDrainInstructions {
6302                    chunk_instructions: scheduled[2].clone(),
6303                    rows_to_take: 1,
6304                    rows_to_skip: 2,
6305                    preamble_action: PreambleAction::Absent,
6306                },
6307                ChunkDrainInstructions {
6308                    chunk_instructions: scheduled[3].clone(),
6309                    rows_to_take: 1,
6310                    rows_to_skip: 0,
6311                    preamble_action: PreambleAction::Take,
6312                },
6313            ]
6314        );
6315
6316        // Regression case.  Need a chunk with preamble, rows, and trailer (the middle chunk here)
6317        let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
6318        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
6319        let repetition_index = MiniBlockRepIndex::decode_from_bytes(&rep_bytes, 2);
6320        let user_ranges = vec![0..28];
6321
6322        // First, schedule the ranges
6323        let scheduled = ChunkInstructions::schedule_instructions(&repetition_index, &user_ranges);
6324
6325        let mut to_drain = VecDeque::from(scheduled.clone());
6326
6327        // Drain first chunk and some of second chunk
6328
6329        let mut need_preamble = false;
6330        let mut skip_in_chunk = 0;
6331
6332        let next_batch =
6333            drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
6334
6335        assert_eq!(
6336            next_batch,
6337            vec![
6338                ChunkDrainInstructions {
6339                    chunk_instructions: scheduled[0].clone(),
6340                    rows_to_take: 6,
6341                    rows_to_skip: 0,
6342                    preamble_action: PreambleAction::Absent,
6343                },
6344                ChunkDrainInstructions {
6345                    chunk_instructions: scheduled[1].clone(),
6346                    rows_to_take: 1,
6347                    rows_to_skip: 0,
6348                    preamble_action: PreambleAction::Take,
6349                },
6350            ]
6351        );
6352
6353        assert!(!need_preamble);
6354        assert_eq!(skip_in_chunk, 1);
6355
6356        // Now, the tricky part.  We drain the second chunk, including the trailer, and need to make sure
6357        // we get a drain task to take the preamble of the third chunk (and nothing else)
6358        let next_batch =
6359            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
6360
6361        assert_eq!(
6362            next_batch,
6363            vec![
6364                ChunkDrainInstructions {
6365                    chunk_instructions: scheduled[1].clone(),
6366                    rows_to_take: 2,
6367                    rows_to_skip: 1,
6368                    preamble_action: PreambleAction::Skip,
6369                },
6370                ChunkDrainInstructions {
6371                    chunk_instructions: scheduled[2].clone(),
6372                    rows_to_take: 0,
6373                    rows_to_skip: 0,
6374                    preamble_action: PreambleAction::Take,
6375                },
6376            ]
6377        );
6378
6379        assert!(!need_preamble);
6380        assert_eq!(skip_in_chunk, 0);
6381    }
6382
6383    #[tokio::test]
6384    async fn test_fullzip_initialize_is_lazy() {
6385        use futures::{FutureExt, future::BoxFuture};
6386        use std::ops::Range;
6387        use std::sync::Mutex;
6388
6389        #[derive(Debug, Clone)]
6390        struct RecordingScheduler {
6391            data: bytes::Bytes,
6392            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6393        }
6394
6395        impl RecordingScheduler {
6396            fn new(data: bytes::Bytes) -> Self {
6397                Self {
6398                    data,
6399                    requests: Arc::new(Mutex::new(Vec::new())),
6400                }
6401            }
6402
6403            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6404                self.requests.lock().unwrap().clone()
6405            }
6406        }
6407
6408        impl crate::EncodingsIo for RecordingScheduler {
6409            fn submit_request(
6410                &self,
6411                ranges: Vec<Range<u64>>,
6412                _priority: u64,
6413            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6414                self.requests.lock().unwrap().push(ranges.clone());
6415                let data = ranges
6416                    .into_iter()
6417                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6418                    .collect::<Vec<_>>();
6419                std::future::ready(Ok(data)).boxed()
6420            }
6421        }
6422
6423        #[derive(Debug)]
6424        struct TestFixedDecompressor;
6425
6426        impl FixedPerValueDecompressor for TestFixedDecompressor {
6427            fn decompress(
6428                &self,
6429                _data: FixedWidthDataBlock,
6430                _num_rows: u64,
6431            ) -> crate::Result<DataBlock> {
6432                unimplemented!("Test decompressor")
6433            }
6434
6435            fn bits_per_value(&self) -> u64 {
6436                32
6437            }
6438        }
6439
6440        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
6441            0;
6442            16 * 1024
6443        ])));
6444        let mut scheduler = FullZipScheduler {
6445            data_buf_position: 0,
6446            data_buf_size: 4096,
6447            rep_index: Some(FullZipRepIndexDetails {
6448                buf_position: 1000,
6449                bytes_per_value: 4,
6450            }),
6451            priority: 0,
6452            rows_in_page: 100,
6453            bits_per_offset: 32,
6454            details: Arc::new(FullZipDecodeDetails {
6455                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6456                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6457                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6458                max_rep: 0,
6459                max_visible_def: 0,
6460            }),
6461            cached_state: None,
6462            enable_cache: false,
6463        };
6464
6465        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6466        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6467
6468        assert!(
6469            cached_data
6470                .as_arc_any()
6471                .downcast_ref::<super::NoCachedPageData>()
6472                .is_some(),
6473            "FullZip initialize should not eagerly load repetition index data"
6474        );
6475        assert!(scheduler.cached_state.is_none());
6476        assert!(
6477            io.requests().is_empty(),
6478            "FullZip initialize should not issue any I/O"
6479        );
6480    }
6481
6482    #[tokio::test]
6483    async fn test_fullzip_read_source_slices_prefetched_page() {
6484        let page_start = 200_u64;
6485        let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
6486        let source = FullZipReadSource::PrefetchedPage {
6487            base_offset: page_start,
6488            data: page_data,
6489        };
6490        let ranges = vec![
6491            page_start..(page_start + 3),
6492            (page_start + 4)..(page_start + 8),
6493        ];
6494        let mut data = source.fetch(&ranges, 0).await.unwrap();
6495        assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
6496        assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
6497    }
6498
6499    #[tokio::test]
6500    async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
6501        use futures::{FutureExt, future::BoxFuture};
6502        use std::ops::Range;
6503        use std::sync::Mutex;
6504
6505        #[derive(Debug, Clone)]
6506        struct RecordingScheduler {
6507            data: bytes::Bytes,
6508            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6509        }
6510
6511        impl RecordingScheduler {
6512            fn new(data: bytes::Bytes) -> Self {
6513                Self {
6514                    data,
6515                    requests: Arc::new(Mutex::new(Vec::new())),
6516                }
6517            }
6518
6519            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6520                self.requests.lock().unwrap().clone()
6521            }
6522        }
6523
6524        impl crate::EncodingsIo for RecordingScheduler {
6525            fn submit_request(
6526                &self,
6527                ranges: Vec<Range<u64>>,
6528                _priority: u64,
6529            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6530                self.requests.lock().unwrap().push(ranges.clone());
6531                let data = ranges
6532                    .into_iter()
6533                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6534                    .collect::<Vec<_>>();
6535                std::future::ready(Ok(data)).boxed()
6536            }
6537        }
6538
6539        #[derive(Debug)]
6540        struct TestFixedDecompressor;
6541
6542        impl FixedPerValueDecompressor for TestFixedDecompressor {
6543            fn decompress(
6544                &self,
6545                _data: FixedWidthDataBlock,
6546                _num_rows: u64,
6547            ) -> crate::Result<DataBlock> {
6548                unimplemented!("Test decompressor")
6549            }
6550
6551            fn bits_per_value(&self) -> u64 {
6552                32
6553            }
6554        }
6555
6556        let rows_in_page = 100_u64;
6557        let bytes_per_value = 4_u64;
6558        let rep_start = 1000_u64;
6559        let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
6560        let mut data = vec![0_u8; 16 * 1024];
6561        data[rep_start as usize..rep_start as usize + rep_size].fill(7);
6562        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
6563
6564        let mut scheduler = FullZipScheduler {
6565            data_buf_position: 0,
6566            data_buf_size: 4096,
6567            rep_index: Some(FullZipRepIndexDetails {
6568                buf_position: rep_start,
6569                bytes_per_value,
6570            }),
6571            priority: 0,
6572            rows_in_page,
6573            bits_per_offset: 32,
6574            details: Arc::new(FullZipDecodeDetails {
6575                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6576                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6577                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6578                max_rep: 0,
6579                max_visible_def: 0,
6580            }),
6581            cached_state: None,
6582            enable_cache: true,
6583        };
6584
6585        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6586        let cached_data = scheduler.initialize(&io_dyn).await.unwrap();
6587        assert!(
6588            cached_data
6589                .as_arc_any()
6590                .downcast_ref::<FullZipCacheableState>()
6591                .is_some()
6592        );
6593        assert!(scheduler.cached_state.is_some());
6594        assert_eq!(
6595            io.requests(),
6596            vec![vec![
6597                rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
6598            ]]
6599        );
6600    }
6601
6602    #[tokio::test]
6603    async fn test_fullzip_full_page_bypasses_rep_index_io() {
6604        use futures::{FutureExt, future::BoxFuture};
6605        use std::ops::Range;
6606        use std::sync::Mutex;
6607
6608        #[derive(Debug, Clone)]
6609        struct RecordingScheduler {
6610            data: bytes::Bytes,
6611            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
6612        }
6613
6614        impl RecordingScheduler {
6615            fn new(data: bytes::Bytes) -> Self {
6616                Self {
6617                    data,
6618                    requests: Arc::new(Mutex::new(Vec::new())),
6619                }
6620            }
6621
6622            fn requests(&self) -> Vec<Vec<Range<u64>>> {
6623                self.requests.lock().unwrap().clone()
6624            }
6625        }
6626
6627        impl crate::EncodingsIo for RecordingScheduler {
6628            fn submit_request(
6629                &self,
6630                ranges: Vec<Range<u64>>,
6631                _priority: u64,
6632            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
6633                self.requests.lock().unwrap().push(ranges.clone());
6634                let data = ranges
6635                    .into_iter()
6636                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
6637                    .collect::<Vec<_>>();
6638                std::future::ready(Ok(data)).boxed()
6639            }
6640        }
6641
6642        #[derive(Debug)]
6643        struct TestFixedDecompressor;
6644
6645        impl FixedPerValueDecompressor for TestFixedDecompressor {
6646            fn decompress(
6647                &self,
6648                _data: FixedWidthDataBlock,
6649                _num_rows: u64,
6650            ) -> crate::Result<DataBlock> {
6651                unimplemented!("Test decompressor")
6652            }
6653
6654            fn bits_per_value(&self) -> u64 {
6655                32
6656            }
6657        }
6658
6659        let rows_in_page = 100_u64;
6660        let data_start = 256_u64;
6661        let data_size = 500_u64;
6662        let rep_start = 4096_u64;
6663        let bytes_per_value = 4_u64;
6664
6665        let mut bytes = vec![0_u8; 16 * 1024];
6666        for i in 0..=rows_in_page {
6667            let offset = (i * 5) as u32;
6668            let pos = rep_start as usize + (i * bytes_per_value) as usize;
6669            bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
6670        }
6671        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
6672
6673        let scheduler = FullZipScheduler {
6674            data_buf_position: data_start,
6675            data_buf_size: data_size,
6676            rep_index: Some(FullZipRepIndexDetails {
6677                buf_position: rep_start,
6678                bytes_per_value,
6679            }),
6680            priority: 0,
6681            rows_in_page,
6682            bits_per_offset: 32,
6683            details: Arc::new(FullZipDecodeDetails {
6684                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
6685                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
6686                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
6687                max_rep: 0,
6688                max_visible_def: 0,
6689            }),
6690            cached_state: None,
6691            enable_cache: false,
6692        };
6693
6694        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
6695        let tasks = scheduler
6696            .schedule_ranges_rep(
6697                &[0..rows_in_page],
6698                &io_dyn,
6699                FullZipRepIndexDetails {
6700                    buf_position: rep_start,
6701                    bytes_per_value,
6702                },
6703            )
6704            .unwrap();
6705
6706        let requests = io.requests();
6707        assert_eq!(requests.len(), 1);
6708        assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
6709
6710        let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
6711        let requests_after_await = io.requests();
6712        assert_eq!(
6713            requests_after_await.len(),
6714            1,
6715            "full page path should not issue rep-index I/O"
6716        );
6717    }
6718
6719    /// This test is used to reproduce fuzz test https://github.com/lancedb/lance/issues/4492
6720    #[tokio::test]
6721    async fn test_fuzz_issue_4492_empty_rep_values() {
6722        use lance_datagen::{RowCount, Seed, array, gen_batch};
6723
6724        let seed = 1823859942947654717u64;
6725        let num_rows = 2741usize;
6726
6727        // Generate the exact same data that caused the failure
6728        let batch_gen = gen_batch().with_seed(Seed::from(seed));
6729        let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
6730        let list_generator = array::rand_list_any(base_generator, false);
6731
6732        let batch = batch_gen
6733            .anon_col(list_generator)
6734            .into_batch_rows(RowCount::from(num_rows as u64))
6735            .unwrap();
6736
6737        let list_array = batch.column(0).clone();
6738
6739        // Force miniblock encoding
6740        let mut metadata = HashMap::new();
6741        metadata.insert(
6742            STRUCTURAL_ENCODING_META_KEY.to_string(),
6743            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6744        );
6745
6746        let test_cases = TestCases::default()
6747            .with_min_file_version(LanceFileVersion::V2_1)
6748            .with_batch_size(100)
6749            .with_range(0..num_rows.min(500) as u64)
6750            .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
6751
6752        check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
6753    }
6754
6755    async fn test_minichunk_size_helper(
6756        string_data: Vec<Option<String>>,
6757        minichunk_size: u64,
6758        file_version: LanceFileVersion,
6759    ) {
6760        use crate::constants::MINICHUNK_SIZE_META_KEY;
6761        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
6762        use arrow_array::{ArrayRef, StringArray};
6763        use std::sync::Arc;
6764
6765        let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
6766
6767        let mut metadata = HashMap::new();
6768        metadata.insert(
6769            MINICHUNK_SIZE_META_KEY.to_string(),
6770            minichunk_size.to_string(),
6771        );
6772        metadata.insert(
6773            STRUCTURAL_ENCODING_META_KEY.to_string(),
6774            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
6775        );
6776
6777        let test_cases = TestCases::default()
6778            .with_min_file_version(file_version)
6779            .with_batch_size(1000);
6780
6781        check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
6782    }
6783
6784    #[tokio::test]
6785    async fn test_minichunk_size_roundtrip() {
6786        // Test that minichunk size can be configured and works correctly in round-trip encoding
6787        let mut string_data = Vec::new();
6788        for i in 0..100 {
6789            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6790        }
6791        // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1
6792        test_minichunk_size_helper(string_data, 64, LanceFileVersion::V2_1).await;
6793    }
6794
6795    #[tokio::test]
6796    async fn test_minichunk_size_128kb_v2_2() {
6797        // Test that minichunk size can be configured to 128KB and works correctly with Lance 2.2
6798        let mut string_data = Vec::new();
6799        // create a 500kb string array
6800        for i in 0..10000 {
6801            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
6802        }
6803        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6804    }
6805
6806    #[tokio::test]
6807    async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
6808        let mut string_data = Vec::new();
6809        // 128kb/chunk / 6 bytes (t_9999) = 21845 > max 4096 items per chunk
6810        for i in 0..10000 {
6811            string_data.push(Some(format!("t_{}", i)));
6812        }
6813        test_minichunk_size_helper(string_data, 128 * 1024, LanceFileVersion::V2_2).await;
6814    }
6815
6816    #[tokio::test]
6817    async fn test_large_dictionary_general_compression() {
6818        use arrow_array::{ArrayRef, StringArray};
6819        use std::collections::HashMap;
6820        use std::sync::Arc;
6821
6822        // Create large string dictionary data (>32KiB) with low cardinality
6823        // Use 100 unique strings, each 500 bytes long = 50KB dictionary
6824        let unique_values: Vec<String> = (0..100)
6825            .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
6826            .collect();
6827
6828        // Repeat these strings many times to create a large array
6829        let repeated_strings: Vec<_> = unique_values
6830            .iter()
6831            .cycle()
6832            .take(100_000)
6833            .map(|s| Some(s.as_str()))
6834            .collect();
6835
6836        let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
6837
6838        // Configure test to use V2_2 and verify encoding
6839        let test_cases = TestCases::default()
6840            .with_min_file_version(LanceFileVersion::V2_2)
6841            .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
6842                assert_eq!(cols.len(), 1);
6843                let col = &cols[0];
6844
6845                // Navigate to the dictionary encoding in the page layout
6846                if let Some(PageEncoding::Structural(page_layout)) =
6847                    &col.final_pages.first().map(|p| &p.description)
6848                    && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
6849                        &page_layout.layout
6850                    && let Some(dictionary_encoding) = &mini_block.dictionary
6851                {
6852                    match dictionary_encoding.compression.as_ref() {
6853                        Some(Compression::General(general)) => {
6854                            // Verify it's using LZ4 or Zstd
6855                            let compression = general.compression.as_ref().unwrap();
6856                            assert!(
6857                                compression.scheme()
6858                                    == pb21::CompressionScheme::CompressionAlgorithmLz4
6859                                    || compression.scheme()
6860                                        == pb21::CompressionScheme::CompressionAlgorithmZstd,
6861                                "Expected LZ4 or Zstd compression for large dictionary"
6862                            );
6863                        }
6864                        _ => panic!("Expected General compression for large dictionary"),
6865                    }
6866                }
6867            }));
6868
6869        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
6870    }
6871
6872    fn dictionary_encoding_from_page(
6873        page: &crate::encoder::EncodedPage,
6874    ) -> &crate::format::pb21::CompressiveEncoding {
6875        let PageEncoding::Structural(layout) = &page.description else {
6876            panic!("Expected structural page encoding");
6877        };
6878        let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
6879        else {
6880            panic!("Expected mini-block layout");
6881        };
6882        layout
6883            .dictionary
6884            .as_ref()
6885            .unwrap_or_else(|| panic!("Expected dictionary encoding"))
6886    }
6887
6888    async fn encode_variable_dict_page(
6889        metadata: HashMap<String, String>,
6890    ) -> crate::encoder::EncodedPage {
6891        use arrow_array::types::Int32Type;
6892        use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
6893
6894        let values = Arc::new(StringArray::from(
6895            (0..128)
6896                .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
6897                .collect::<Vec<_>>(),
6898        )) as ArrayRef;
6899        let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
6900        let dict_array =
6901            Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
6902
6903        let field = arrow_schema::Field::new(
6904            "dict_col",
6905            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6906            false,
6907        )
6908        .with_metadata(metadata);
6909
6910        encode_first_page(field, dict_array, LanceFileVersion::V2_2).await
6911    }
6912
6913    async fn encode_auto_fixed_dict_page(
6914        metadata: HashMap<String, String>,
6915    ) -> crate::encoder::EncodedPage {
6916        use arrow_array::{ArrayRef, Decimal128Array};
6917
6918        // 128-bit fixed-width values with low cardinality to trigger dictionary encoding.
6919        let values = (0..20_000)
6920            .map(|i| match i % 3 {
6921                0 => 10_i128,
6922                1 => 20_i128,
6923                _ => 30_i128,
6924            })
6925            .collect::<Vec<_>>();
6926        let decimal = Decimal128Array::from_iter_values(values)
6927            .with_precision_and_scale(38, 0)
6928            .unwrap();
6929        let decimal = Arc::new(decimal) as ArrayRef;
6930
6931        let mut field_metadata = metadata;
6932        // Strongly encourage dictionary encoding for this synthetic test data.
6933        field_metadata.insert(
6934            "lance-encoding:dict-size-ratio".to_string(),
6935            "0.99".to_string(),
6936        );
6937        let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
6938            .with_metadata(field_metadata);
6939
6940        encode_first_page(field, decimal, LanceFileVersion::V2_2).await
6941    }
6942
6943    #[tokio::test]
6944    async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
6945        let page = encode_variable_dict_page(HashMap::new()).await;
6946        let dictionary_encoding = dictionary_encoding_from_page(&page);
6947        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6948            panic!("Expected General compression for dictionary values");
6949        };
6950        let compression = general.compression.as_ref().unwrap();
6951        assert_eq!(
6952            compression.scheme(),
6953            pb21::CompressionScheme::CompressionAlgorithmLz4
6954        );
6955    }
6956
6957    #[tokio::test]
6958    async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
6959        let page = encode_auto_fixed_dict_page(HashMap::new()).await;
6960        let dictionary_encoding = dictionary_encoding_from_page(&page);
6961        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6962            panic!("Expected General compression for dictionary values");
6963        };
6964        let compression = general.compression.as_ref().unwrap();
6965        assert_eq!(
6966            compression.scheme(),
6967            pb21::CompressionScheme::CompressionAlgorithmLz4
6968        );
6969    }
6970
6971    #[tokio::test]
6972    async fn test_dict_values_general_compression_zstd() {
6973        let mut metadata = HashMap::new();
6974        metadata.insert(
6975            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
6976            "zstd".to_string(),
6977        );
6978        let page = encode_variable_dict_page(metadata).await;
6979        let dictionary_encoding = dictionary_encoding_from_page(&page);
6980        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
6981            panic!("Expected General compression for dictionary values");
6982        };
6983        let compression = general.compression.as_ref().unwrap();
6984        assert_eq!(
6985            compression.scheme(),
6986            pb21::CompressionScheme::CompressionAlgorithmZstd
6987        );
6988    }
6989
6990    #[tokio::test]
6991    async fn test_dict_values_general_compression_none() {
6992        let mut metadata = HashMap::new();
6993        metadata.insert(
6994            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
6995            "none".to_string(),
6996        );
6997        let page = encode_variable_dict_page(metadata).await;
6998        let dictionary_encoding = dictionary_encoding_from_page(&page);
6999        assert!(
7000            !matches!(
7001                dictionary_encoding.compression.as_ref(),
7002                Some(Compression::General(_))
7003            ),
7004            "Expected dictionary values to avoid General compression"
7005        );
7006    }
7007
7008    #[test]
7009    fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
7010        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7011            &HashMap::new(),
7012            None,
7013            None,
7014        );
7015        assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
7016        assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
7017    }
7018
7019    #[test]
7020    fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
7021        let field_metadata = HashMap::from([
7022            (
7023                DICT_VALUES_COMPRESSION_META_KEY.to_string(),
7024                "none".to_string(),
7025            ),
7026            (
7027                DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
7028                "7".to_string(),
7029            ),
7030        ]);
7031        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7032            &field_metadata,
7033            Some("zstd".to_string()),
7034            Some("3".to_string()),
7035        );
7036        assert_eq!(
7037            metadata.get(COMPRESSION_META_KEY),
7038            Some(&"none".to_string()),
7039        );
7040        assert_eq!(
7041            metadata.get(COMPRESSION_LEVEL_META_KEY),
7042            Some(&"7".to_string()),
7043        );
7044    }
7045
7046    #[test]
7047    fn test_resolve_dict_values_compression_metadata_env_fallback() {
7048        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
7049            &HashMap::new(),
7050            Some("zstd".to_string()),
7051            Some("9".to_string()),
7052        );
7053        assert_eq!(
7054            metadata.get(COMPRESSION_META_KEY),
7055            Some(&"zstd".to_string()),
7056        );
7057        assert_eq!(
7058            metadata.get(COMPRESSION_LEVEL_META_KEY),
7059            Some(&"9".to_string()),
7060        );
7061    }
7062
7063    #[tokio::test]
7064    async fn test_dictionary_encode_int64() {
7065        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7066        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7067        use crate::version::LanceFileVersion;
7068        use arrow_array::{ArrayRef, Int64Array};
7069        use std::collections::HashMap;
7070        use std::sync::Arc;
7071
7072        // Low cardinality with poor RLE opportunity.
7073        let values = (0..1000)
7074            .map(|i| match i % 3 {
7075                0 => 10i64,
7076                1 => 20i64,
7077                _ => 30i64,
7078            })
7079            .collect::<Vec<_>>();
7080        let array = Arc::new(Int64Array::from(values)) as ArrayRef;
7081
7082        let mut metadata = HashMap::new();
7083        metadata.insert(
7084            STRUCTURAL_ENCODING_META_KEY.to_string(),
7085            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7086        );
7087        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7088
7089        let test_cases = TestCases::default()
7090            .with_min_file_version(LanceFileVersion::V2_2)
7091            .with_batch_size(1000)
7092            .with_range(0..1000)
7093            .with_indices(vec![0, 1, 10, 999])
7094            .with_expected_encoding("dictionary");
7095
7096        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7097    }
7098
7099    #[tokio::test]
7100    async fn test_dictionary_encode_float64() {
7101        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
7102        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
7103        use crate::version::LanceFileVersion;
7104        use arrow_array::{ArrayRef, Float64Array};
7105        use std::collections::HashMap;
7106        use std::sync::Arc;
7107
7108        // Low cardinality with poor RLE opportunity.
7109        let values = (0..1000)
7110            .map(|i| match i % 3 {
7111                0 => 0.1f64,
7112                1 => 0.2f64,
7113                _ => 0.3f64,
7114            })
7115            .collect::<Vec<_>>();
7116        let array = Arc::new(Float64Array::from(values)) as ArrayRef;
7117
7118        let mut metadata = HashMap::new();
7119        metadata.insert(
7120            STRUCTURAL_ENCODING_META_KEY.to_string(),
7121            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
7122        );
7123        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
7124
7125        let test_cases = TestCases::default()
7126            .with_min_file_version(LanceFileVersion::V2_2)
7127            .with_batch_size(1000)
7128            .with_range(0..1000)
7129            .with_indices(vec![0, 1, 10, 999])
7130            .with_expected_encoding("dictionary");
7131
7132        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
7133    }
7134
7135    #[test]
7136    fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
7137        let rows = 10_000;
7138        let unique_values = 2_000;
7139
7140        let dictionary_encoding =
7141            ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
7142        let layout = pb21::MiniBlockLayout {
7143            rep_compression: None,
7144            def_compression: None,
7145            value_compression: Some(ProtobufUtils21::flat(64, None)),
7146            dictionary: Some(dictionary_encoding),
7147            num_dictionary_items: unique_values,
7148            layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
7149            num_buffers: 1,
7150            repetition_index_depth: 0,
7151            num_items: rows,
7152            has_large_chunk: false,
7153        };
7154
7155        let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
7156        let scheduler = super::MiniBlockScheduler::try_new(
7157            &buffer_offsets_and_sizes,
7158            /*priority=*/ 0,
7159            /*items_in_page=*/ rows,
7160            &layout,
7161            &DefaultDecompressionStrategy::default(),
7162        )
7163        .unwrap();
7164
7165        let dictionary = scheduler.dictionary.unwrap();
7166        assert_eq!(dictionary.num_dictionary_items, unique_values);
7167        assert_eq!(
7168            dictionary.dictionary_data_alignment,
7169            crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
7170        );
7171    }
7172
7173    // Dictionary encoding decision tests
7174    fn create_test_fixed_data_block(
7175        num_values: u64,
7176        cardinality: u64,
7177        bits_per_value: u64,
7178    ) -> DataBlock {
7179        assert!(cardinality > 0);
7180        assert!(cardinality <= num_values);
7181        let block_info = BlockInfo::default();
7182
7183        assert_eq!(bits_per_value % 8, 0);
7184        let data = match bits_per_value {
7185            32 => {
7186                let values = (0..num_values)
7187                    .map(|i| (i % cardinality) as u32)
7188                    .collect::<Vec<_>>();
7189                crate::buffer::LanceBuffer::reinterpret_vec(values)
7190            }
7191            64 => {
7192                let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
7193                crate::buffer::LanceBuffer::reinterpret_vec(values)
7194            }
7195            128 => {
7196                let values = (0..num_values)
7197                    .map(|i| (i % cardinality) as u128)
7198                    .collect::<Vec<_>>();
7199                crate::buffer::LanceBuffer::reinterpret_vec(values)
7200            }
7201            _ => unreachable!(),
7202        };
7203        DataBlock::FixedWidth(FixedWidthDataBlock {
7204            bits_per_value,
7205            data,
7206            num_values,
7207            block_info,
7208        })
7209    }
7210
7211    /// Helper to create VariableWidth (string) test data block with exact cardinality
7212    fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
7213        use arrow_array::StringArray;
7214
7215        assert!(cardinality <= num_values && cardinality > 0);
7216
7217        let mut values = Vec::with_capacity(num_values as usize);
7218        for i in 0..num_values {
7219            values.push(format!("value_{:016}", i % cardinality));
7220        }
7221
7222        let array = StringArray::from(values);
7223        DataBlock::from_array(Arc::new(array) as ArrayRef)
7224    }
7225
7226    #[test]
7227    fn test_should_dictionary_encode() {
7228        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7229        use lance_core::datatypes::Field as LanceField;
7230
7231        // Create data where dict encoding saves space
7232        let block = create_test_variable_width_block(1000, 10);
7233
7234        let mut metadata = HashMap::new();
7235        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7236        let arrow_field =
7237            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7238        let field = LanceField::try_from(&arrow_field).unwrap();
7239
7240        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7241            &block,
7242            &field,
7243            LanceFileVersion::V2_1,
7244        );
7245
7246        assert!(
7247            result.is_some(),
7248            "Should use dictionary encode based on size"
7249        );
7250    }
7251
7252    #[test]
7253    fn test_should_not_dictionary_encode_unsupported_bits() {
7254        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7255        use lance_core::datatypes::Field as LanceField;
7256
7257        let block = create_test_fixed_data_block(1000, 1000, 32);
7258
7259        let mut metadata = HashMap::new();
7260        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
7261        let arrow_field =
7262            arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
7263        let field = LanceField::try_from(&arrow_field).unwrap();
7264
7265        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7266            &block,
7267            &field,
7268            LanceFileVersion::V2_1,
7269        );
7270
7271        assert!(
7272            result.is_none(),
7273            "Should not use dictionary encode for unsupported bit width"
7274        );
7275    }
7276
7277    #[test]
7278    fn test_should_not_dictionary_encode_near_unique_sample() {
7279        use crate::constants::DICT_SIZE_RATIO_META_KEY;
7280        use lance_core::datatypes::Field as LanceField;
7281
7282        let num_values = 5000;
7283        let block = create_test_variable_width_block(num_values, num_values);
7284
7285        let mut metadata = HashMap::new();
7286        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
7287        let arrow_field =
7288            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
7289        let field = LanceField::try_from(&arrow_field).unwrap();
7290
7291        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
7292            &block,
7293            &field,
7294            LanceFileVersion::V2_1,
7295        );
7296
7297        assert!(
7298            result.is_none(),
7299            "Should not probe dictionary encoding for near-unique data"
7300        );
7301    }
7302
7303    async fn encode_first_page(
7304        field: arrow_schema::Field,
7305        array: ArrayRef,
7306        version: LanceFileVersion,
7307    ) -> crate::encoder::EncodedPage {
7308        use crate::encoder::{
7309            ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
7310            default_encoding_strategy,
7311        };
7312        use crate::repdef::RepDefBuilder;
7313
7314        let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
7315        let encoding_strategy = default_encoding_strategy(version);
7316        let mut column_index_seq = ColumnIndexSequence::default();
7317        let encoding_options = EncodingOptions {
7318            cache_bytes_per_column: 1,
7319            max_page_bytes: 32 * 1024 * 1024,
7320            keep_original_array: true,
7321            buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
7322            version,
7323        };
7324
7325        let mut encoder = encoding_strategy
7326            .create_field_encoder(
7327                encoding_strategy.as_ref(),
7328                &lance_field,
7329                &mut column_index_seq,
7330                &encoding_options,
7331            )
7332            .unwrap();
7333
7334        let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
7335        let repdef = RepDefBuilder::default();
7336        let num_rows = array.len() as u64;
7337        let mut pages = Vec::new();
7338        for task in encoder
7339            .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
7340            .unwrap()
7341        {
7342            pages.push(task.await.unwrap());
7343        }
7344        for task in encoder.flush(&mut external_buffers).unwrap() {
7345            pages.push(task.await.unwrap());
7346        }
7347        pages.into_iter().next().unwrap()
7348    }
7349
7350    #[tokio::test]
7351    async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
7352        use crate::format::pb21::page_layout::Layout;
7353
7354        let val = vec![0xABu8; 33];
7355        let arr: ArrayRef = Arc::new(
7356            arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7357                std::iter::repeat_n(Some(val.as_slice()), 256),
7358                33,
7359            )
7360            .unwrap(),
7361        );
7362        let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
7363        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7364
7365        let PageEncoding::Structural(layout) = &page.description else {
7366            panic!("Expected structural encoding");
7367        };
7368        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7369            panic!("Expected constant layout in slot 2");
7370        };
7371        assert!(layout.inline_value.is_none());
7372        assert_eq!(page.data.len(), 1);
7373
7374        let test_cases = TestCases::default()
7375            .with_min_file_version(LanceFileVersion::V2_2)
7376            .with_max_file_version(LanceFileVersion::V2_2)
7377            .with_page_sizes(vec![4096]);
7378        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7379    }
7380
7381    #[tokio::test]
7382    async fn test_constant_layout_out_of_line_utf8_v2_2() {
7383        use crate::format::pb21::page_layout::Layout;
7384
7385        let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
7386            std::iter::repeat_n("hello", 512),
7387        ));
7388        let field = arrow_schema::Field::new("c", DataType::Utf8, true);
7389        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7390
7391        let PageEncoding::Structural(layout) = &page.description else {
7392            panic!("Expected structural encoding");
7393        };
7394        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7395            panic!("Expected constant layout in slot 2");
7396        };
7397        assert!(layout.inline_value.is_none());
7398        assert_eq!(page.data.len(), 1);
7399
7400        let test_cases = TestCases::default()
7401            .with_min_file_version(LanceFileVersion::V2_2)
7402            .with_max_file_version(LanceFileVersion::V2_2)
7403            .with_page_sizes(vec![4096]);
7404        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7405    }
7406
7407    #[tokio::test]
7408    async fn test_constant_layout_nullable_item_v2_2() {
7409        use crate::format::pb21::page_layout::Layout;
7410
7411        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
7412            Some(7),
7413            None,
7414            Some(7),
7415            None,
7416            Some(7),
7417        ]));
7418        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7419        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7420
7421        let PageEncoding::Structural(layout) = &page.description else {
7422            panic!("Expected structural encoding");
7423        };
7424        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7425            panic!("Expected constant layout in slot 2");
7426        };
7427        assert!(layout.inline_value.is_some());
7428        assert_eq!(page.data.len(), 2);
7429
7430        let test_cases = TestCases::default()
7431            .with_min_file_version(LanceFileVersion::V2_2)
7432            .with_max_file_version(LanceFileVersion::V2_2)
7433            .with_page_sizes(vec![4096]);
7434        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7435    }
7436
7437    #[tokio::test]
7438    async fn test_constant_layout_list_repdef_v2_2() {
7439        use crate::format::pb21::page_layout::Layout;
7440        use arrow_array::builder::{Int32Builder, ListBuilder};
7441
7442        let mut builder = ListBuilder::new(Int32Builder::new());
7443        builder.values().append_value(7);
7444        builder.values().append_null();
7445        builder.values().append_value(7);
7446        builder.append(true);
7447
7448        builder.append(true);
7449
7450        builder.values().append_value(7);
7451        builder.append(true);
7452
7453        builder.append_null();
7454
7455        let arr: ArrayRef = Arc::new(builder.finish());
7456        let field = arrow_schema::Field::new(
7457            "c",
7458            DataType::List(Arc::new(arrow_schema::Field::new(
7459                "item",
7460                DataType::Int32,
7461                true,
7462            ))),
7463            true,
7464        );
7465        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7466
7467        let PageEncoding::Structural(layout) = &page.description else {
7468            panic!("Expected structural encoding");
7469        };
7470        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7471            panic!("Expected constant layout in slot 2");
7472        };
7473        assert!(layout.inline_value.is_some());
7474        assert_eq!(page.data.len(), 2);
7475
7476        let test_cases = TestCases::default()
7477            .with_min_file_version(LanceFileVersion::V2_2)
7478            .with_max_file_version(LanceFileVersion::V2_2)
7479            .with_page_sizes(vec![4096]);
7480        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7481    }
7482
7483    #[tokio::test]
7484    async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
7485        use crate::format::pb21::page_layout::Layout;
7486        use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
7487
7488        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
7489        for _ in 0..64 {
7490            builder.values().append_value(1);
7491            builder.values().append_null();
7492            builder.values().append_value(3);
7493            builder.append(true);
7494        }
7495        let arr: ArrayRef = Arc::new(builder.finish());
7496        let field = arrow_schema::Field::new(
7497            "c",
7498            DataType::FixedSizeList(
7499                Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
7500                3,
7501            ),
7502            true,
7503        );
7504        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7505
7506        if let PageEncoding::Structural(layout) = &page.description {
7507            assert!(
7508                !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7509                "FixedSizeList should not use constant layout yet"
7510            );
7511        }
7512
7513        let test_cases = TestCases::default()
7514            .with_min_file_version(LanceFileVersion::V2_2)
7515            .with_max_file_version(LanceFileVersion::V2_2)
7516            .with_page_sizes(vec![4096]);
7517        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7518    }
7519
7520    #[tokio::test]
7521    async fn test_constant_layout_not_written_before_v2_2() {
7522        use crate::format::pb21::page_layout::Layout;
7523
7524        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
7525        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7526        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_1).await;
7527
7528        let PageEncoding::Structural(layout) = &page.description else {
7529            return;
7530        };
7531        assert!(
7532            !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
7533            "Should not emit constant layout before v2.2"
7534        );
7535
7536        let test_cases = TestCases::default()
7537            .with_min_file_version(LanceFileVersion::V2_1)
7538            .with_max_file_version(LanceFileVersion::V2_1)
7539            .with_page_sizes(vec![4096]);
7540        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7541    }
7542
7543    #[tokio::test]
7544    async fn test_all_null_constant_layout_still_works_v2_2() {
7545        use crate::format::pb21::page_layout::Layout;
7546
7547        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
7548        let field = arrow_schema::Field::new("c", DataType::Int32, true);
7549        let page = encode_first_page(field, arr.clone(), LanceFileVersion::V2_2).await;
7550
7551        let PageEncoding::Structural(layout) = &page.description else {
7552            panic!("Expected structural encoding");
7553        };
7554        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
7555            panic!("Expected layout in slot 2");
7556        };
7557        assert!(layout.inline_value.is_none());
7558        assert_eq!(page.data.len(), 0);
7559
7560        let test_cases = TestCases::default()
7561            .with_min_file_version(LanceFileVersion::V2_2)
7562            .with_max_file_version(LanceFileVersion::V2_2)
7563            .with_page_sizes(vec![4096]);
7564        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
7565    }
7566
7567    #[test]
7568    fn test_encode_decode_complex_all_null_vals_roundtrip() {
7569        use crate::compression::{
7570            DecompressionStrategy, DefaultCompressionStrategy, DefaultDecompressionStrategy,
7571        };
7572
7573        let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
7574
7575        let compression_strategy = DefaultCompressionStrategy::default();
7576        let decompression_strategy = DefaultDecompressionStrategy::default();
7577
7578        let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
7579            &values,
7580            &compression_strategy,
7581        )
7582        .unwrap();
7583
7584        let decompressor = decompression_strategy
7585            .create_block_decompressor(&encoding)
7586            .unwrap();
7587        let decompressed = decompressor
7588            .decompress(compressed_buf, values.len() as u64)
7589            .unwrap();
7590        let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
7591        assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
7592        assert_eq!(decompressed_fixed_width.bits_per_value, 16);
7593        let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
7594        assert_eq!(rep_result.as_ref(), values.as_ref());
7595    }
7596
7597    #[tokio::test]
7598    async fn test_complex_all_null_compression_gated_by_version() {
7599        use crate::format::pb21::page_layout::Layout;
7600        use arrow_array::ListArray;
7601
7602        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7603            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7604        );
7605        let arr: ArrayRef = Arc::new(list_array);
7606        let field = arrow_schema::Field::new(
7607            "c",
7608            DataType::List(Arc::new(arrow_schema::Field::new(
7609                "item",
7610                DataType::Int32,
7611                true,
7612            ))),
7613            true,
7614        );
7615
7616        let page_v21 = encode_first_page(field.clone(), arr.clone(), LanceFileVersion::V2_1).await;
7617        let PageEncoding::Structural(layout_v21) = &page_v21.description else {
7618            panic!("Expected structural encoding");
7619        };
7620        let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
7621            panic!("Expected constant layout");
7622        };
7623        assert!(layout_v21.rep_compression.is_none());
7624        assert!(layout_v21.def_compression.is_none());
7625        assert_eq!(layout_v21.num_rep_values, 0);
7626        assert_eq!(layout_v21.num_def_values, 0);
7627
7628        let page_v22 = encode_first_page(field, arr, LanceFileVersion::V2_2).await;
7629        let PageEncoding::Structural(layout_v22) = &page_v22.description else {
7630            panic!("Expected structural encoding");
7631        };
7632        let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
7633            panic!("Expected constant layout");
7634        };
7635        assert!(layout_v22.def_compression.is_some());
7636        assert!(layout_v22.num_def_values > 0);
7637    }
7638
7639    #[tokio::test]
7640    async fn test_complex_all_null_round_trip() {
7641        use arrow_array::ListArray;
7642
7643        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
7644            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
7645        );
7646
7647        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_2);
7648        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
7649            .await;
7650    }
7651
7652    // https://github.com/lance-format/lance/issues/6681
7653    #[tokio::test]
7654    async fn test_sparse_boolean_list_roundtrip() {
7655        use arrow_array::builder::{BooleanBuilder, ListBuilder};
7656
7657        let mut list_builder = ListBuilder::new(BooleanBuilder::new());
7658        for i in 0..1000i32 {
7659            if i % 64 == 0 {
7660                // Alternate true/false so the array is not constant (constant path avoids the bug).
7661                list_builder.values().append_value(i % 128 == 0);
7662                list_builder.append(true);
7663            } else {
7664                list_builder.append(false);
7665            }
7666        }
7667        let list_array = Arc::new(list_builder.finish());
7668
7669        let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);
7670        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
7671    }
7672
7673    fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
7674        use crate::compression::VariablePerValueDecompressor;
7675        use crate::encodings::physical::binary::VariableDecoder;
7676        use crate::repdef::{ControlWordParser, DefinitionInterpretation};
7677        use std::sync::Arc;
7678        Arc::new(super::FullZipDecodeDetails {
7679            value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
7680                VariableDecoder::default(),
7681            )
7682                as Arc<dyn VariablePerValueDecompressor>),
7683            def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
7684            ctrl_word_parser: ControlWordParser::new(0, 0),
7685            max_rep: 0,
7686            max_visible_def: 0,
7687        })
7688    }
7689
7690    fn decode_variable_full_zip(
7691        buf: Vec<u8>,
7692        bits_per_offset: u8,
7693    ) -> lance_core::Result<super::VariableFullZipDecoder> {
7694        use std::collections::VecDeque;
7695        let mut data = VecDeque::new();
7696        data.push_back(crate::buffer::LanceBuffer::from(buf));
7697        super::VariableFullZipDecoder::new(
7698            truncated_tail_details(),
7699            data,
7700            1,
7701            bits_per_offset,
7702            bits_per_offset,
7703        )
7704    }
7705
7706    /// A page whose item walk ends with a partial length prefix must surface a
7707    /// corrupt-file error rather than read past the end of the buffer.
7708    ///
7709    /// This asserts the error variant and message rather than merely expecting a
7710    /// panic: before the length prefix was bounds checked, the read was
7711    /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here
7712    /// (which a `#[should_panic]` test would have accepted as a pass) while a
7713    /// release build read up to 8 bytes out of a 4 byte allocation.
7714    #[test]
7715    fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
7716        use lance_core::Error;
7717
7718        for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
7719            let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
7720                .expect_err("a truncated length prefix must not decode");
7721            assert!(
7722                matches!(err, Error::CorruptFile { .. }),
7723                "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
7724                bits,
7725                buf_len,
7726                err
7727            );
7728            let msg = err.to_string();
7729            assert!(
7730                msg.contains("truncated length prefix"),
7731                "error should say what is wrong, got: {msg}"
7732            );
7733        }
7734    }
7735}