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        STRUCTURAL_ENCODING_SPARSE,
19    },
20    data::DictionaryDataBlock,
21    encodings::logical::primitive::blob::{BlobDescriptionPageScheduler, BlobPageScheduler},
22    format::{
23        ProtobufUtils21,
24        pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression},
25    },
26};
27use arrow_array::{
28    Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, new_null_array, types::UInt64Type,
29};
30use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer};
31use arrow_schema::{DataType, Field as ArrowField};
32use bytes::Bytes;
33use futures::{FutureExt, TryStreamExt, future::BoxFuture, stream::FuturesOrdered};
34use itertools::Itertools;
35use lance_arrow::DataTypeExt;
36use lance_arrow::deepcopy::deep_copy_nulls;
37use lance_core::{
38    cache::{CacheKey, CacheKeySchema, Context, DeepSizeOf, KeyBuilder, LanceCache},
39    error::{Error, LanceOptionExt},
40    utils::bit::pad_bytes,
41};
42use log::{debug, trace};
43
44use crate::encodings::logical::primitive::miniblock::MiniBlockChunk;
45use crate::encodings::physical::rle::{RleDecompressor, RleRuns};
46use crate::utils::bytepack::ByteUnpacker;
47use crate::{
48    compression::{
49        BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor,
50        compress_required_block, create_rle_decompressor,
51    },
52    data::{AllNullDataBlock, DataBlock, VariableWidthBlock},
53    decoder::{IO_REQUEST_BATCH_BYTES, RequestBatch},
54    utils::bytepack::BytepackedIntegerEncoder,
55};
56use crate::{
57    compression::{FixedPerValueDecompressor, VariablePerValueDecompressor},
58    encodings::logical::primitive::fullzip::PerValueDataBlock,
59};
60use crate::{
61    encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext},
62    statistics::{ComputeStat, GetStat, Stat},
63};
64use crate::{
65    repdef::{
66        CompositeRepDefUnraveler, ControlWordIterator, ControlWordParser, DefinitionInterpretation,
67        MiniBlockRepDefBudget, NormalizedStructuralPlan, RepDefSlicer, SerializedRepDefs,
68        build_control_word_iterator, max_visible_level,
69    },
70    utils::accumulation::AccumulationQueue,
71};
72use lance_core::{Result, datatypes::Field, utils::tokio::spawn_cpu};
73
74use crate::constants::{
75    COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY,
76    DICT_SIZE_RATIO_META_KEY, DICT_VALUES_COMPRESSION_ENV_VAR,
77    DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
78    DICT_VALUES_COMPRESSION_META_KEY,
79};
80use crate::{
81    EncodingsIo,
82    buffer::LanceBuffer,
83    data::{BlockInfo, DataBlockBuilder, FixedWidthDataBlock},
84    decoder::{
85        ColumnInfo, DecodePageTask, DecodedArray, DecodedPage, FilterExpression, LoadedPageShard,
86        MessageType, PageEncoding, PageInfo, ScheduledScanLine, SchedulerContext,
87        StructuralDecodeArrayTask, StructuralFieldDecoder, StructuralFieldScheduler,
88        StructuralPageDecoder, StructuralSchedulingJob, UnloadedPageShard,
89    },
90    encoder::{
91        EncodeTask, EncodedColumn, EncodedPage, EncodingOptions, FieldEncoder, OutOfLineBuffers,
92    },
93    repdef::{LevelBuffer, RepDefBuilder, RepDefUnraveler},
94};
95
96pub mod blob;
97mod chunk_index;
98pub mod constant;
99pub mod dict;
100pub mod fullzip;
101mod layout;
102pub mod miniblock;
103pub(crate) mod sparse;
104
105use chunk_index::{ItemCounts, MiniBlockChunkIndex, PrefixSums, RowMapping, parse_nested_rep};
106
107const FILL_BYTE: u8 = 0xFE;
108const DEFAULT_DICT_DIVISOR: u64 = 2;
109const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000;
110const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8;
111const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4";
112
113/// Largest level count a direct legacy u16-value/U8-run RLE frame can prove.
114///
115/// Mini-block level payload sizes are u16. After the eight-byte frame header, each run needs a
116/// two-byte value and a one-byte length, and each U8 run can represent at most 255 levels.
117const MAX_LEGACY_RLE_LEVELS: u64 = ((u16::MAX as u64 - std::mem::size_of::<u64>() as u64)
118    / (std::mem::size_of::<u16>() as u64 + std::mem::size_of::<u8>() as u64))
119    * u8::MAX as u64;
120
121struct PageLoadTask {
122    decoder_fut: BoxFuture<'static, Result<Box<dyn StructuralPageDecoder>>>,
123    num_rows: u64,
124}
125
126/// The ordered initialization-buffer slots for one page.
127///
128/// Optional slots preserve the semantic position of absent buffers so range
129/// declaration and buffer consumption share one representation.
130#[derive(Debug)]
131struct PageInitialization {
132    buffer_ranges: Vec<Option<Range<u64>>>,
133}
134
135impl PageInitialization {
136    fn new(buffer_ranges: impl IntoIterator<Item = Option<Range<u64>>>) -> Self {
137        Self {
138            buffer_ranges: buffer_ranges.into_iter().collect(),
139        }
140    }
141
142    fn required_ranges(&self) -> impl Iterator<Item = &Range<u64>> {
143        self.buffer_ranges.iter().flatten()
144    }
145
146    fn num_required_ranges(&self) -> usize {
147        self.required_ranges().count()
148    }
149
150    fn with_buffers(self, buffers: Vec<Bytes>) -> Result<PageInitializationBuffers> {
151        let expected = self.num_required_ranges();
152        if buffers.len() != expected {
153            return Err(Error::invalid_input_source(
154                format!(
155                    "Page initialization received {} buffers for {expected} declared ranges",
156                    buffers.len()
157                )
158                .into(),
159            ));
160        }
161
162        let mut buffers = buffers.into_iter();
163        let mut buffer_slots = Vec::with_capacity(self.buffer_ranges.len());
164        for range in self.buffer_ranges {
165            if range.is_some() {
166                let buffer = buffers.next().ok_or_else(|| {
167                    Error::internal("Page initialization buffer was not assigned")
168                })?;
169                buffer_slots.push(Some(buffer));
170            } else {
171                buffer_slots.push(None);
172            }
173        }
174        debug_assert!(
175            buffers.next().is_none(),
176            "validated page initialization buffer count must be exhausted"
177        );
178        Ok(PageInitializationBuffers { buffer_slots })
179    }
180}
181
182#[derive(Debug)]
183struct PageInitializationBuffers {
184    buffer_slots: Vec<Option<Bytes>>,
185}
186
187impl PageInitializationBuffers {
188    fn try_into_array<const N: usize>(self, layout: &str) -> Result<[Option<Bytes>; N]> {
189        let actual = self.buffer_slots.len();
190        self.buffer_slots.try_into().map_err(|_| {
191            Error::internal(format!(
192                "{layout} page initialization declared {actual} buffer slots, expected {N}"
193            ))
194        })
195    }
196
197    fn into_required(self, layout: &str) -> Result<Vec<Bytes>> {
198        self.buffer_slots
199            .into_iter()
200            .enumerate()
201            .map(|(index, buffer)| {
202                buffer.ok_or_else(|| {
203                    Error::internal(format!(
204                        "{layout} page initialization buffer slot {index} is absent"
205                    ))
206                })
207            })
208            .collect()
209    }
210}
211
212/// A trait for figuring out how to schedule the data within
213/// a single page.
214trait StructuralPageScheduler: std::fmt::Debug + Send {
215    /// Whether this page requires initialized metadata before scheduling.
216    fn needs_initialization(&self) -> bool {
217        true
218    }
219
220    /// Describes the file byte ranges this page must read to initialize its
221    /// metadata (chunk metadata, dictionary, repetition index, ...).
222    ///
223    /// The field scheduler sorts the combined I/O request by file offset and
224    /// restores buffers to the declared slots afterward.
225    fn init_layout(&self) -> Result<PageInitialization>;
226
227    /// Returns initialized state from the bytes declared by [`Self::init_layout`].
228    ///
229    /// The caller attaches the returned state with [`Self::try_load`].
230    fn init_from_buffers<'a>(
231        &'a mut self,
232        buffers: PageInitializationBuffers,
233        io: &Arc<dyn EncodingsIo>,
234    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>>;
235
236    /// Tries to attach metadata from a previous initialization.
237    fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()>;
238
239    // Identifies the requested decode shape (e.g. blob descriptor struct vs
240    // raw bytes). Blob columns can produce multiple page scheduler variants
241    // for the same physical column depending on the target field's data type,
242    // and the cached page state types differ per variant. The stable view is
243    // mixed into the cache key so different variants do not collide.
244    fn cache_view(&self) -> PageCacheView {
245        PageCacheView::Native
246    }
247
248    /// Schedules the read of the given ranges in the page
249    ///
250    /// The read may be split into multiple "shards" if the page is extremely large.
251    /// Each shard maps to one or more rows and can be decoded independently.
252    ///
253    /// Note: this sharding is for splitting up very large pages into smaller reads to
254    /// avoid buffering too much data in memory.  It is not related to the batch size or
255    /// compute units in any way.
256    fn schedule_ranges(
257        &self,
258        ranges: &[Range<u64>],
259        io: &Arc<dyn EncodingsIo>,
260    ) -> Result<Vec<PageLoadTask>>;
261}
262
263/// Metadata describing the decoded size of a mini-block
264#[derive(Debug)]
265struct ChunkMeta {
266    num_values: u64,
267    chunk_size_bytes: u64,
268    offset_bytes: u64,
269}
270
271/// A mini-block chunk that has been decoded and decompressed
272#[derive(Debug, Clone)]
273struct DecodedMiniBlockChunk {
274    rep: Option<ScalarBuffer<u16>>,
275    def: Option<ScalarBuffer<u16>>,
276    values: DataBlock,
277}
278
279/// A task to decode a one or more mini-blocks of data into an output batch
280///
281/// Note: Two batches might share the same mini-block of data.  When this happens
282/// then each batch gets a copy of the block and each batch decodes the block independently.
283///
284/// This means we have duplicated work but it is necessary to avoid having to synchronize
285/// the decoding of the block. (TODO: test this theory)
286#[derive(Debug)]
287struct DecodeMiniBlockTask {
288    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
289    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
290    value_decompressor: Arc<dyn MiniBlockDecompressor>,
291    dictionary_data: Option<Arc<DataBlock>>,
292    def_meaning: Arc<[DefinitionInterpretation]>,
293    num_buffers: u64,
294    max_visible_level: u16,
295    instructions: Vec<(ChunkDrainInstructions, LoadedChunk)>,
296    has_large_chunk: bool,
297}
298
299impl DecodeMiniBlockTask {
300    fn decoded_size_bytes(&self) -> Option<u64> {
301        if self.rep_decompressor.is_some() || self.def_decompressor.is_some() {
302            return None;
303        }
304        let num_values = self
305            .instructions
306            .iter()
307            .try_fold(0_u64, |total, (instruction, _)| {
308                total.checked_add(instruction.rows_to_take)
309            })?;
310        self.value_decompressor.decoded_size_bytes(num_values)
311    }
312
313    fn resolve_num_levels(
314        rep_decompressor: Option<&dyn BlockDecompressor>,
315        rep_levels: Option<&LanceBuffer>,
316        def_decompressor: Option<&dyn BlockDecompressor>,
317        def_levels: Option<&LanceBuffer>,
318        declared_num_levels: u16,
319    ) -> Result<u64> {
320        let rep_num_levels = match (rep_decompressor, rep_levels) {
321            (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?,
322            (None, None) => None,
323            _ => {
324                return Err(Error::invalid_input_source(
325                    "miniblock repetition codec and payload presence disagree".into(),
326                ));
327            }
328        };
329        let def_num_levels = match (def_decompressor, def_levels) {
330            (Some(decompressor), Some(levels)) => decompressor.infer_num_values(levels)?,
331            (None, None) => None,
332            _ => {
333                return Err(Error::invalid_input_source(
334                    "miniblock definition codec and payload presence disagree".into(),
335                ));
336            }
337        };
338
339        let inferred_num_levels = match (rep_num_levels, def_num_levels) {
340            (Some(rep_num_levels), Some(def_num_levels)) if rep_num_levels != def_num_levels => {
341                return Err(Error::invalid_input_source(
342                    format!(
343                        "miniblock structural streams disagree on the level count: repetition inferred {rep_num_levels}, definition inferred {def_num_levels}"
344                    )
345                    .into(),
346                ));
347            }
348            (Some(num_levels), _) | (_, Some(num_levels)) => num_levels,
349            (None, None) => return Ok(u64::from(declared_num_levels)),
350        };
351
352        let declared_num_levels = u64::from(declared_num_levels);
353        if inferred_num_levels == declared_num_levels {
354            return Ok(inferred_num_levels);
355        }
356        let u16_modulus = u64::from(u16::MAX) + 1;
357        if inferred_num_levels <= declared_num_levels
358            || inferred_num_levels % u16_modulus != declared_num_levels
359        {
360            return Err(Error::invalid_input_source(
361                format!(
362                    "miniblock payload proves {inferred_num_levels} levels but the header declared {declared_num_levels}; the counts are not congruent modulo 65536"
363                )
364                .into(),
365            ));
366        }
367        if inferred_num_levels > MAX_LEGACY_RLE_LEVELS {
368            return Err(Error::invalid_input_source(
369                format!(
370                    "miniblock payload proves {inferred_num_levels} levels, exceeding the legacy RLE payload bound of {MAX_LEGACY_RLE_LEVELS}"
371                )
372                .into(),
373            ));
374        }
375        Ok(inferred_num_levels)
376    }
377
378    fn decode_levels(
379        decompressor: &dyn BlockDecompressor,
380        levels: LanceBuffer,
381        expected_num_levels: u64,
382    ) -> Result<ScalarBuffer<u16>> {
383        let levels = decompressor.decompress(Some(levels), expected_num_levels)?;
384        let levels = levels.as_fixed_width().ok_or_else(|| {
385            Error::invalid_input_source(
386                "miniblock levels did not decode to fixed-width data".into(),
387            )
388        })?;
389        if levels.num_values != expected_num_levels {
390            return Err(Error::invalid_input_source(
391                format!(
392                    "miniblock levels decoded {} values, expected {expected_num_levels}",
393                    levels.num_values
394                )
395                .into(),
396            ));
397        }
398        if levels.bits_per_value != 16 {
399            return Err(Error::invalid_input_source(
400                format!(
401                    "miniblock levels decoded {} bits per value, expected 16",
402                    levels.bits_per_value
403                )
404                .into(),
405            ));
406        }
407        let expected_bytes = expected_num_levels.checked_mul(2).ok_or_else(|| {
408            Error::invalid_input_source(
409                format!("miniblock level byte count overflowed for {expected_num_levels} levels")
410                    .into(),
411            )
412        })?;
413        if levels.data.len() as u64 != expected_bytes {
414            return Err(Error::invalid_input_source(
415                format!(
416                    "miniblock levels decoded {} bytes, expected {expected_bytes}",
417                    levels.data.len()
418                )
419                .into(),
420            ));
421        }
422        Ok(levels.data.borrow_to_typed_slice::<u16>())
423    }
424
425    // We are building a LevelBuffer (levels) and want to copy into it `total_len`
426    // values from `level_buf` starting at `offset`.
427    //
428    // We need to handle both the case where `levels` is None (no nulls encountered
429    // yet) and the case where `level_buf` is None (the input we are copying from has
430    // no nulls)
431    fn extend_levels(
432        range: Range<u64>,
433        levels: &mut Option<LevelBuffer>,
434        level_buf: &Option<impl AsRef<[u16]>>,
435        dest_offset: usize,
436    ) {
437        if let Some(level_buf) = level_buf {
438            if levels.is_none() {
439                // This is the first non-empty def buf we've hit, fill in the past
440                // with 0 (valid)
441                let mut new_levels_vec =
442                    LevelBuffer::with_capacity(dest_offset + (range.end - range.start) as usize);
443                new_levels_vec.extend(iter::repeat_n(0, dest_offset));
444                *levels = Some(new_levels_vec);
445            }
446            levels.as_mut().unwrap().extend(
447                level_buf.as_ref()[range.start as usize..range.end as usize]
448                    .iter()
449                    .copied(),
450            );
451        } else if let Some(levels) = levels {
452            let num_values = (range.end - range.start) as usize;
453            // This is an all-valid level_buf but we had nulls earlier and so we
454            // need to materialize it
455            levels.extend(iter::repeat_n(0, num_values));
456        }
457    }
458
459    /// Maps a range of rows to a range of items and a range of levels
460    ///
461    /// If there is no repetition information this just returns the range as-is.
462    ///
463    /// If there is repetition information then we need to do some work to figure out what
464    /// range of items corresponds to the requested range of rows.
465    ///
466    /// For example, if the data is [[1, 2, 3], [4, 5], [6, 7]] and the range is 1..2 (i.e. just row
467    /// 1) then the user actually wants items 3..5.  In the above case the rep levels would be:
468    ///
469    /// Idx: 0 1 2 3 4 5 6
470    /// Rep: 1 0 0 1 0 1 0
471    ///
472    /// So the start (1) maps to the second 1 (idx=3) and the end (2) maps to the third 1 (idx=5)
473    ///
474    /// If there are invisible items then we don't count them when calculating the range of items we
475    /// are interested in but we do count them when calculating the range of levels we are interested
476    /// in.  As a result we have to return both the item range (first return value) and the level range
477    /// (second return value).
478    ///
479    /// For example, if the data is [[1, 2, 3], [4, 5], NULL, [6, 7, 8]] and the range is 2..4 then the
480    /// user wants items 5..8 but they want levels 5..9.  In the above case the rep/def levels would be:
481    ///
482    /// Idx: 0 1 2 3 4 5 6 7 8
483    /// Rep: 1 0 0 1 0 1 1 0 0
484    /// Def: 0 0 0 0 0 1 0 0 0
485    /// Itm: 1 2 3 4 5 6 7 8
486    ///
487    /// Finally, we have to contend with the fact that chunks may or may not start with a "preamble" of
488    /// trailing values that finish up a list from the previous chunk.  In this case the first item does
489    /// not start at max_rep because it is a continuation of the previous chunk.  For our purposes we do
490    /// not consider this a "row" and so the range 0..1 will refer to the first row AFTER the preamble.
491    ///
492    /// We have a separate parameter (`preamble_action`) to control whether we want the preamble or not.
493    ///
494    /// Note that the "trailer" is considered a "row" and if we want it we should include it in the range.
495    fn map_range(
496        range: Range<u64>,
497        rep: Option<&impl AsRef<[u16]>>,
498        def: Option<&impl AsRef<[u16]>>,
499        max_rep: u16,
500        max_visible_def: u16,
501        // The total number of items (not rows) in the chunk.  This is not quite the same as
502        // rep.len() / def.len() because it doesn't count invisible items
503        total_items: u64,
504        preamble_action: PreambleAction,
505    ) -> (Range<u64>, Range<u64>) {
506        if let Some(rep) = rep {
507            let mut rep = rep.as_ref();
508            // If there is a preamble and we need to skip it then do that first.  The work is the same
509            // whether there is def information or not
510            let mut items_in_preamble = 0_u64;
511            let first_row_start = match preamble_action {
512                PreambleAction::Skip | PreambleAction::Take => {
513                    let first_row_start = if let Some(def) = def.as_ref() {
514                        let mut first_row_start = None;
515                        for (idx, (rep, def)) in rep.iter().zip(def.as_ref()).enumerate() {
516                            if *rep == max_rep {
517                                first_row_start = Some(idx as u64);
518                                break;
519                            }
520                            if *def <= max_visible_def {
521                                items_in_preamble += 1;
522                            }
523                        }
524                        first_row_start
525                    } else {
526                        let first_row_start =
527                            rep.iter().position(|&r| r == max_rep).map(|r| r as u64);
528                        items_in_preamble = first_row_start.unwrap_or(rep.len() as u64);
529                        first_row_start
530                    };
531                    // It is possible for a chunk to be entirely partial values but if it is then it
532                    // should never show up as a preamble to skip
533                    if first_row_start.is_none() {
534                        assert!(preamble_action == PreambleAction::Take);
535                        return (0..total_items, 0..rep.len() as u64);
536                    }
537                    let first_row_start = first_row_start.unwrap();
538                    rep = &rep[first_row_start as usize..];
539                    first_row_start
540                }
541                PreambleAction::Absent => {
542                    debug_assert!(rep[0] == max_rep);
543                    0
544                }
545            };
546
547            // We hit this case when all we needed was the preamble
548            if range.start == range.end {
549                debug_assert!(preamble_action == PreambleAction::Take);
550                debug_assert!(items_in_preamble <= total_items);
551                return (0..items_in_preamble, 0..first_row_start);
552            }
553            assert!(range.start < range.end);
554
555            let mut rows_seen = 0;
556            let mut new_start = 0;
557            let mut new_levels_start = 0;
558
559            if let Some(def) = def {
560                let def = &def.as_ref()[first_row_start as usize..];
561
562                // range.start == 0 always maps to 0 (even with invis items), otherwise we need to walk
563                let mut lead_invis_seen = 0;
564
565                if range.start > 0 {
566                    if def[0] > max_visible_def {
567                        lead_invis_seen += 1;
568                    }
569                    for (idx, (rep, def)) in rep.iter().zip(def).skip(1).enumerate() {
570                        if *rep == max_rep {
571                            rows_seen += 1;
572                            if rows_seen == range.start {
573                                new_start = idx as u64 + 1 - lead_invis_seen;
574                                new_levels_start = idx as u64 + 1;
575                                break;
576                            }
577                        }
578                        if *def > max_visible_def {
579                            lead_invis_seen += 1;
580                        }
581                    }
582                }
583
584                rows_seen += 1;
585
586                let mut new_end = u64::MAX;
587                let mut new_levels_end = rep.len() as u64;
588                let new_start_is_visible = def[new_levels_start as usize] <= max_visible_def;
589                let mut tail_invis_seen = if new_start_is_visible { 0 } else { 1 };
590                for (idx, (rep, def)) in rep[(new_levels_start + 1) as usize..]
591                    .iter()
592                    .zip(&def[(new_levels_start + 1) as usize..])
593                    .enumerate()
594                {
595                    if *rep == max_rep {
596                        rows_seen += 1;
597                        if rows_seen == range.end + 1 {
598                            new_end = idx as u64 + new_start + 1 - tail_invis_seen;
599                            new_levels_end = idx as u64 + new_levels_start + 1;
600                            break;
601                        }
602                    }
603                    if *def > max_visible_def {
604                        tail_invis_seen += 1;
605                    }
606                }
607
608                if new_end == u64::MAX {
609                    new_levels_end = rep.len() as u64;
610                    let total_invis_seen = lead_invis_seen + tail_invis_seen;
611                    new_end = rep.len() as u64 - total_invis_seen;
612                }
613
614                assert_ne!(new_end, u64::MAX);
615
616                // Adjust for any skipped preamble
617                if preamble_action == PreambleAction::Skip {
618                    new_start += items_in_preamble;
619                    new_end += items_in_preamble;
620                    new_levels_start += first_row_start;
621                    new_levels_end += first_row_start;
622                } else if preamble_action == PreambleAction::Take {
623                    debug_assert_eq!(new_start, 0);
624                    debug_assert_eq!(new_levels_start, 0);
625                    new_end += items_in_preamble;
626                    new_levels_end += first_row_start;
627                }
628
629                debug_assert!(new_end <= total_items);
630                (new_start..new_end, new_levels_start..new_levels_end)
631            } else {
632                // Easy case, there are no invisible items, so we don't need to check for them
633                // The items range and levels range will be the same.  We do still need to walk
634                // the rep levels to find the row boundaries
635
636                // range.start == 0 always maps to 0, otherwise we need to walk
637                if range.start > 0 {
638                    for (idx, rep) in rep.iter().skip(1).enumerate() {
639                        if *rep == max_rep {
640                            rows_seen += 1;
641                            if rows_seen == range.start {
642                                new_start = idx as u64 + 1;
643                                break;
644                            }
645                        }
646                    }
647                }
648                let mut new_end = rep.len() as u64;
649                // range.end == max_items always maps to rep.len(), otherwise we need to walk
650                if range.end < total_items {
651                    for (idx, rep) in rep[(new_start + 1) as usize..].iter().enumerate() {
652                        if *rep == max_rep {
653                            rows_seen += 1;
654                            if rows_seen == range.end {
655                                new_end = idx as u64 + new_start + 1;
656                                break;
657                            }
658                        }
659                    }
660                }
661
662                // Adjust for any skipped preamble
663                if preamble_action == PreambleAction::Skip {
664                    new_start += first_row_start;
665                    new_end += first_row_start;
666                } else if preamble_action == PreambleAction::Take {
667                    debug_assert_eq!(new_start, 0);
668                    new_end += first_row_start;
669                }
670
671                debug_assert!(new_end <= total_items);
672                (new_start..new_end, new_start..new_end)
673            }
674        } else {
675            // No repetition info, easy case, just use the range as-is and the item
676            // and level ranges are the same
677            (range.clone(), range)
678        }
679    }
680
681    // read `num_buffers` buffer sizes from `buf` starting at `offset`
682    fn read_buffer_sizes<const LARGE: bool>(
683        buf: &[u8],
684        offset: &mut usize,
685        num_buffers: u64,
686    ) -> Vec<u32> {
687        let read_size = if LARGE { 4 } else { 2 };
688        (0..num_buffers)
689            .map(|_| {
690                let bytes = &buf[*offset..*offset + read_size];
691                let size = if LARGE {
692                    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
693                } else {
694                    // the buffer size is read from u16 but is stored as u32 after decoding for consistency
695                    u16::from_le_bytes([bytes[0], bytes[1]]) as u32
696                };
697                *offset += read_size;
698                size
699            })
700            .collect()
701    }
702
703    // Unserialize a miniblock into a collection of vectors
704    fn decode_miniblock_chunk(
705        &self,
706        buf: &LanceBuffer,
707        items_in_chunk: u64,
708    ) -> Result<DecodedMiniBlockChunk> {
709        let mut offset = 0;
710        let num_levels = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
711        offset += 2;
712
713        let rep_size = if self.rep_decompressor.is_some() {
714            let rep_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
715            offset += 2;
716            Some(rep_size)
717        } else {
718            None
719        };
720        let def_size = if self.def_decompressor.is_some() {
721            let def_size = u16::from_le_bytes([buf[offset], buf[offset + 1]]);
722            offset += 2;
723            Some(def_size)
724        } else {
725            None
726        };
727
728        let buffer_sizes = if self.has_large_chunk {
729            Self::read_buffer_sizes::<true>(buf, &mut offset, self.num_buffers)
730        } else {
731            Self::read_buffer_sizes::<false>(buf, &mut offset, self.num_buffers)
732        };
733
734        offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
735
736        let rep = rep_size.map(|rep_size| {
737            let rep = buf.slice_with_length(offset, rep_size as usize);
738            offset += rep_size as usize;
739            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
740            rep
741        });
742
743        let def = def_size.map(|def_size| {
744            let def = buf.slice_with_length(offset, def_size as usize);
745            offset += def_size as usize;
746            offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
747            def
748        });
749
750        let num_levels = Self::resolve_num_levels(
751            self.rep_decompressor.as_deref(),
752            rep.as_ref(),
753            self.def_decompressor.as_deref(),
754            def.as_ref(),
755            num_levels,
756        )?;
757
758        let buffers = buffer_sizes
759            .into_iter()
760            .map(|buf_size| {
761                let buf = buf.slice_with_length(offset, buf_size as usize);
762                offset += buf_size as usize;
763                offset += pad_bytes::<MINIBLOCK_ALIGNMENT>(offset);
764                buf
765            })
766            .collect::<Vec<_>>();
767
768        let values = self
769            .value_decompressor
770            .decompress(buffers, items_in_chunk)?;
771
772        let rep = rep
773            .map(|rep| {
774                Self::decode_levels(
775                    self.rep_decompressor.as_ref().unwrap().as_ref(),
776                    rep,
777                    num_levels,
778                )
779            })
780            .transpose()?;
781        let def = def
782            .map(|def| {
783                Self::decode_levels(
784                    self.def_decompressor.as_ref().unwrap().as_ref(),
785                    def,
786                    num_levels,
787                )
788            })
789            .transpose()?;
790
791        if let (Some(rep), Some(def)) = (&rep, &def)
792            && rep.len() != def.len()
793        {
794            return Err(Error::invalid_input_source(
795                format!(
796                    "miniblock structural streams decoded different level counts: repetition {}, definition {}",
797                    rep.len(),
798                    def.len()
799                )
800                .into(),
801            ));
802        }
803
804        Ok(DecodedMiniBlockChunk { rep, def, values })
805    }
806}
807
808impl DecodePageTask for DecodeMiniBlockTask {
809    fn decode(self: Box<Self>) -> Result<DecodedPage> {
810        // First, we create output buffers for the rep and def and data
811        let mut repbuf: Option<LevelBuffer> = None;
812        let mut defbuf: Option<LevelBuffer> = None;
813
814        let max_rep = self.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
815
816        let estimated_size_bytes = self.decoded_size_bytes().unwrap_or_else(|| {
817            // Variable-width and rep/def encoded output sizes are not known before decoding.
818            self.instructions
819                .iter()
820                .map(|(_, chunk)| chunk.data.len() as u64)
821                .sum::<u64>()
822                * 2
823        });
824        let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes);
825
826        // We need to keep track of the offset into repbuf/defbuf that we are building up
827        let mut level_offset = 0;
828
829        // Pre-compute caching needs for each chunk by checking if the next chunk is the same
830        let needs_caching: Vec<bool> = self
831            .instructions
832            .windows(2)
833            .map(|w| w[0].1.chunk_idx == w[1].1.chunk_idx)
834            .chain(std::iter::once(false)) // the last one never needs caching
835            .collect();
836
837        // Cache for storing decoded chunks when beneficial
838        let mut chunk_cache: Option<(usize, DecodedMiniBlockChunk)> = None;
839
840        // Now we iterate through each instruction and process it
841        for (idx, (instructions, chunk)) in self.instructions.iter().enumerate() {
842            let should_cache_this_chunk = needs_caching[idx];
843
844            let decoded_chunk = match &chunk_cache {
845                Some((cached_chunk_idx, cached_chunk)) if *cached_chunk_idx == chunk.chunk_idx => {
846                    // Clone only when we have a cache hit (much cheaper than decoding)
847                    cached_chunk.clone()
848                }
849                _ => {
850                    // Cache miss, need to decode
851                    let decoded = self.decode_miniblock_chunk(&chunk.data, chunk.items_in_chunk)?;
852
853                    // Only update cache if this chunk will benefit the next access
854                    if should_cache_this_chunk {
855                        chunk_cache = Some((chunk.chunk_idx, decoded.clone()));
856                    }
857                    decoded
858                }
859            };
860
861            let DecodedMiniBlockChunk { rep, def, values } = decoded_chunk;
862
863            // Our instructions tell us which rows we want to take from this chunk
864            let row_range_start =
865                instructions.rows_to_skip + instructions.chunk_instructions.rows_to_skip;
866            let row_range_end = row_range_start + instructions.rows_to_take;
867
868            // We use the rep info to map the row range to an item range / levels range
869            let (item_range, level_range) = Self::map_range(
870                row_range_start..row_range_end,
871                rep.as_ref(),
872                def.as_ref(),
873                max_rep,
874                self.max_visible_level,
875                chunk.items_in_chunk,
876                instructions.preamble_action,
877            );
878            if item_range.end - item_range.start > chunk.items_in_chunk {
879                return Err(lance_core::Error::internal(format!(
880                    "Item range {:?} is greater than chunk items in chunk {:?}",
881                    item_range, chunk.items_in_chunk
882                )));
883            }
884
885            // Now we append the data to the output buffers
886            Self::extend_levels(level_range.clone(), &mut repbuf, &rep, level_offset);
887            Self::extend_levels(level_range.clone(), &mut defbuf, &def, level_offset);
888            level_offset += (level_range.end - level_range.start) as usize;
889            data_builder.append(&values, item_range)?;
890        }
891
892        let mut data = data_builder.finish();
893
894        let unraveler =
895            RepDefUnraveler::new(repbuf, defbuf, self.def_meaning.clone(), data.num_values());
896
897        if let Some(dictionary) = &self.dictionary_data {
898            // Don't decode here, that happens later (if needed)
899            let DataBlock::FixedWidth(indices) = data else {
900                return Err(lance_core::Error::internal(format!(
901                    "Expected FixedWidth DataBlock for dictionary indices, got {:?}",
902                    data
903                )));
904            };
905            data = DataBlock::Dictionary(DictionaryDataBlock::from_parts(
906                indices,
907                dictionary.as_ref().clone(),
908            ));
909        }
910
911        Ok(DecodedPage {
912            data,
913            repdef: unraveler,
914        })
915    }
916}
917
918/// A chunk that has been loaded by the miniblock scheduler (but not
919/// yet decoded)
920#[derive(Debug)]
921struct LoadedChunk {
922    data: LanceBuffer,
923    items_in_chunk: u64,
924    byte_range: Range<u64>,
925    chunk_idx: usize,
926}
927
928impl Clone for LoadedChunk {
929    fn clone(&self) -> Self {
930        Self {
931            // Safe as we always create borrowed buffers here
932            data: self.data.clone(),
933            items_in_chunk: self.items_in_chunk,
934            byte_range: self.byte_range.clone(),
935            chunk_idx: self.chunk_idx,
936        }
937    }
938}
939
940/// Decodes mini-block formatted data.  See [`PrimitiveStructuralEncoder`] for more
941/// details on the different layouts.
942#[derive(Debug)]
943struct MiniBlockDecoder {
944    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
945    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
946    value_decompressor: Arc<dyn MiniBlockDecompressor>,
947    def_meaning: Arc<[DefinitionInterpretation]>,
948    loaded_chunks: VecDeque<LoadedChunk>,
949    instructions: VecDeque<ChunkInstructions>,
950    offset_in_current_chunk: u64,
951    num_rows: u64,
952    num_buffers: u64,
953    dictionary: Option<Arc<DataBlock>>,
954    has_large_chunk: bool,
955}
956
957/// See [`MiniBlockScheduler`] for more details on the scheduling and decoding
958/// process for miniblock encoded data.
959impl StructuralPageDecoder for MiniBlockDecoder {
960    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
961        let mut items_desired = num_rows;
962        let mut need_preamble = false;
963        let mut skip_in_chunk = self.offset_in_current_chunk;
964        let mut drain_instructions = Vec::new();
965        while items_desired > 0 || need_preamble {
966            let (instructions, consumed) = self
967                .instructions
968                .front()
969                .unwrap()
970                .drain_from_instruction(&mut items_desired, &mut need_preamble, &mut skip_in_chunk);
971
972            while self.loaded_chunks.front().unwrap().chunk_idx
973                != instructions.chunk_instructions.chunk_idx
974            {
975                self.loaded_chunks.pop_front();
976            }
977            drain_instructions.push((instructions, self.loaded_chunks.front().unwrap().clone()));
978            if consumed {
979                self.instructions.pop_front();
980            }
981        }
982        // We can throw away need_preamble here because it must be false.  If it were true it would mean
983        // we were still in the middle of loading rows.  We do need to latch skip_in_chunk though.
984        self.offset_in_current_chunk = skip_in_chunk;
985
986        let max_visible_level = max_visible_level(&self.def_meaning);
987
988        Ok(Box::new(DecodeMiniBlockTask {
989            instructions: drain_instructions,
990            def_decompressor: self.def_decompressor.clone(),
991            rep_decompressor: self.rep_decompressor.clone(),
992            value_decompressor: self.value_decompressor.clone(),
993            dictionary_data: self.dictionary.clone(),
994            def_meaning: self.def_meaning.clone(),
995            num_buffers: self.num_buffers,
996            max_visible_level,
997            has_large_chunk: self.has_large_chunk,
998        }))
999    }
1000
1001    fn num_rows(&self) -> u64 {
1002        self.num_rows
1003    }
1004}
1005
1006/// How a complex-all-null page's rep/def level buffer is compressed on disk.
1007/// Captured at scheduler construction so `initialize` can keep RLE levels in run
1008/// form instead of expanding them.
1009#[derive(Debug, Clone)]
1010pub(crate) enum LevelCodec {
1011    /// Raw little-endian u16 levels (no block compression).
1012    Uncompressed,
1013    /// RLE-compressed levels; the validated physical runs select their cached representation.
1014    Rle(Arc<RleDecompressor>),
1015    /// Any other block compression; decoded eagerly into [`LazyLevels::Dense`]
1016    /// (these encodings don't expand, so laziness buys nothing).
1017    Block(Arc<dyn BlockDecompressor>),
1018}
1019
1020impl LevelCodec {
1021    fn try_new(
1022        encoding: Option<&CompressiveEncoding>,
1023        decompression_strategy: &dyn DecompressionStrategy,
1024    ) -> Result<Self> {
1025        match encoding {
1026            None => Ok(Self::Uncompressed),
1027            Some(encoding) => match encoding.compression.as_ref() {
1028                Some(Compression::Rle(rle)) => Ok(Self::Rle(Arc::new(create_rle_decompressor(
1029                    rle,
1030                    decompression_strategy,
1031                )?))),
1032                _ => Ok(Self::Block(Arc::from(
1033                    decompression_strategy.create_block_decompressor(encoding)?,
1034                ))),
1035            },
1036        }
1037    }
1038}
1039
1040#[derive(Debug)]
1041enum RunEnds {
1042    U16(Box<[u16]>),
1043    U32(Box<[u32]>),
1044    U64(Box<[u64]>),
1045}
1046
1047impl RunEnds {
1048    fn width_for(num_values: usize) -> usize {
1049        if u16::try_from(num_values).is_ok() {
1050            std::mem::size_of::<u16>()
1051        } else if u32::try_from(num_values).is_ok() {
1052            std::mem::size_of::<u32>()
1053        } else {
1054            std::mem::size_of::<u64>()
1055        }
1056    }
1057
1058    fn len(&self) -> usize {
1059        match self {
1060            Self::U16(ends) => ends.len(),
1061            Self::U32(ends) => ends.len(),
1062            Self::U64(ends) => ends.len(),
1063        }
1064    }
1065
1066    fn get(&self, run: usize) -> usize {
1067        match self {
1068            Self::U16(ends) => ends[run] as usize,
1069            Self::U32(ends) => ends[run] as usize,
1070            Self::U64(ends) => ends[run] as usize,
1071        }
1072    }
1073
1074    fn partition_point(&self, logical_index: usize) -> usize {
1075        match self {
1076            Self::U16(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1077            Self::U32(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1078            Self::U64(ends) => ends.partition_point(|&end| end as usize <= logical_index),
1079        }
1080    }
1081
1082    fn deep_size(&self) -> usize {
1083        match self {
1084            Self::U16(ends) => std::mem::size_of_val(ends.as_ref()),
1085            Self::U32(ends) => std::mem::size_of_val(ends.as_ref()),
1086            Self::U64(ends) => std::mem::size_of_val(ends.as_ref()),
1087        }
1088    }
1089}
1090
1091enum RunEndsBuilder {
1092    U16(Vec<u16>),
1093    U32(Vec<u32>),
1094    U64(Vec<u64>),
1095}
1096
1097impl RunEndsBuilder {
1098    fn with_capacity(num_values: usize, capacity: usize) -> Self {
1099        if u16::try_from(num_values).is_ok() {
1100            Self::U16(Vec::with_capacity(capacity))
1101        } else if u32::try_from(num_values).is_ok() {
1102            Self::U32(Vec::with_capacity(capacity))
1103        } else {
1104            Self::U64(Vec::with_capacity(capacity))
1105        }
1106    }
1107
1108    fn push(&mut self, end: usize) -> Result<()> {
1109        match self {
1110            Self::U16(ends) => ends.push(
1111                u16::try_from(end)
1112                    .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?,
1113            ),
1114            Self::U32(ends) => ends.push(
1115                u32::try_from(end)
1116                    .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?,
1117            ),
1118            Self::U64(ends) => ends.push(end as u64),
1119        }
1120        Ok(())
1121    }
1122
1123    fn set_last(&mut self, end: usize) -> Result<()> {
1124        match self {
1125            Self::U16(ends) => {
1126                let last = ends.last_mut().ok_or_else(|| {
1127                    Error::internal("Cannot extend an empty coalesced run buffer")
1128                })?;
1129                *last = u16::try_from(end)
1130                    .map_err(|_| Error::internal(format!("Run end {end} does not fit in u16")))?;
1131            }
1132            Self::U32(ends) => {
1133                let last = ends.last_mut().ok_or_else(|| {
1134                    Error::internal("Cannot extend an empty coalesced run buffer")
1135                })?;
1136                *last = u32::try_from(end)
1137                    .map_err(|_| Error::internal(format!("Run end {end} does not fit in u32")))?;
1138            }
1139            Self::U64(ends) => {
1140                let last = ends.last_mut().ok_or_else(|| {
1141                    Error::internal("Cannot extend an empty coalesced run buffer")
1142                })?;
1143                *last = end as u64;
1144            }
1145        }
1146        Ok(())
1147    }
1148
1149    fn finish(self) -> RunEnds {
1150        match self {
1151            Self::U16(ends) => RunEnds::U16(ends.into_boxed_slice()),
1152            Self::U32(ends) => RunEnds::U32(ends.into_boxed_slice()),
1153            Self::U64(ends) => RunEnds::U64(ends.into_boxed_slice()),
1154        }
1155    }
1156}
1157
1158#[derive(Debug)]
1159enum RunStorage {
1160    Physical(RleRuns),
1161    Coalesced { values: Box<[u16]>, ends: RunEnds },
1162}
1163
1164impl RunStorage {
1165    fn len(&self) -> usize {
1166        match self {
1167            Self::Physical(runs) => runs.num_values(),
1168            Self::Coalesced { ends, .. } => ends.get(ends.len() - 1),
1169        }
1170    }
1171
1172    fn num_runs(&self) -> usize {
1173        match self {
1174            Self::Physical(runs) => runs.num_runs(),
1175            Self::Coalesced { values, .. } => values.len(),
1176        }
1177    }
1178
1179    fn value(&self, run: usize) -> u16 {
1180        match self {
1181            Self::Physical(runs) => runs.value(run),
1182            Self::Coalesced { values, .. } => values[run],
1183        }
1184    }
1185
1186    fn first_value_above(&self, max: u16) -> Option<(usize, u16)> {
1187        (0..self.num_runs()).find_map(|run| {
1188            let value = self.value(run);
1189            (value > max).then_some((run, value))
1190        })
1191    }
1192
1193    fn seek(&self, position: &mut RunPosition, logical_index: usize) {
1194        if logical_index >= self.len() {
1195            *position = RunPosition {
1196                run: self.num_runs(),
1197                start: self.len(),
1198                end: self.len(),
1199            };
1200            return;
1201        }
1202
1203        match self {
1204            Self::Physical(runs) => {
1205                if position.run >= runs.num_runs()
1206                    || position.end == 0
1207                    || logical_index < position.start
1208                {
1209                    *position = RunPosition {
1210                        run: 0,
1211                        start: 0,
1212                        end: runs.length(0),
1213                    };
1214                }
1215                while position.end <= logical_index {
1216                    self.advance(position);
1217                }
1218            }
1219            Self::Coalesced { ends, .. } => {
1220                if logical_index < position.start || logical_index >= position.end {
1221                    let run = ends.partition_point(logical_index);
1222                    *position = RunPosition {
1223                        run,
1224                        start: if run == 0 { 0 } else { ends.get(run - 1) },
1225                        end: ends.get(run),
1226                    };
1227                }
1228            }
1229        }
1230    }
1231
1232    fn advance(&self, position: &mut RunPosition) {
1233        let next_run = position.run + 1;
1234        if next_run >= self.num_runs() {
1235            *position = RunPosition {
1236                run: self.num_runs(),
1237                start: self.len(),
1238                end: self.len(),
1239            };
1240            return;
1241        }
1242
1243        let start = position.end;
1244        position.run = next_run;
1245        position.start = start;
1246        position.end = match self {
1247            Self::Physical(runs) => start + runs.length(next_run),
1248            Self::Coalesced { ends, .. } => ends.get(next_run),
1249        };
1250    }
1251
1252    fn deep_size(&self) -> usize {
1253        match self {
1254            Self::Physical(runs) => runs.deep_size(),
1255            Self::Coalesced { values, ends } => {
1256                std::mem::size_of_val(values.as_ref()) + ends.deep_size()
1257            }
1258        }
1259    }
1260}
1261
1262/// Rep/def levels for a complex-all-null page.
1263///
1264/// RLE pages retain the smallest of their validated physical runs, coalesced
1265/// runs, and dense values. The decoder materializes only the per-drain slices
1266/// it touches.
1267#[derive(Debug, Clone)]
1268enum LazyLevels {
1269    Dense(ScalarBuffer<u16>),
1270    Runs(Arc<RunStorage>),
1271}
1272
1273#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1274enum LevelPlan {
1275    Physical,
1276    Coalesced,
1277    Dense,
1278}
1279
1280#[derive(Debug, Default, Clone, Copy)]
1281struct RunPosition {
1282    run: usize,
1283    start: usize,
1284    end: usize,
1285}
1286
1287/// Monotonic forward cursor into a [`LazyLevels`] sequence.
1288///
1289/// Drains seek to strictly increasing rows, so each [`LazyLevels::seek_row_start`]
1290/// resumes from the last position instead of rescanning — every run is visited at
1291/// most once per page while locating and counting level ranges.
1292#[derive(Debug, Default, Clone, Copy)]
1293struct LevelCursor {
1294    /// Logical level index where the current row begins.
1295    level: usize,
1296    /// Row index at `level` (the number of `max_rep` occurrences before it).
1297    row: u64,
1298    /// Run containing `level`. Unused for [`LazyLevels::Dense`].
1299    run: RunPosition,
1300}
1301
1302impl LazyLevels {
1303    fn from_rle_runs(runs: RleRuns) -> Result<Self> {
1304        let plan = Self::select_plan(&runs);
1305        match plan {
1306            LevelPlan::Physical => Ok(Self::Runs(Arc::new(RunStorage::Physical(
1307                runs.into_owned(),
1308            )))),
1309            LevelPlan::Coalesced => Self::build_coalesced(runs),
1310            LevelPlan::Dense => Self::build_dense(runs),
1311        }
1312    }
1313
1314    /// Minimize retained payload bytes first, then expected traversal work.
1315    /// If both are equal, keep the physical runs and avoid another allocation.
1316    fn select_plan(runs: &RleRuns) -> LevelPlan {
1317        if runs.num_values() == 0 {
1318            return LevelPlan::Dense;
1319        }
1320
1321        let run_storage_size = std::mem::size_of::<RunStorage>() as u128;
1322        let physical_size = run_storage_size + runs.owned_size() as u128;
1323        let coalesced_size = run_storage_size
1324            + (runs.coalesced_runs() as u128)
1325                * (std::mem::size_of::<u16>() + RunEnds::width_for(runs.num_values())) as u128;
1326        let dense_size = (runs.num_values() as u128) * std::mem::size_of::<u16>() as u128;
1327        [
1328            (physical_size, runs.num_runs(), 0usize, LevelPlan::Physical),
1329            (
1330                coalesced_size,
1331                runs.coalesced_runs(),
1332                1usize,
1333                LevelPlan::Coalesced,
1334            ),
1335            (dense_size, runs.num_values(), 2usize, LevelPlan::Dense),
1336        ]
1337        .into_iter()
1338        .min_by_key(|(size, traversal, priority, _)| (*size, *traversal, *priority))
1339        .map(|(_, _, _, plan)| plan)
1340        .unwrap_or(LevelPlan::Dense)
1341    }
1342
1343    fn build_coalesced(runs: RleRuns) -> Result<Self> {
1344        let mut values = Vec::with_capacity(runs.coalesced_runs());
1345        let mut ends = RunEndsBuilder::with_capacity(runs.num_values(), runs.coalesced_runs());
1346        let mut logical_end = 0usize;
1347        for (value, length) in runs.iter() {
1348            logical_end = logical_end
1349                .checked_add(length)
1350                .ok_or_else(|| Error::internal("Validated RLE run length sum overflowed usize"))?;
1351            if values.last().copied() == Some(value) {
1352                ends.set_last(logical_end)?;
1353            } else {
1354                values.push(value);
1355                ends.push(logical_end)?;
1356            }
1357        }
1358        Ok(Self::Runs(Arc::new(RunStorage::Coalesced {
1359            values: values.into_boxed_slice(),
1360            ends: ends.finish(),
1361        })))
1362    }
1363
1364    fn build_dense(runs: RleRuns) -> Result<Self> {
1365        let mut values = Vec::new();
1366        values.try_reserve_exact(runs.num_values()).map_err(|_| {
1367            Error::internal(format!(
1368                "Cannot allocate {} dense repetition/definition levels",
1369                runs.num_values()
1370            ))
1371        })?;
1372        for (value, length) in runs.iter() {
1373            values.resize(values.len() + length, value);
1374        }
1375        Ok(Self::Dense(ScalarBuffer::from(values)))
1376    }
1377
1378    fn len(&self) -> usize {
1379        match self {
1380            Self::Dense(buf) => buf.len(),
1381            Self::Runs(runs) => runs.len(),
1382        }
1383    }
1384
1385    fn validate_max(&self, level_type: &str, max: u16) -> Result<()> {
1386        let invalid = match self {
1387            Self::Dense(levels) => levels
1388                .iter()
1389                .enumerate()
1390                .find_map(|(index, &value)| (value > max).then_some(("index", index, value))),
1391            Self::Runs(runs) => runs
1392                .first_value_above(max)
1393                .map(|(run, value)| ("run", run, value)),
1394        };
1395        if let Some((position_type, position, value)) = invalid {
1396            return Err(Error::invalid_input_source(
1397                format!(
1398                    "Invalid {level_type} level {value} at {position_type} {position}: maximum is {max}"
1399                )
1400                .into(),
1401            ));
1402        }
1403        Ok(())
1404    }
1405
1406    /// Advance `cursor` to the start of row `target_row`, returning that row's
1407    /// starting level index.
1408    ///
1409    /// Rows begin at `max_rep` positions, so this finds the `target_row`-th one.
1410    /// `target_row` must be `>= cursor.row`: the cursor only moves forward, which
1411    /// is what keeps a full page decode O(runs) rather than O(rows).
1412    fn seek_row_start(
1413        &self,
1414        cursor: &mut LevelCursor,
1415        target_row: u64,
1416        max_rep: u16,
1417    ) -> Result<usize> {
1418        let mut need = target_row.checked_sub(cursor.row).ok_or_else(|| {
1419            Error::internal(format!(
1420                "Complex all-null row ranges are not sorted: target row {target_row} follows {}",
1421                cursor.row
1422            ))
1423        })?;
1424        if need == 0 {
1425            return Ok(cursor.level);
1426        }
1427        match self {
1428            Self::Dense(buf) => {
1429                let mut level = cursor.level;
1430                while need > 0 {
1431                    if level >= buf.len() {
1432                        return Err(Error::internal(
1433                            "Invalid complex all-null layout: repetition buffer too short",
1434                        ));
1435                    }
1436                    if buf[level] != max_rep {
1437                        return Err(Error::internal(
1438                            "Invalid complex all-null layout: row did not start at max repetition level",
1439                        ));
1440                    }
1441                    level += 1;
1442                    while level < buf.len() && buf[level] != max_rep {
1443                        level += 1;
1444                    }
1445                    need -= 1;
1446                }
1447                cursor.level = level;
1448                cursor.row = target_row;
1449                Ok(level)
1450            }
1451            Self::Runs(runs) => {
1452                let mut level = cursor.level;
1453                let mut run = cursor.run;
1454                runs.seek(&mut run, level);
1455                while need > 0 {
1456                    if run.run >= runs.num_runs() {
1457                        return Err(Error::internal(
1458                            "Invalid complex all-null layout: repetition buffer too short",
1459                        ));
1460                    }
1461                    if runs.value(run.run) != max_rep {
1462                        return Err(Error::internal(
1463                            "Invalid complex all-null layout: row did not start at max repetition level",
1464                        ));
1465                    }
1466                    let avail = (run.end - level) as u64;
1467                    if need < avail {
1468                        // Target lands inside this max-rep run.
1469                        level += need as usize;
1470                        need = 0;
1471                    } else {
1472                        // Consume every row start in this run, then skip the
1473                        // trailing non-max-rep runs to reach the next row start.
1474                        need -= avail;
1475                        runs.advance(&mut run);
1476                        while run.run < runs.num_runs() && runs.value(run.run) != max_rep {
1477                            runs.advance(&mut run);
1478                        }
1479                        level = if run.run < runs.num_runs() {
1480                            run.start
1481                        } else {
1482                            self.len()
1483                        };
1484                    }
1485                }
1486                cursor.level = level;
1487                cursor.row = target_row;
1488                cursor.run = run;
1489                Ok(level)
1490            }
1491        }
1492    }
1493
1494    /// Count of levels in `range` that are `<= max`, resuming from `*run_cursor`
1495    /// and leaving it on the last run that overlaps `range`.
1496    ///
1497    /// Successive calls must pass ascending, non-overlapping ranges (`range.start
1498    /// >=` the previous `range.end`) so runs are swept at most once per page.
1499    fn count_le_cursor(
1500        &self,
1501        run_cursor: &mut RunPosition,
1502        range: Range<usize>,
1503        max: u16,
1504    ) -> (u64, RunPosition) {
1505        if range.is_empty() {
1506            return (0, *run_cursor);
1507        }
1508        match self {
1509            Self::Dense(buf) => (
1510                buf[range].iter().filter(|&&d| d <= max).count() as u64,
1511                RunPosition::default(),
1512            ),
1513            Self::Runs(runs) => {
1514                // Advance to the first run overlapping the range.
1515                runs.seek(run_cursor, range.start);
1516                let start = *run_cursor;
1517                let mut count = 0u64;
1518                let mut current = *run_cursor;
1519                while current.run < runs.num_runs() && current.start < range.end {
1520                    if runs.value(current.run) <= max {
1521                        let lo = current.start.max(range.start);
1522                        let hi = current.end.min(range.end);
1523                        count += (hi - lo) as u64;
1524                    }
1525                    if current.end >= range.end {
1526                        break;
1527                    }
1528                    runs.advance(&mut current);
1529                }
1530                // Resume the next (ascending) range from the last overlapping run;
1531                // `current` remains valid because `range` is non-empty.
1532                *run_cursor = current;
1533                (count, start)
1534            }
1535        }
1536    }
1537
1538    fn extend_into(&self, range: Range<usize>, run: RunPosition, out: &mut Vec<u16>) {
1539        if range.is_empty() {
1540            return;
1541        }
1542        match self {
1543            Self::Dense(buf) => out.extend_from_slice(&buf[range]),
1544            Self::Runs(runs) => {
1545                let mut current = run;
1546                runs.seek(&mut current, range.start);
1547                while current.run < runs.num_runs() && current.start < range.end {
1548                    let lo = current.start.max(range.start);
1549                    let hi = current.end.min(range.end);
1550                    if hi > lo {
1551                        out.resize(out.len() + (hi - lo), runs.value(current.run));
1552                    }
1553                    runs.advance(&mut current);
1554                }
1555            }
1556        }
1557    }
1558
1559    #[cfg(test)]
1560    fn deep_size(&self) -> usize {
1561        self.deep_size_of_children(&mut Context::new())
1562    }
1563}
1564
1565impl DeepSizeOf for LazyLevels {
1566    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1567        match self {
1568            Self::Dense(buf) => buf.deep_size_of_children(ctx),
1569            Self::Runs(runs) => {
1570                let pointer = Arc::as_ptr(runs) as *const () as usize;
1571                if ctx.mark_seen(pointer) {
1572                    std::mem::size_of_val(runs.as_ref()) + runs.deep_size()
1573                } else {
1574                    0
1575                }
1576            }
1577        }
1578    }
1579}
1580
1581fn validate_complex_all_null_levels(
1582    rep: &Option<LazyLevels>,
1583    def: &Option<LazyLevels>,
1584    max_rep: u16,
1585    max_def: u16,
1586) -> Result<()> {
1587    if let Some(rep) = rep {
1588        rep.validate_max("repetition", max_rep)?;
1589    }
1590    if let Some(def) = def {
1591        def.validate_max("definition", max_def)?;
1592    }
1593    if let (Some(rep), Some(def)) = (rep, def)
1594        && rep.len() != def.len()
1595    {
1596        return Err(Error::invalid_input_source(
1597            format!(
1598                "Mismatched complex all-null level counts: repetition has {}, definition has {}",
1599                rep.len(),
1600                def.len()
1601            )
1602            .into(),
1603        ));
1604    }
1605    Ok(())
1606}
1607
1608fn expected_level_bytes(num_values: u64, level_type: &str) -> Result<usize> {
1609    usize::try_from(num_values)
1610        .ok()
1611        .and_then(|num_values| num_values.checked_mul(std::mem::size_of::<u16>()))
1612        .ok_or_else(|| {
1613            Error::invalid_input_source(
1614                format!("{level_type} level count {num_values} does not fit in memory").into(),
1615            )
1616        })
1617}
1618
1619fn dense_levels_from_block(
1620    decompressed: DataBlock,
1621    num_values: u64,
1622    level_type: &str,
1623) -> Result<LazyLevels> {
1624    let DataBlock::FixedWidth(block) = decompressed else {
1625        return Err(Error::invalid_input_source(
1626            format!("Expected fixed-width data block for {level_type} levels").into(),
1627        ));
1628    };
1629    if block.num_values != num_values {
1630        return Err(Error::invalid_input_source(
1631            format!(
1632                "Unexpected {level_type} level count after decompression: expected {num_values}, got {}",
1633                block.num_values
1634            )
1635            .into(),
1636        ));
1637    }
1638    if block.bits_per_value != 16 {
1639        return Err(Error::invalid_input_source(
1640            format!(
1641                "Unexpected {level_type} level bit width after decompression: expected 16, got {}",
1642                block.bits_per_value
1643            )
1644            .into(),
1645        ));
1646    }
1647    let expected_bytes = expected_level_bytes(num_values, level_type)?;
1648    if block.data.len() != expected_bytes {
1649        return Err(Error::invalid_input_source(
1650            format!(
1651                "Unexpected decompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1652                block.data.len()
1653            )
1654            .into(),
1655        ));
1656    }
1657    Ok(LazyLevels::Dense(block.data.borrow_to_typed_slice::<u16>()))
1658}
1659
1660#[derive(Debug)]
1661struct CachedComplexAllNullState {
1662    rep: Option<LazyLevels>,
1663    def: Option<LazyLevels>,
1664}
1665
1666impl DeepSizeOf for CachedComplexAllNullState {
1667    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
1668        self.rep.deep_size_of_children(ctx) + self.def.deep_size_of_children(ctx)
1669    }
1670}
1671
1672impl CachedPageData for CachedComplexAllNullState {
1673    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
1674        self
1675    }
1676}
1677
1678/// A scheduler for all-null data that has repetition and definition levels
1679///
1680/// We still need to do some I/O in this case because we need to figure out what kind of null we
1681/// are dealing with (null list, null struct, what level null struct, etc.)
1682///
1683/// TODO: Right now we just load the entire rep/def at initialization time and cache it.  This is a touch
1684/// RAM aggressive and maybe we want something more lazy in the future.  On the other hand, it's simple
1685/// and fast so...maybe not :)
1686#[derive(Debug)]
1687pub struct ComplexAllNullScheduler {
1688    // Set from protobuf
1689    buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1690    def_meaning: Arc<[DefinitionInterpretation]>,
1691    repdef: Option<Arc<CachedComplexAllNullState>>,
1692    max_rep: u16,
1693    max_def: u16,
1694    max_visible_level: u16,
1695    rep_codec: LevelCodec,
1696    def_codec: LevelCodec,
1697    num_rep_values: u64,
1698    num_def_values: u64,
1699}
1700
1701impl ComplexAllNullScheduler {
1702    pub(crate) fn new(
1703        buffer_offsets_and_sizes: Arc<[(u64, u64)]>,
1704        def_meaning: Arc<[DefinitionInterpretation]>,
1705        rep_codec: LevelCodec,
1706        def_codec: LevelCodec,
1707        num_rep_values: u64,
1708        num_def_values: u64,
1709    ) -> Self {
1710        let max_rep = def_meaning.iter().filter(|l| l.is_list()).count() as u16;
1711        let max_def = def_meaning
1712            .iter()
1713            .map(|meaning| meaning.num_def_levels())
1714            .sum::<u16>();
1715        let max_visible_level = max_visible_level(&def_meaning);
1716        Self {
1717            buffer_offsets_and_sizes,
1718            def_meaning,
1719            repdef: None,
1720            max_rep,
1721            max_def,
1722            max_visible_level,
1723            rep_codec,
1724            def_codec,
1725            num_rep_values,
1726            num_def_values,
1727        }
1728    }
1729}
1730
1731impl StructuralPageScheduler for ComplexAllNullScheduler {
1732    fn init_layout(&self) -> Result<PageInitialization> {
1733        // Fully load the rep & def buffers, as needed
1734        let (rep_pos, rep_size) = self.buffer_offsets_and_sizes[0];
1735        let (def_pos, def_size) = self.buffer_offsets_and_sizes[1];
1736
1737        Ok(PageInitialization::new([
1738            (rep_size > 0).then_some(rep_pos..rep_pos + rep_size),
1739            (def_size > 0).then_some(def_pos..def_pos + def_size),
1740        ]))
1741    }
1742
1743    fn init_from_buffers<'a>(
1744        &'a mut self,
1745        buffers: PageInitializationBuffers,
1746        _io: &Arc<dyn EncodingsIo>,
1747    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
1748        let rep_codec = self.rep_codec.clone();
1749        let def_codec = self.def_codec.clone();
1750        let num_rep_values = self.num_rep_values;
1751        let num_def_values = self.num_def_values;
1752        let max_rep = self.max_rep;
1753        let max_def = self.max_def;
1754
1755        async move {
1756            let [rep_bytes, def_bytes] = buffers.try_into_array("complex all-null")?;
1757
1758            // RLE levels select the smallest validated cache representation;
1759            // everything else expands eagerly to `LazyLevels::Dense`.
1760            let build_levels = |compressed_bytes: Bytes,
1761                                codec: &LevelCodec,
1762                                num_values: u64,
1763                                level_type: &str|
1764             -> Result<LazyLevels> {
1765                match codec {
1766                    LevelCodec::Uncompressed => {
1767                        if num_values == 0 {
1768                            if !compressed_bytes
1769                                .len()
1770                                .is_multiple_of(std::mem::size_of::<u16>())
1771                            {
1772                                return Err(Error::invalid_input_source(
1773                                    format!(
1774                                        "Unexpected uncompressed {level_type} level size: {} bytes is not divisible by {}",
1775                                        compressed_bytes.len(),
1776                                        std::mem::size_of::<u16>()
1777                                    )
1778                                    .into(),
1779                                ));
1780                            }
1781                        } else {
1782                            let expected_bytes = expected_level_bytes(num_values, level_type)?;
1783                            if compressed_bytes.len() != expected_bytes {
1784                                return Err(Error::invalid_input_source(
1785                                    format!(
1786                                        "Unexpected uncompressed {level_type} level size: expected {expected_bytes} bytes for {num_values} values, got {}",
1787                                        compressed_bytes.len()
1788                                    )
1789                                    .into(),
1790                                ));
1791                            }
1792                        }
1793                        let buffer = LanceBuffer::from_bytes(compressed_bytes, 2);
1794                        Ok(LazyLevels::Dense(buffer.borrow_to_typed_slice::<u16>()))
1795                    }
1796                    LevelCodec::Rle(decompressor) => {
1797                        let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1798                        let runs = decompressor.decode_u16_runs(frame, num_values)?;
1799                        LazyLevels::from_rle_runs(runs)
1800                    }
1801                    LevelCodec::Block(decompressor) => {
1802                        let frame = LanceBuffer::from_bytes(compressed_bytes, 1);
1803                        let decompressed = decompressor.decompress(Some(frame), num_values)?;
1804                        dense_levels_from_block(decompressed, num_values, level_type)
1805                    }
1806                }
1807            };
1808
1809            let rep = if let Some(rep) = rep_bytes {
1810                Some(build_levels(rep, &rep_codec, num_rep_values, "repetition")?)
1811            } else {
1812                None
1813            };
1814
1815            let def = if let Some(def) = def_bytes {
1816                Some(build_levels(def, &def_codec, num_def_values, "definition")?)
1817            } else {
1818                None
1819            };
1820
1821            validate_complex_all_null_levels(&rep, &def, max_rep, max_def)?;
1822            let repdef = Arc::new(CachedComplexAllNullState { rep, def });
1823
1824            Ok(repdef as Arc<dyn CachedPageData>)
1825        }
1826        .boxed()
1827    }
1828
1829    fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
1830        self.repdef = Some(data.clone().as_arc_any().downcast().map_err(|_| {
1831            Error::invalid_input_source(
1832                "Cached complex all-null page data has an unexpected type".into(),
1833            )
1834        })?);
1835        Ok(())
1836    }
1837
1838    fn schedule_ranges(
1839        &self,
1840        ranges: &[Range<u64>],
1841        _io: &Arc<dyn EncodingsIo>,
1842    ) -> Result<Vec<PageLoadTask>> {
1843        let ranges = VecDeque::from_iter(ranges.iter().cloned());
1844        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1845        let decoder = Box::new(ComplexAllNullPageDecoder {
1846            ranges,
1847            rep: self.repdef.as_ref().unwrap().rep.clone(),
1848            def: self.repdef.as_ref().unwrap().def.clone(),
1849            num_rows,
1850            def_meaning: self.def_meaning.clone(),
1851            max_rep: self.max_rep,
1852            max_visible_level: self.max_visible_level,
1853            rep_cursor: LevelCursor::default(),
1854            def_run_cursor: RunPosition::default(),
1855        }) as Box<dyn StructuralPageDecoder>;
1856        let page_load_task = PageLoadTask {
1857            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
1858            num_rows,
1859        };
1860        Ok(vec![page_load_task])
1861    }
1862}
1863
1864#[derive(Debug)]
1865pub struct ComplexAllNullPageDecoder {
1866    ranges: VecDeque<Range<u64>>,
1867    rep: Option<LazyLevels>,
1868    def: Option<LazyLevels>,
1869    num_rows: u64,
1870    def_meaning: Arc<[DefinitionInterpretation]>,
1871    max_rep: u16,
1872    max_visible_level: u16,
1873    /// Monotonic cursor into `rep` tracking the current row's level start.
1874    rep_cursor: LevelCursor,
1875    /// Monotonic run cursor into `def` for `count_le_cursor`.
1876    def_run_cursor: RunPosition,
1877}
1878
1879impl ComplexAllNullPageDecoder {
1880    fn drain_ranges(&mut self, num_rows: u64) -> Vec<Range<u64>> {
1881        let mut rows_desired = num_rows;
1882        let mut ranges = Vec::with_capacity(self.ranges.len());
1883        while rows_desired > 0 {
1884            let front = self.ranges.front_mut().unwrap();
1885            let avail = front.end - front.start;
1886            if avail > rows_desired {
1887                ranges.push(front.start..front.start + rows_desired);
1888                front.start += rows_desired;
1889                rows_desired = 0;
1890            } else {
1891                ranges.push(self.ranges.pop_front().unwrap());
1892                rows_desired -= avail;
1893            }
1894        }
1895        ranges
1896    }
1897
1898    /// Level index at which row `target_row` starts, advancing the monotonic
1899    /// repetition cursor. Callers must request non-decreasing `target_row`.
1900    fn seek_row_start(&mut self, target_row: u64) -> Result<usize> {
1901        match &self.rep {
1902            Some(rep) => rep.seek_row_start(&mut self.rep_cursor, target_row, self.max_rep),
1903            None => {
1904                // Without repetition every level is its own row.
1905                self.rep_cursor.row = target_row;
1906                self.rep_cursor.level = target_row as usize;
1907                Ok(target_row as usize)
1908            }
1909        }
1910    }
1911
1912    /// Number of visible items in the level range `levels` (definition levels
1913    /// `<= max_visible_level`), advancing the monotonic definition cursor.
1914    fn count_visible(&mut self, levels: Range<usize>) -> Result<(u64, RunPosition)> {
1915        match &self.def {
1916            Some(def) => {
1917                if levels.end > def.len() {
1918                    return Err(Error::internal(
1919                        "Invalid complex all-null layout: definition buffer too short",
1920                    ));
1921                }
1922                Ok(def.count_le_cursor(&mut self.def_run_cursor, levels, self.max_visible_level))
1923            }
1924            None => Ok(((levels.end - levels.start) as u64, RunPosition::default())),
1925        }
1926    }
1927}
1928
1929impl StructuralPageDecoder for ComplexAllNullPageDecoder {
1930    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
1931        let drained_ranges = self.drain_ranges(num_rows);
1932        let mut level_slices: Vec<LevelSlice> = Vec::with_capacity(drained_ranges.len());
1933        let mut visible_items_total = 0;
1934
1935        // Each row range is one contiguous level slice `[start_row_level,
1936        // end_row_level)`, so we seek both boundaries and count its visibility at
1937        // once rather than per row. The cursors only move forward, so locating and
1938        // counting all requested ranges visits each intervening run at most once.
1939        for range in drained_ranges {
1940            let level_start = self.seek_row_start(range.start)?;
1941            let rep_run = self.rep_cursor.run;
1942            let level_end = self.seek_row_start(range.end)?;
1943            let (visible_items, def_run) = self.count_visible(level_start..level_end)?;
1944            visible_items_total += visible_items;
1945            if let Some(last) = level_slices.last_mut()
1946                && last.range.end == level_start
1947            {
1948                last.range.end = level_end;
1949            } else {
1950                level_slices.push(LevelSlice {
1951                    range: level_start..level_end,
1952                    rep_run,
1953                    def_run,
1954                });
1955            }
1956        }
1957
1958        Ok(Box::new(DecodeComplexAllNullTask {
1959            level_slices,
1960            visible_items_total,
1961            rep: self.rep.clone(),
1962            def: self.def.clone(),
1963            def_meaning: self.def_meaning.clone(),
1964            max_visible_level: self.max_visible_level,
1965        }))
1966    }
1967
1968    fn num_rows(&self) -> u64 {
1969        self.num_rows
1970    }
1971}
1972
1973/// We use `level_slices` to slice into `rep` and `def` and create rep/def buffers
1974/// for the null data.
1975#[derive(Debug, Clone)]
1976struct LevelSlice {
1977    range: Range<usize>,
1978    rep_run: RunPosition,
1979    def_run: RunPosition,
1980}
1981
1982#[derive(Clone, Copy)]
1983enum LevelKind {
1984    Repetition,
1985    Definition,
1986}
1987
1988impl LevelSlice {
1989    fn run(&self, kind: LevelKind) -> RunPosition {
1990        match kind {
1991            LevelKind::Repetition => self.rep_run,
1992            LevelKind::Definition => self.def_run,
1993        }
1994    }
1995}
1996
1997#[derive(Debug)]
1998pub struct DecodeComplexAllNullTask {
1999    level_slices: Vec<LevelSlice>,
2000    visible_items_total: u64,
2001    rep: Option<LazyLevels>,
2002    def: Option<LazyLevels>,
2003    def_meaning: Arc<[DefinitionInterpretation]>,
2004    max_visible_level: u16,
2005}
2006
2007impl DecodeComplexAllNullTask {
2008    fn decode_level(&self, levels: &Option<LazyLevels>, kind: LevelKind) -> Option<Vec<u16>> {
2009        levels.as_ref().map(|levels| {
2010            let num_levels = self
2011                .level_slices
2012                .iter()
2013                .map(|slice| slice.range.end - slice.range.start)
2014                .sum();
2015            let mut referenced_levels = Vec::with_capacity(num_levels);
2016            for slice in &self.level_slices {
2017                levels.extend_into(slice.range.clone(), slice.run(kind), &mut referenced_levels);
2018            }
2019            referenced_levels
2020        })
2021    }
2022}
2023
2024impl DecodePageTask for DecodeComplexAllNullTask {
2025    fn decode(self: Box<Self>) -> Result<DecodedPage> {
2026        let rep = self.decode_level(&self.rep, LevelKind::Repetition);
2027        let def = self.decode_level(&self.def, LevelKind::Definition);
2028
2029        // If there are definition levels there may be empty / null lists which are not visible
2030        // in the items array.  We need to account for that here to figure out how many values
2031        // should be in the items array.
2032        let num_values = if let Some(def) = &def {
2033            def.iter().filter(|&d| *d <= self.max_visible_level).count() as u64
2034        } else {
2035            self.visible_items_total
2036        };
2037
2038        let data = DataBlock::AllNull(AllNullDataBlock { num_values });
2039        let unraveler = RepDefUnraveler::new(rep, def, self.def_meaning, num_values);
2040        Ok(DecodedPage {
2041            data,
2042            repdef: unraveler,
2043        })
2044    }
2045}
2046
2047/// A scheduler for simple all-null data
2048///
2049/// "simple" all-null data is data that is all null and only has a single level of definition and
2050/// no repetition.  We don't need to read any data at all in this case.
2051#[derive(Debug, Default)]
2052pub struct SimpleAllNullScheduler {}
2053
2054impl StructuralPageScheduler for SimpleAllNullScheduler {
2055    fn needs_initialization(&self) -> bool {
2056        false
2057    }
2058
2059    fn init_layout(&self) -> Result<PageInitialization> {
2060        Ok(PageInitialization::new([]))
2061    }
2062
2063    fn init_from_buffers<'a>(
2064        &'a mut self,
2065        _buffers: PageInitializationBuffers,
2066        _io: &Arc<dyn EncodingsIo>,
2067    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2068        std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
2069    }
2070
2071    fn try_load(&mut self, _cache: &Arc<dyn CachedPageData>) -> Result<()> {
2072        Ok(())
2073    }
2074
2075    fn schedule_ranges(
2076        &self,
2077        ranges: &[Range<u64>],
2078        _io: &Arc<dyn EncodingsIo>,
2079    ) -> Result<Vec<PageLoadTask>> {
2080        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
2081        let decoder =
2082            Box::new(SimpleAllNullPageDecoder { num_rows }) as Box<dyn StructuralPageDecoder>;
2083        let page_load_task = PageLoadTask {
2084            decoder_fut: std::future::ready(Ok(decoder)).boxed(),
2085            num_rows,
2086        };
2087        Ok(vec![page_load_task])
2088    }
2089}
2090
2091/// A page decode task for all-null data without any
2092/// repetition and only a single level of definition
2093#[derive(Debug)]
2094struct SimpleAllNullDecodePageTask {
2095    num_values: u64,
2096}
2097impl DecodePageTask for SimpleAllNullDecodePageTask {
2098    fn decode(self: Box<Self>) -> Result<DecodedPage> {
2099        let unraveler = RepDefUnraveler::new(
2100            None,
2101            Some(vec![1; self.num_values as usize]),
2102            Arc::new([DefinitionInterpretation::NullableItem]),
2103            self.num_values,
2104        );
2105        Ok(DecodedPage {
2106            data: DataBlock::AllNull(AllNullDataBlock {
2107                num_values: self.num_values,
2108            }),
2109            repdef: unraveler,
2110        })
2111    }
2112}
2113
2114#[derive(Debug)]
2115pub struct SimpleAllNullPageDecoder {
2116    num_rows: u64,
2117}
2118
2119impl StructuralPageDecoder for SimpleAllNullPageDecoder {
2120    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
2121        Ok(Box::new(SimpleAllNullDecodePageTask {
2122            num_values: num_rows,
2123        }))
2124    }
2125
2126    fn num_rows(&self) -> u64 {
2127        self.num_rows
2128    }
2129}
2130
2131#[derive(Debug, Clone)]
2132struct MiniBlockSchedulerDictionary {
2133    // These come from the protobuf
2134    dictionary_decompressor: Arc<dyn BlockDecompressor>,
2135    dictionary_buf_position_and_size: (u64, u64),
2136    dictionary_data_alignment: u64,
2137    num_dictionary_items: u64,
2138}
2139
2140/// State that is loaded once and cached for future lookups
2141#[derive(Debug)]
2142struct MiniBlockCacheableState {
2143    /// Compact per-chunk index (byte ranges + row/item mapping) for the page
2144    chunk_index: MiniBlockChunkIndex,
2145    /// The dictionary for the page, if any
2146    dictionary: Option<Arc<DataBlock>>,
2147}
2148
2149impl DeepSizeOf for MiniBlockCacheableState {
2150    fn deep_size_of_children(&self, context: &mut Context) -> usize {
2151        self.chunk_index.deep_size_of_children(context)
2152            + self
2153                .dictionary
2154                .as_ref()
2155                .map(|dict| dict.data_size() as usize)
2156                .unwrap_or(0)
2157    }
2158}
2159
2160impl CachedPageData for MiniBlockCacheableState {
2161    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
2162        self
2163    }
2164}
2165
2166/// A scheduler for a page that has been encoded with the mini-block layout
2167///
2168/// Scheduling mini-block encoded data is simple in concept and somewhat complex
2169/// in practice.
2170///
2171/// First, during initialization, we load the chunk metadata, the repetition index,
2172/// and the dictionary (these last two may not be present)
2173///
2174/// Then, during scheduling, we use the user's requested row ranges and the repetition
2175/// index to determine which chunks we need and which rows we need from those chunks.
2176///
2177/// For example, if the repetition index is: [50, 3], [50, 0], [10, 0] and the range
2178/// from the user is 40..60 then we need to:
2179///
2180///  - Read the first chunk and skip the first 40 rows, then read 10 full rows, and
2181///    then read 3 items for the 11th row of our range.
2182///  - Read the second chunk and read the remaining items in our 11th row and then read
2183///    the remaining 9 full rows.
2184///
2185/// Then, if we are going to decode that in batches of 5, we need to make decode tasks.
2186/// The first two decode tasks will just need the first chunk.  The third decode task will
2187/// need the first chunk (for the trailer which has the 11th row in our range) and the second
2188/// chunk.  The final decode task will just need the second chunk.
2189///
2190/// The above prose descriptions are what are represented by `ChunkInstructions` and
2191/// `ChunkDrainInstructions`.
2192#[derive(Debug)]
2193pub struct MiniBlockScheduler {
2194    // These come from the protobuf
2195    buffer_offsets_and_sizes: Vec<(u64, u64)>,
2196    priority: u64,
2197    items_in_page: u64,
2198    repetition_index_depth: u16,
2199    num_buffers: u64,
2200    rep_decompressor: Option<Arc<dyn BlockDecompressor>>,
2201    def_decompressor: Option<Arc<dyn BlockDecompressor>>,
2202    value_decompressor: Arc<dyn MiniBlockDecompressor>,
2203    def_meaning: Arc<[DefinitionInterpretation]>,
2204    dictionary: Option<MiniBlockSchedulerDictionary>,
2205    // This is set after initialization
2206    page_meta: Option<Arc<MiniBlockCacheableState>>,
2207    has_large_chunk: bool,
2208}
2209
2210impl MiniBlockScheduler {
2211    fn try_new(
2212        buffer_offsets_and_sizes: &[(u64, u64)],
2213        priority: u64,
2214        items_in_page: u64,
2215        layout: &pb21::MiniBlockLayout,
2216        decompressors: &dyn DecompressionStrategy,
2217    ) -> Result<Self> {
2218        let rep_decompressor = layout
2219            .rep_compression
2220            .as_ref()
2221            .map(|rep_compression| {
2222                decompressors
2223                    .create_block_decompressor(rep_compression)
2224                    .map(Arc::from)
2225            })
2226            .transpose()?;
2227        let def_decompressor = layout
2228            .def_compression
2229            .as_ref()
2230            .map(|def_compression| {
2231                decompressors
2232                    .create_block_decompressor(def_compression)
2233                    .map(Arc::from)
2234            })
2235            .transpose()?;
2236        let def_meaning = layout
2237            .layers
2238            .iter()
2239            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
2240            .collect::<Vec<_>>();
2241        let value_decompressor = decompressors.create_miniblock_decompressor(
2242            layout.value_compression.as_ref().unwrap(),
2243            decompressors,
2244        )?;
2245
2246        let dictionary = if let Some(dictionary_encoding) = layout.dictionary.as_ref() {
2247            let num_dictionary_items = layout.num_dictionary_items;
2248            let dictionary_decompressor = decompressors
2249                .create_block_decompressor(dictionary_encoding)?
2250                .into();
2251            let dictionary_data_alignment = match dictionary_encoding.compression.as_ref().unwrap()
2252            {
2253                Compression::Variable(_) => 4,
2254                Compression::Flat(_) => 16,
2255                Compression::General(_) => 1,
2256                Compression::InlineBitpacking(_) | Compression::OutOfLineBitpacking(_) => {
2257                    crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
2258                }
2259                _ => {
2260                    return Err(Error::invalid_input_source(
2261                        format!(
2262                            "Unsupported mini-block dictionary encoding: {:?}",
2263                            dictionary_encoding.compression.as_ref().unwrap()
2264                        )
2265                        .into(),
2266                    ));
2267                }
2268            };
2269            Some(MiniBlockSchedulerDictionary {
2270                dictionary_decompressor,
2271                dictionary_buf_position_and_size: buffer_offsets_and_sizes[2],
2272                dictionary_data_alignment,
2273                num_dictionary_items,
2274            })
2275        } else {
2276            None
2277        };
2278
2279        Ok(Self {
2280            buffer_offsets_and_sizes: buffer_offsets_and_sizes.to_vec(),
2281            rep_decompressor,
2282            def_decompressor,
2283            value_decompressor: value_decompressor.into(),
2284            repetition_index_depth: layout.repetition_index_depth as u16,
2285            num_buffers: layout.num_buffers,
2286            priority,
2287            items_in_page,
2288            dictionary,
2289            def_meaning: def_meaning.into(),
2290            page_meta: None,
2291            has_large_chunk: layout.has_large_chunk,
2292        })
2293    }
2294
2295    fn lookup_chunks(&self, chunk_indices: &[usize]) -> Vec<LoadedChunk> {
2296        let chunk_index = &self.page_meta.as_ref().unwrap().chunk_index;
2297        chunk_indices
2298            .iter()
2299            .map(|&chunk_idx| LoadedChunk {
2300                byte_range: chunk_index.byte_range(chunk_idx),
2301                items_in_chunk: chunk_index.items_in_chunk(chunk_idx),
2302                chunk_idx,
2303                data: LanceBuffer::empty(),
2304            })
2305            .collect()
2306    }
2307}
2308
2309#[derive(Debug, PartialEq, Eq, Clone, Copy)]
2310enum PreambleAction {
2311    Take,
2312    Skip,
2313    Absent,
2314}
2315
2316// When we schedule a chunk we use the repetition index (or, if none exists, just the # of items
2317// in each chunk) to map a user requested range into a set of ChunkInstruction objects which tell
2318// us how exactly to read from the chunk.
2319//
2320// Examples:
2321//
2322// | Chunk 0     | Chunk 1   | Chunk 2   | Chunk 3 |
2323// | xxxxyyyyzzz | zzzzzzzzz | zzzzzzzzz | aaabbcc |
2324//
2325// Full read (0..6)
2326//
2327// Chunk 0: (several rows, ends with trailer)
2328//   preamble: absent
2329//   rows_to_skip: 0
2330//   rows_to_take: 3 (x, y, z)
2331//   take_trailer: true
2332//
2333// Chunk 1: (all preamble, ends with trailer)
2334//   preamble: take
2335//   rows_to_skip: 0
2336//   rows_to_take: 0
2337//   take_trailer: true
2338//
2339// Chunk 2: (all preamble, no trailer)
2340//   preamble: take
2341//   rows_to_skip: 0
2342//   rows_to_take: 0
2343//   take_trailer: false
2344//
2345// Chunk 3: (several rows, no trailer or preamble)
2346//   preamble: absent
2347//   rows_to_skip: 0
2348//   rows_to_take: 3 (a, b, c)
2349//   take_trailer: false
2350#[derive(Clone, Debug, PartialEq, Eq)]
2351struct ChunkInstructions {
2352    // The index of the chunk to read
2353    chunk_idx: usize,
2354    // A "preamble" is when a chunk begins with a continuation of the previous chunk's list.  If there
2355    // is no repetition index there is never a preamble.
2356    //
2357    // It's possible for a chunk to be entirely premable.  For example, if there is a really large list
2358    // that spans several chunks.
2359    preamble: PreambleAction,
2360    // How many complete rows (not including the preamble or trailer) to skip
2361    //
2362    // If this is non-zero then premable must not be Take
2363    rows_to_skip: u64,
2364    // How many rows to take.  If a row splits across chunks then we will count the row in the first
2365    // chunk that contains the row.
2366    rows_to_take: u64,
2367    // A "trailer" is when a chunk ends with a partial list.  If there is no repetition index there is
2368    // never a trailer.
2369    //
2370    // A chunk that is all preamble may or may not have a trailer.
2371    //
2372    // If this is true then we want to include the trailer
2373    take_trailer: bool,
2374}
2375
2376// First, we schedule a bunch of [`ChunkInstructions`] based on the users ranges.  Then we
2377// start decoding them, based on a batch size, which might not align with what we scheduled.
2378//
2379// This results in `ChunkDrainInstructions` which targets a contiguous slice of a `ChunkInstructions`
2380//
2381// So if `ChunkInstructions` is "skip preamble, skip 10, take 50, take trailer" and we are decoding in
2382// batches of size 10 we might have a `ChunkDrainInstructions` that targets that chunk and has its own
2383// skip of 17 and take of 10.  This would mean we decode the chunk, skip the preamble and 27 rows, and
2384// then take 10 rows.
2385//
2386// One very confusing bit is that `rows_to_take` includes the trailer.  So if we have two chunks:
2387//  -no preamble, skip 5, take 10, take trailer
2388//  -take preamble, skip 0, take 50, no trailer
2389//
2390// and we are draining 20 rows then the drain instructions for the first batch will be:
2391//  - no preamble, skip 0 (from chunk 0), take 11 (from chunk 0)
2392//  - take preamble (from chunk 1), skip 0 (from chunk 1), take 9 (from chunk 1)
2393#[derive(Debug, PartialEq, Eq)]
2394struct ChunkDrainInstructions {
2395    chunk_instructions: ChunkInstructions,
2396    rows_to_skip: u64,
2397    rows_to_take: u64,
2398    preamble_action: PreambleAction,
2399}
2400
2401impl ChunkInstructions {
2402    // Given a repetition index and a set of user ranges we need to figure out how to read from the chunks
2403    //
2404    // We assume that `user_ranges` are in sorted order and non-overlapping
2405    //
2406    // The output will be a set of `ChunkInstructions` which tell us how to read from the chunks
2407    fn schedule_instructions(
2408        chunk_index: &MiniBlockChunkIndex,
2409        user_ranges: &[Range<u64>],
2410    ) -> Vec<Self> {
2411        // Bind the per-page chunk count once; re-deriving it each iteration
2412        // costs a width match plus a length read.
2413        let num_chunks = chunk_index.num_chunks();
2414        // This is an in-exact capacity guess but pretty good.  The actual capacity can be
2415        // smaller if instructions are merged.  It can be larger if there are multiple instructions
2416        // per row which can happen with lists.
2417        let mut chunk_instructions = Vec::with_capacity(user_ranges.len());
2418
2419        for user_range in user_ranges {
2420            let mut rows_needed = user_range.end - user_range.start;
2421            let mut need_preamble = false;
2422
2423            // Need to find the first chunk with a first row >= user_range.start.  If there are
2424            // multiple chunks with the same first row we need to take the first one.
2425            let mut block_index = chunk_index.find_chunk(user_range.start);
2426
2427            let mut to_skip = user_range.start - chunk_index.first_row(block_index);
2428
2429            while rows_needed > 0 || need_preamble {
2430                // Check if we've gone past the last block (should not happen)
2431                if block_index >= num_chunks {
2432                    log::warn!(
2433                        "schedule_instructions inconsistency: block_index >= num_chunks, exiting early"
2434                    );
2435                    break;
2436                }
2437
2438                let starts_including_trailer = chunk_index.rows_in_chunk(block_index);
2439                let has_preamble = chunk_index.has_preamble(block_index);
2440                let has_trailer = chunk_index.has_trailer(block_index);
2441                let rows_avail = starts_including_trailer.saturating_sub(to_skip);
2442
2443                // Handle blocks that are entirely preamble (rows_avail = 0)
2444                // These blocks have no rows to take but may have a preamble we need
2445                // We only look for preamble if to_skip == 0 (we're not skipping rows)
2446                if rows_avail == 0 && to_skip == 0 {
2447                    // Only process if this chunk has a preamble we need
2448                    if has_preamble && need_preamble {
2449                        chunk_instructions.push(Self {
2450                            chunk_idx: block_index,
2451                            preamble: PreambleAction::Take,
2452                            rows_to_skip: 0,
2453                            rows_to_take: 0,
2454                            // We still need to look at has_trailer to distinguish between "all preamble
2455                            // and row ends at end of chunk" and "all preamble and row bleeds into next
2456                            // chunk".  Both cases will have 0 rows available.
2457                            take_trailer: has_trailer,
2458                        });
2459                        // Only set need_preamble = false if the chunk has at least one row,
2460                        // Or we are reaching the last block,
2461                        // Otherwise, the chunk is entirely preamble and we need the next chunk's preamble too
2462                        if starts_including_trailer > 0 || block_index == num_chunks - 1 {
2463                            need_preamble = false;
2464                        }
2465                    }
2466                    // Move to next block
2467                    block_index += 1;
2468                    continue;
2469                }
2470
2471                // Edge case: if rows_avail == 0 but to_skip > 0
2472                // This theoretically shouldn't happen (binary search should avoid it)
2473                // but handle it for safety
2474                if rows_avail == 0 && to_skip > 0 {
2475                    // This block doesn't have enough rows to skip, move to next block
2476                    // Adjust to_skip by the number of rows in this block
2477                    to_skip -= starts_including_trailer;
2478                    block_index += 1;
2479                    continue;
2480                }
2481
2482                let rows_to_take = rows_avail.min(rows_needed);
2483                rows_needed -= rows_to_take;
2484
2485                let mut take_trailer = false;
2486                let preamble = if has_preamble {
2487                    if need_preamble {
2488                        PreambleAction::Take
2489                    } else {
2490                        PreambleAction::Skip
2491                    }
2492                } else {
2493                    PreambleAction::Absent
2494                };
2495
2496                // Are we taking the trailer?  If so, make sure we mark that we need the preamble
2497                if rows_to_take == rows_avail && has_trailer {
2498                    take_trailer = true;
2499                    need_preamble = true;
2500                } else {
2501                    need_preamble = false;
2502                };
2503
2504                chunk_instructions.push(Self {
2505                    preamble,
2506                    chunk_idx: block_index,
2507                    rows_to_skip: to_skip,
2508                    rows_to_take,
2509                    take_trailer,
2510                });
2511
2512                to_skip = 0;
2513                block_index += 1;
2514            }
2515        }
2516
2517        // If there were multiple ranges we may have multiple instructions for a single chunk.  Merge them now if they
2518        // are _adjacent_ (i.e. don't merge "take first row of chunk 0" and "take third row of chunk 0" into "take 2
2519        // rows of chunk 0 starting at 0")
2520        if user_ranges.len() > 1 {
2521            // Merge adjacent instructions in place.  `write` indexes the last
2522            // retained instruction; each following instruction is either folded
2523            // into it (contiguous within the same chunk) or compacted forward.
2524            let mut write = 0;
2525            for read in 1..chunk_instructions.len() {
2526                let merges = {
2527                    let last = &chunk_instructions[write];
2528                    let candidate = &chunk_instructions[read];
2529                    last.chunk_idx == candidate.chunk_idx
2530                        && last.rows_to_take + last.rows_to_skip == candidate.rows_to_skip
2531                };
2532                if merges {
2533                    let rows_to_take = chunk_instructions[read].rows_to_take;
2534                    let take_trailer = chunk_instructions[read].take_trailer;
2535                    let last = &mut chunk_instructions[write];
2536                    last.rows_to_take += rows_to_take;
2537                    last.take_trailer |= take_trailer;
2538                } else {
2539                    write += 1;
2540                    if write != read {
2541                        chunk_instructions.swap(write, read);
2542                    }
2543                }
2544            }
2545            chunk_instructions.truncate(write + 1);
2546        }
2547        chunk_instructions
2548    }
2549
2550    fn drain_from_instruction(
2551        &self,
2552        rows_desired: &mut u64,
2553        need_preamble: &mut bool,
2554        skip_in_chunk: &mut u64,
2555    ) -> (ChunkDrainInstructions, bool) {
2556        // If we need the premable then we shouldn't be skipping anything
2557        debug_assert!(!*need_preamble || *skip_in_chunk == 0);
2558        let rows_avail = self.rows_to_take - *skip_in_chunk;
2559        let has_preamble = self.preamble != PreambleAction::Absent;
2560        let preamble_action = match (*need_preamble, has_preamble) {
2561            (true, true) => PreambleAction::Take,
2562            (true, false) => panic!("Need preamble but there isn't one"),
2563            (false, true) => PreambleAction::Skip,
2564            (false, false) => PreambleAction::Absent,
2565        };
2566
2567        // How many rows are we actually taking in this take step (including the preamble
2568        // and trailer both as individual rows)
2569        let rows_taking = if *rows_desired >= rows_avail {
2570            // We want all the rows.  If there is a trailer we are grabbing it and will need
2571            // the preamble of the next chunk
2572            // If there is a trailer and we are taking all the rows then we need the preamble
2573            // of the next chunk.
2574            //
2575            // Also, if this chunk is entirely preamble (rows_avail == 0 && !take_trailer) then we
2576            // need the preamble of the next chunk.
2577            *need_preamble = self.take_trailer;
2578            rows_avail
2579        } else {
2580            // We aren't taking all the rows.  Even if there is a trailer we aren't taking
2581            // it so we will not need the preamble
2582            *need_preamble = false;
2583            *rows_desired
2584        };
2585        let rows_skipped = *skip_in_chunk;
2586
2587        // Update the state for the next iteration
2588        let consumed_chunk = if *rows_desired >= rows_avail {
2589            *rows_desired -= rows_avail;
2590            *skip_in_chunk = 0;
2591            true
2592        } else {
2593            *skip_in_chunk += *rows_desired;
2594            *rows_desired = 0;
2595            false
2596        };
2597
2598        (
2599            ChunkDrainInstructions {
2600                chunk_instructions: self.clone(),
2601                rows_to_skip: rows_skipped,
2602                rows_to_take: rows_taking,
2603                preamble_action,
2604            },
2605            consumed_chunk,
2606        )
2607    }
2608}
2609
2610enum Words {
2611    U16(ScalarBuffer<u16>),
2612    U32(ScalarBuffer<u32>),
2613}
2614
2615struct WordsIter<'a> {
2616    iter: Box<dyn Iterator<Item = u32> + 'a>,
2617}
2618
2619impl Words {
2620    pub fn len(&self) -> usize {
2621        match self {
2622            Self::U16(b) => b.len(),
2623            Self::U32(b) => b.len(),
2624        }
2625    }
2626
2627    pub fn iter(&self) -> WordsIter<'_> {
2628        match self {
2629            Self::U16(buf) => WordsIter {
2630                iter: Box::new(buf.iter().map(|&x| x as u32)),
2631            },
2632            Self::U32(buf) => WordsIter {
2633                iter: Box::new(buf.iter().copied()),
2634            },
2635        }
2636    }
2637
2638    pub fn from_bytes(bytes: Bytes, has_large_chunk: bool) -> Result<Self> {
2639        let bytes_per_value = if has_large_chunk { 4 } else { 2 };
2640        assert_eq!(bytes.len() % bytes_per_value, 0);
2641        let buffer = LanceBuffer::from_bytes(bytes, bytes_per_value as u64);
2642        if has_large_chunk {
2643            Ok(Self::U32(buffer.borrow_to_typed_slice::<u32>()))
2644        } else {
2645            Ok(Self::U16(buffer.borrow_to_typed_slice::<u16>()))
2646        }
2647    }
2648}
2649
2650impl<'a> Iterator for WordsIter<'a> {
2651    type Item = u32;
2652
2653    fn next(&mut self) -> Option<Self::Item> {
2654        self.iter.next()
2655    }
2656}
2657
2658/// Per-chunk leaf value-count analysis derived from the metadata words.
2659///
2660/// `values_per_chunk` is the count shared by every non-last chunk (meaningful
2661/// when `uniform`), and `last_chunk_values` is the final chunk's count.
2662struct FlatValueCounts {
2663    logs: Vec<u8>,
2664    uniform: bool,
2665    values_per_chunk: u64,
2666    last_chunk_values: u64,
2667}
2668
2669fn analyze_value_counts(words: &Words, items_in_page: u64) -> Result<FlatValueCounts> {
2670    let num_chunks = words.len();
2671    let logs = words.iter().map(|w| (w & 0x0F) as u8).collect::<Vec<_>>();
2672    let mut counted = 0u64;
2673    for (chunk_index, &log) in logs.iter().take(num_chunks.saturating_sub(1)).enumerate() {
2674        if log == 0 {
2675            return Err(Error::corrupt_file_named(
2676                "miniblock_metadata",
2677                format!(
2678                    "non-final chunk {chunk_index} of {num_chunks} has invalid log_num_values=0"
2679                ),
2680            ));
2681        }
2682        counted = counted.checked_add(1u64 << log).ok_or_else(|| {
2683            Error::corrupt_file_named(
2684                "miniblock_metadata",
2685                format!(
2686                    "value count overflow at chunk {chunk_index}: counted_values={counted}, \
2687                     log_num_values={log}, items_in_page={items_in_page}"
2688                ),
2689            )
2690        })?;
2691    }
2692    let last_chunk_values = items_in_page.checked_sub(counted).ok_or_else(|| {
2693        Error::corrupt_file_named(
2694            "miniblock_metadata",
2695            format!(
2696                "non-final chunks account for counted_values={counted}, exceeding \
2697                 items_in_page={items_in_page}"
2698            ),
2699        )
2700    })?;
2701    if let Some(&last_log) = logs.last()
2702        && last_log != 0
2703        && (1u64 << last_log) != last_chunk_values
2704    {
2705        return Err(Error::corrupt_file_named(
2706            "miniblock_metadata",
2707            format!(
2708                "final chunk log_num_values={last_log} does not match \
2709                 last_chunk_values={last_chunk_values}: counted_values={counted}, \
2710                 items_in_page={items_in_page}"
2711            ),
2712        ));
2713    }
2714    let uniform = num_chunks <= 1 || logs[..num_chunks - 1].iter().all(|&log| log == logs[0]);
2715    // A single-chunk page has no "non-last" chunk to derive a stride from; use the
2716    // page item count (min 1 so it stays a valid divisor in `find_chunk`).
2717    let values_per_chunk = if num_chunks <= 1 {
2718        items_in_page.max(1)
2719    } else {
2720        1u64 << logs[0]
2721    };
2722    Ok(FlatValueCounts {
2723        logs,
2724        uniform,
2725        values_per_chunk,
2726        last_chunk_values,
2727    })
2728}
2729
2730/// Iterator over per-chunk value counts for a non-uniform flat page.  Non-last
2731/// chunks yield `1 << log`; the last yields the validated remaining item count.
2732fn flat_value_counts_iter(logs: &[u8], last_chunk_values: u64) -> impl Iterator<Item = u64> + '_ {
2733    let num_chunks = logs.len();
2734    (0..num_chunks).map(move |i| {
2735        if i + 1 < num_chunks {
2736            1u64 << logs[i]
2737        } else {
2738            last_chunk_values
2739        }
2740    })
2741}
2742
2743/// Builds the compact per-chunk index from the metadata words and, for nested
2744/// pages, the raw repetition-index bytes.  The row axis is picked by page shape:
2745/// `UniformFlat` when all non-last chunks share a value count (fixed-width /
2746/// bitpacking), `Flat` for non-uniform flat pages (RLE / FSST), else `Nested`.
2747fn build_chunk_index(
2748    words: &Words,
2749    items_in_page: u64,
2750    base: u64,
2751    data_buf_size: u64,
2752    rep_index_bytes: Option<&[u8]>,
2753    repetition_index_depth: u16,
2754) -> Result<MiniBlockChunkIndex> {
2755    let num_chunks = words.len();
2756    // Validate item counts before byte sizes because both share a metadata word,
2757    // and an invalid count must not reach the final-chunk subtraction.
2758    let value_counts = analyze_value_counts(words, items_in_page)?;
2759
2760    // Each chunk stores `(divided_bytes + 1) * MINIBLOCK_ALIGNMENT` bytes, so the
2761    // deltas are the chunk sizes and their grand total is the data buffer size.
2762    let byte_starts = PrefixSums::from_deltas(
2763        words
2764            .iter()
2765            .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64),
2766        num_chunks,
2767        data_buf_size,
2768    );
2769
2770    // Nested pages track rows via the repetition index and keep leaf item counts
2771    // separately; flat pages have row == value index, so value counts are rows.
2772    let rows = if let Some(rep_index_data) = rep_index_bytes {
2773        assert!(rep_index_data.len() % 8 == 0);
2774        let stride = repetition_index_depth as usize + 1;
2775        let (row_starts, has_trailer) = parse_nested_rep(rep_index_data, stride);
2776        let item_counts = if value_counts.uniform {
2777            ItemCounts::Uniform {
2778                values_per_chunk: value_counts.values_per_chunk,
2779                last_chunk_values: value_counts.last_chunk_values,
2780            }
2781        } else {
2782            ItemCounts::PerChunkLog {
2783                logs: value_counts.logs,
2784                last_chunk_values: value_counts.last_chunk_values,
2785            }
2786        };
2787        RowMapping::Nested {
2788            row_starts,
2789            has_trailer,
2790            item_counts,
2791        }
2792    } else {
2793        if value_counts.uniform {
2794            RowMapping::UniformFlat {
2795                values_per_chunk: value_counts.values_per_chunk,
2796                last_chunk_values: value_counts.last_chunk_values,
2797                num_chunks,
2798            }
2799        } else {
2800            let value_starts = PrefixSums::from_deltas(
2801                flat_value_counts_iter(&value_counts.logs, value_counts.last_chunk_values),
2802                num_chunks,
2803                items_in_page,
2804            );
2805            RowMapping::Flat { value_starts }
2806        }
2807    };
2808
2809    Ok(MiniBlockChunkIndex::new(base, byte_starts, rows))
2810}
2811
2812impl StructuralPageScheduler for MiniBlockScheduler {
2813    fn init_layout(&self) -> Result<PageInitialization> {
2814        // We always need to fetch chunk metadata.  We may also need to fetch a dictionary and
2815        // we may also need to fetch the repetition index.  Here, we gather what buffers we
2816        // need.
2817        let (meta_buf_position, meta_buf_size) = self.buffer_offsets_and_sizes[0];
2818        let dictionary_range = self.dictionary.as_ref().map(|dictionary| {
2819            dictionary.dictionary_buf_position_and_size.0
2820                ..dictionary.dictionary_buf_position_and_size.0
2821                    + dictionary.dictionary_buf_position_and_size.1
2822        });
2823        let repetition_index_range = if self.repetition_index_depth > 0 {
2824            let (rep_index_pos, rep_index_size) =
2825                self.buffer_offsets_and_sizes.last().ok_or_else(|| {
2826                    Error::invalid_input_source(
2827                        "Mini-block page has a repetition index but no buffers".into(),
2828                    )
2829                })?;
2830            Some(*rep_index_pos..*rep_index_pos + *rep_index_size)
2831        } else {
2832            None
2833        };
2834        Ok(PageInitialization::new([
2835            Some(meta_buf_position..meta_buf_position + meta_buf_size),
2836            dictionary_range,
2837            repetition_index_range,
2838        ]))
2839    }
2840
2841    fn init_from_buffers<'a>(
2842        &'a mut self,
2843        buffers: PageInitializationBuffers,
2844        _io: &Arc<dyn EncodingsIo>,
2845    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
2846        let base = self.buffer_offsets_and_sizes[1].0;
2847        let data_buf_size = self.buffer_offsets_and_sizes[1].1;
2848
2849        async move {
2850            let [meta_bytes, dictionary_bytes, rep_index_bytes] =
2851                buffers.try_into_array("mini-block")?;
2852            let meta_bytes = meta_bytes.ok_or_else(|| {
2853                Error::internal("Mini-block page initialization is missing chunk metadata")
2854            })?;
2855
2856            let words = Words::from_bytes(meta_bytes, self.has_large_chunk)?;
2857            let chunk_index = build_chunk_index(
2858                &words,
2859                self.items_in_page,
2860                base,
2861                data_buf_size,
2862                rep_index_bytes.as_deref(),
2863                self.repetition_index_depth,
2864            )?;
2865
2866            // decode dictionary
2867            let dictionary = match (self.dictionary.as_ref(), dictionary_bytes) {
2868                (Some(dictionary), Some(dictionary_data)) => {
2869                    Some(Arc::new(dictionary.dictionary_decompressor.decompress(
2870                        Some(LanceBuffer::from_bytes(
2871                            dictionary_data,
2872                            dictionary.dictionary_data_alignment,
2873                        )),
2874                        dictionary.num_dictionary_items,
2875                    )?))
2876                }
2877                (None, None) => None,
2878                _ => {
2879                    return Err(Error::internal(
2880                        "Mini-block page initialization dictionary slot does not match its layout",
2881                    ));
2882                }
2883            };
2884
2885            let page_meta = Arc::new(MiniBlockCacheableState {
2886                chunk_index,
2887                dictionary,
2888            });
2889            Ok(page_meta as Arc<dyn CachedPageData>)
2890        }
2891        .boxed()
2892    }
2893
2894    fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> Result<()> {
2895        self.page_meta = Some(data.clone().as_arc_any().downcast().map_err(|_| {
2896            Error::invalid_input_source("Cached mini-block page data has an unexpected type".into())
2897        })?);
2898        Ok(())
2899    }
2900
2901    fn schedule_ranges(
2902        &self,
2903        ranges: &[Range<u64>],
2904        io: &Arc<dyn EncodingsIo>,
2905    ) -> Result<Vec<PageLoadTask>> {
2906        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
2907
2908        let page_meta = self.page_meta.as_ref().unwrap();
2909
2910        let chunk_instructions =
2911            ChunkInstructions::schedule_instructions(&page_meta.chunk_index, ranges);
2912
2913        debug_assert_eq!(
2914            num_rows,
2915            chunk_instructions
2916                .iter()
2917                .map(|ci| ci.rows_to_take)
2918                .sum::<u64>()
2919        );
2920
2921        let chunks_needed = chunk_instructions
2922            .iter()
2923            .map(|ci| ci.chunk_idx)
2924            .unique()
2925            .collect::<Vec<_>>();
2926
2927        let mut loaded_chunks = self.lookup_chunks(&chunks_needed);
2928        let chunk_ranges = loaded_chunks
2929            .iter()
2930            .map(|c| c.byte_range.clone())
2931            .collect::<Vec<_>>();
2932        let loaded_chunk_data = io.submit_request(chunk_ranges, self.priority);
2933
2934        let rep_decompressor = self.rep_decompressor.clone();
2935        let def_decompressor = self.def_decompressor.clone();
2936        let value_decompressor = self.value_decompressor.clone();
2937        let num_buffers = self.num_buffers;
2938        let has_large_chunk = self.has_large_chunk;
2939        let dictionary = page_meta
2940            .dictionary
2941            .as_ref()
2942            .map(|dictionary| dictionary.clone());
2943        let def_meaning = self.def_meaning.clone();
2944
2945        let res = async move {
2946            let loaded_chunk_data = loaded_chunk_data.await?;
2947            for (loaded_chunk, chunk_data) in loaded_chunks.iter_mut().zip(loaded_chunk_data) {
2948                loaded_chunk.data = LanceBuffer::from_bytes(chunk_data, 1);
2949            }
2950
2951            Ok(Box::new(MiniBlockDecoder {
2952                rep_decompressor,
2953                def_decompressor,
2954                value_decompressor,
2955                def_meaning,
2956                loaded_chunks: VecDeque::from_iter(loaded_chunks),
2957                instructions: VecDeque::from(chunk_instructions),
2958                offset_in_current_chunk: 0,
2959                dictionary,
2960                num_rows,
2961                num_buffers,
2962                has_large_chunk,
2963            }) as Box<dyn StructuralPageDecoder>)
2964        }
2965        .boxed();
2966        let page_load_task = PageLoadTask {
2967            decoder_fut: res,
2968            num_rows,
2969        };
2970        Ok(vec![page_load_task])
2971    }
2972}
2973
2974#[derive(Debug, Clone, Copy)]
2975struct FullZipRepIndexDetails {
2976    buf_position: u64,
2977    bytes_per_value: u64, // Will be 1, 2, 4, or 8
2978}
2979
2980#[derive(Debug)]
2981enum PerValueDecompressor {
2982    Fixed(Arc<dyn FixedPerValueDecompressor>),
2983    Variable(Arc<dyn VariablePerValueDecompressor>),
2984}
2985
2986#[derive(Debug)]
2987struct FullZipDecodeDetails {
2988    value_decompressor: PerValueDecompressor,
2989    def_meaning: Arc<[DefinitionInterpretation]>,
2990    ctrl_word_parser: ControlWordParser,
2991    max_rep: u16,
2992    max_visible_def: u16,
2993}
2994
2995/// Describes where FullZip byte ranges should be read from.
2996///
2997/// FullZip decoding always needs a list of byte ranges, but those bytes can come
2998/// from two different places:
2999/// - Remote I/O (normal path): ranges are fetched from the underlying `EncodingsIo`.
3000/// - A prefetched full page (full scan fast path): the entire page has already been
3001///   loaded once and ranges should be sliced from memory.
3002///
3003/// This abstraction keeps scheduling code focused on "which ranges are needed"
3004/// instead of "how bytes are fetched", and it lets full-page scans avoid the
3005/// two-stage rep-index -> data I/O pipeline.
3006#[derive(Debug, Clone)]
3007enum FullZipReadSource {
3008    /// Fetch ranges from the storage backend through the encoding I/O interface.
3009    Remote(Arc<dyn EncodingsIo>),
3010    /// Slice ranges from an already-loaded FullZip page buffer.
3011    PrefetchedPage { base_offset: u64, data: LanceBuffer },
3012}
3013
3014impl FullZipReadSource {
3015    /// Materialize the requested ranges as decode-ready `LanceBuffer`s.
3016    ///
3017    /// The returned buffers preserve the input range order.
3018    fn fetch(
3019        &self,
3020        ranges: &[Range<u64>],
3021        priority: u64,
3022    ) -> BoxFuture<'static, Result<VecDeque<LanceBuffer>>> {
3023        match self {
3024            Self::Remote(io) => {
3025                let io = io.clone();
3026                let ranges = ranges.to_vec();
3027                async move {
3028                    let data = io.submit_request(ranges, priority).await?;
3029                    Ok(data
3030                        .into_iter()
3031                        .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
3032                        .collect::<VecDeque<_>>())
3033                }
3034                .boxed()
3035            }
3036            Self::PrefetchedPage { base_offset, data } => {
3037                let base_offset = *base_offset;
3038                let data = data.clone();
3039                let page_end = base_offset + data.len() as u64;
3040                std::future::ready(
3041                    ranges
3042                        .iter()
3043                        .map(|range| {
3044                            if range.start > range.end
3045                                || range.start < base_offset
3046                                || range.end > page_end
3047                            {
3048                                return Err(Error::internal(format!(
3049                                    "Requested range {:?} is outside page range {}..{}",
3050                                    range, base_offset, page_end
3051                                )));
3052                            }
3053                            let start = (range.start - base_offset) as usize;
3054                            let len = (range.end - range.start) as usize;
3055                            Ok(data.slice_with_length(start, len))
3056                        })
3057                        .collect::<Result<VecDeque<_>>>(),
3058                )
3059                .boxed()
3060            }
3061        }
3062    }
3063}
3064
3065/// A scheduler for full-zip encoded data
3066///
3067/// When the data type has a fixed-width then we simply need to map from
3068/// row ranges to byte ranges using the fixed-width of the data type.
3069///
3070/// When the data type is variable-width or has any repetition then a
3071/// repetition index is required.
3072#[derive(Debug)]
3073pub struct FullZipScheduler {
3074    data_buf_position: u64,
3075    data_buf_size: u64,
3076    rep_index: Option<FullZipRepIndexDetails>,
3077    priority: u64,
3078    rows_in_page: u64,
3079    bits_per_offset: u8,
3080    details: Arc<FullZipDecodeDetails>,
3081    /// Cached state containing the decoded repetition index
3082    cached_state: Option<Arc<FullZipCacheableState>>,
3083    /// Whether repetition index metadata should be cached during initialize.
3084    enable_cache: bool,
3085}
3086
3087impl FullZipScheduler {
3088    fn try_new(
3089        buffer_offsets_and_sizes: &[(u64, u64)],
3090        priority: u64,
3091        rows_in_page: u64,
3092        layout: &pb21::FullZipLayout,
3093        decompressors: &dyn DecompressionStrategy,
3094    ) -> Result<Self> {
3095        let (data_buf_position, data_buf_size) = buffer_offsets_and_sizes[0];
3096        let rep_index = buffer_offsets_and_sizes.get(1).map(|(pos, len)| {
3097            let num_reps = rows_in_page + 1;
3098            let bytes_per_rep = len / num_reps;
3099            debug_assert_eq!(len % num_reps, 0);
3100            debug_assert!(
3101                bytes_per_rep == 1
3102                    || bytes_per_rep == 2
3103                    || bytes_per_rep == 4
3104                    || bytes_per_rep == 8
3105            );
3106            FullZipRepIndexDetails {
3107                buf_position: *pos,
3108                bytes_per_value: bytes_per_rep,
3109            }
3110        });
3111
3112        let value_decompressor = match layout.details {
3113            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => {
3114                let decompressor = decompressors.create_fixed_per_value_decompressor(
3115                    layout.value_compression.as_ref().unwrap(),
3116                )?;
3117                PerValueDecompressor::Fixed(decompressor.into())
3118            }
3119            Some(pb21::full_zip_layout::Details::BitsPerOffset(_)) => {
3120                let decompressor = decompressors.create_variable_per_value_decompressor(
3121                    layout.value_compression.as_ref().unwrap(),
3122                )?;
3123                PerValueDecompressor::Variable(decompressor.into())
3124            }
3125            None => {
3126                panic!("Full-zip layout must have a `details` field");
3127            }
3128        };
3129        let ctrl_word_parser = ControlWordParser::new(
3130            layout.bits_rep.try_into().unwrap(),
3131            layout.bits_def.try_into().unwrap(),
3132        );
3133        let def_meaning = layout
3134            .layers
3135            .iter()
3136            .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
3137            .collect::<Vec<_>>();
3138
3139        let max_rep = def_meaning.iter().filter(|d| d.is_list()).count() as u16;
3140        let max_visible_def = max_visible_level(&def_meaning);
3141
3142        let bits_per_offset = match layout.details {
3143            Some(pb21::full_zip_layout::Details::BitsPerValue(_)) => 32,
3144            Some(pb21::full_zip_layout::Details::BitsPerOffset(bits_per_offset)) => {
3145                bits_per_offset as u8
3146            }
3147            None => panic!("Full-zip layout must have a `details` field"),
3148        };
3149
3150        let details = Arc::new(FullZipDecodeDetails {
3151            value_decompressor,
3152            def_meaning: def_meaning.into(),
3153            ctrl_word_parser,
3154            max_rep,
3155            max_visible_def,
3156        });
3157        Ok(Self {
3158            data_buf_position,
3159            data_buf_size,
3160            rep_index,
3161            details,
3162            priority,
3163            rows_in_page,
3164            bits_per_offset,
3165            cached_state: None,
3166            enable_cache: false,
3167        })
3168    }
3169
3170    fn covers_entire_page(ranges: &[Range<u64>], rows_in_page: u64) -> bool {
3171        if ranges.is_empty() {
3172            return false;
3173        }
3174        let mut expected_start = 0;
3175        for range in ranges {
3176            if range.start != expected_start || range.end > rows_in_page || range.end < range.start
3177            {
3178                return false;
3179            }
3180            expected_start = range.end;
3181        }
3182        expected_start == rows_in_page
3183    }
3184
3185    fn create_page_load_task(
3186        io_future: BoxFuture<'static, Result<Vec<Bytes>>>,
3187        num_rows: u64,
3188        details: Arc<FullZipDecodeDetails>,
3189        bits_per_offset: u8,
3190    ) -> PageLoadTask {
3191        let load_task = async move {
3192            let buffers = io_future.await?;
3193            let data = buffers
3194                .into_iter()
3195                .map(|bytes| LanceBuffer::from_bytes(bytes, 1))
3196                .collect::<VecDeque<_>>();
3197            Self::create_decoder(details, data, num_rows, bits_per_offset)
3198        }
3199        .boxed();
3200        PageLoadTask {
3201            decoder_fut: load_task,
3202            num_rows,
3203        }
3204    }
3205
3206    /// Creates a decoder from the loaded data
3207    fn create_decoder(
3208        details: Arc<FullZipDecodeDetails>,
3209        data: VecDeque<LanceBuffer>,
3210        num_rows: u64,
3211        bits_per_offset: u8,
3212    ) -> Result<Box<dyn StructuralPageDecoder>> {
3213        match &details.value_decompressor {
3214            PerValueDecompressor::Fixed(decompressor) => {
3215                let bits_per_value = decompressor.bits_per_value();
3216                if bits_per_value % 8 != 0 {
3217                    return Err(lance_core::Error::not_supported_source("Bit-packed full-zip encoding (non-byte-aligned values) is not yet implemented".into()));
3218                }
3219                let bytes_per_value = bits_per_value / 8;
3220                let total_bytes_per_value =
3221                    bytes_per_value as usize + details.ctrl_word_parser.bytes_per_word();
3222                // total_bytes_per_value == 0 is valid for constant-null FSL pages written by
3223                // earlier encoders that produced bits_per_value=0 with no ctrl-word bytes.
3224                // FixedFullZipDecoder::drain handles this case by producing AllNull output
3225                // without touching the (empty) data buffer.
3226                Ok(Box::new(FixedFullZipDecoder {
3227                    details,
3228                    data,
3229                    num_rows,
3230                    offset_in_current: 0,
3231                    bytes_per_value: bytes_per_value as usize,
3232                    total_bytes_per_value,
3233                }) as Box<dyn StructuralPageDecoder>)
3234            }
3235            PerValueDecompressor::Variable(_decompressor) => {
3236                Ok(Box::new(VariableFullZipDecoder::new(
3237                    details,
3238                    data,
3239                    num_rows,
3240                    bits_per_offset,
3241                    bits_per_offset,
3242                )?))
3243            }
3244        }
3245    }
3246
3247    /// Extracts byte ranges from a repetition index buffer
3248    /// The buffer contains pairs of (start, end) values for each range
3249    fn extract_byte_ranges_from_pairs(
3250        buffer: LanceBuffer,
3251        bytes_per_value: u64,
3252        data_buf_position: u64,
3253    ) -> Vec<Range<u64>> {
3254        ByteUnpacker::new(buffer, bytes_per_value as usize)
3255            .chunks(2)
3256            .into_iter()
3257            .map(|mut c| {
3258                let start = c.next().unwrap() + data_buf_position;
3259                let end = c.next().unwrap() + data_buf_position;
3260                start..end
3261            })
3262            .collect::<Vec<_>>()
3263    }
3264
3265    /// Extracts byte ranges from a cached repetition index buffer
3266    /// The buffer contains all values and we need to extract specific ranges
3267    fn extract_byte_ranges_from_cached(
3268        buffer: &LanceBuffer,
3269        ranges: &[Range<u64>],
3270        bytes_per_value: u64,
3271        data_buf_position: u64,
3272    ) -> Vec<Range<u64>> {
3273        ranges
3274            .iter()
3275            .map(|r| {
3276                let start_offset = (r.start * bytes_per_value) as usize;
3277                let end_offset = (r.end * bytes_per_value) as usize;
3278
3279                let start_slice = &buffer[start_offset..start_offset + bytes_per_value as usize];
3280                let start_val =
3281                    ByteUnpacker::new(start_slice.iter().copied(), bytes_per_value as usize)
3282                        .next()
3283                        .unwrap();
3284
3285                let end_slice = &buffer[end_offset..end_offset + bytes_per_value as usize];
3286                let end_val =
3287                    ByteUnpacker::new(end_slice.iter().copied(), bytes_per_value as usize)
3288                        .next()
3289                        .unwrap();
3290
3291                (data_buf_position + start_val)..(data_buf_position + end_val)
3292            })
3293            .collect()
3294    }
3295
3296    /// Computes the ranges in the repetition index that need to be loaded
3297    fn compute_rep_index_ranges(
3298        ranges: &[Range<u64>],
3299        rep_index: &FullZipRepIndexDetails,
3300    ) -> Vec<Range<u64>> {
3301        ranges
3302            .iter()
3303            .flat_map(|r| {
3304                let first_val_start =
3305                    rep_index.buf_position + (r.start * rep_index.bytes_per_value);
3306                let first_val_end = first_val_start + rep_index.bytes_per_value;
3307                let last_val_start = rep_index.buf_position + (r.end * rep_index.bytes_per_value);
3308                let last_val_end = last_val_start + rep_index.bytes_per_value;
3309                [first_val_start..first_val_end, last_val_start..last_val_end]
3310            })
3311            .collect()
3312    }
3313
3314    /// Schedules ranges in the presence of a repetition index
3315    fn schedule_ranges_rep(
3316        &self,
3317        ranges: &[Range<u64>],
3318        io: &Arc<dyn EncodingsIo>,
3319        rep_index: FullZipRepIndexDetails,
3320    ) -> Result<Vec<PageLoadTask>> {
3321        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3322        let data_buf_position = self.data_buf_position;
3323        let priority = self.priority;
3324        let details = self.details.clone();
3325        let bits_per_offset = self.bits_per_offset;
3326
3327        if Self::covers_entire_page(ranges, self.rows_in_page) {
3328            let full_range = self.data_buf_position..(self.data_buf_position + self.data_buf_size);
3329            let page_data = io.submit_single(full_range.clone(), priority);
3330            let load_task = async move {
3331                let page_data = page_data.await?;
3332                let source = FullZipReadSource::PrefetchedPage {
3333                    base_offset: full_range.start,
3334                    data: LanceBuffer::from_bytes(page_data, 1),
3335                };
3336                let read_ranges = vec![full_range];
3337                let data = source.fetch(&read_ranges, priority).await?;
3338                Self::create_decoder(details, data, num_rows, bits_per_offset)
3339            }
3340            .boxed();
3341            let page_load_task = PageLoadTask {
3342                decoder_fut: load_task,
3343                num_rows,
3344            };
3345            return Ok(vec![page_load_task]);
3346        }
3347
3348        if let Some(cached_state) = &self.cached_state {
3349            let byte_ranges = Self::extract_byte_ranges_from_cached(
3350                &cached_state.rep_index_buffer,
3351                ranges,
3352                rep_index.bytes_per_value,
3353                data_buf_position,
3354            );
3355            let io_future = io.submit_request(byte_ranges, priority);
3356            let page_load_task =
3357                Self::create_page_load_task(io_future, num_rows, details, bits_per_offset);
3358            return Ok(vec![page_load_task]);
3359        }
3360
3361        let rep_ranges = Self::compute_rep_index_ranges(ranges, &rep_index);
3362        let rep_data = io.submit_request(rep_ranges, priority);
3363        let io_clone = io.clone();
3364        let load_task = async move {
3365            let rep_data = rep_data.await?;
3366            let rep_buffer = LanceBuffer::concat(
3367                &rep_data
3368                    .into_iter()
3369                    .map(|d| LanceBuffer::from_bytes(d, 1))
3370                    .collect::<Vec<_>>(),
3371            );
3372            let byte_ranges = Self::extract_byte_ranges_from_pairs(
3373                rep_buffer,
3374                rep_index.bytes_per_value,
3375                data_buf_position,
3376            );
3377            let source = FullZipReadSource::Remote(io_clone);
3378            let data = source.fetch(&byte_ranges, priority).await?;
3379            Self::create_decoder(details, data, num_rows, bits_per_offset)
3380        }
3381        .boxed();
3382        let page_load_task = PageLoadTask {
3383            decoder_fut: load_task,
3384            num_rows,
3385        };
3386        Ok(vec![page_load_task])
3387    }
3388
3389    // In the simple case there is no repetition and we just have large fixed-width
3390    // rows of data.  We can just map row ranges to byte ranges directly using the
3391    // fixed-width of the data type.
3392    fn schedule_ranges_simple(
3393        &self,
3394        ranges: &[Range<u64>],
3395        io: &Arc<dyn EncodingsIo>,
3396    ) -> Result<Vec<PageLoadTask>> {
3397        // Convert row ranges to item ranges (i.e. multiply by items per row)
3398        let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
3399
3400        let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
3401            unreachable!()
3402        };
3403
3404        // Convert item ranges to byte ranges (i.e. multiply by bytes per item)
3405        let bits_per_value = decompressor.bits_per_value();
3406        if !bits_per_value.is_multiple_of(8) {
3407            return Err(Error::invalid_input_source(
3408                format!(
3409                    "Full-zip fixed-width values must be byte aligned, got {} bits per value",
3410                    bits_per_value
3411                )
3412                .into(),
3413            ));
3414        }
3415        let bytes_per_value = bits_per_value / 8;
3416        let bytes_per_cw = self.details.ctrl_word_parser.bytes_per_word();
3417        let total_bytes_per_value = bytes_per_value + bytes_per_cw as u64;
3418        let byte_ranges = ranges
3419            .iter()
3420            .map(|r| {
3421                debug_assert!(r.end <= self.rows_in_page);
3422                let start = self.data_buf_position + r.start * total_bytes_per_value;
3423                let end = self.data_buf_position + r.end * total_bytes_per_value;
3424                start..end
3425            })
3426            .collect::<Vec<_>>();
3427
3428        let io_future = io.submit_request(byte_ranges, self.priority);
3429        let page_load_task = Self::create_page_load_task(
3430            io_future,
3431            num_rows,
3432            self.details.clone(),
3433            self.bits_per_offset,
3434        );
3435        Ok(vec![page_load_task])
3436    }
3437}
3438
3439/// Cacheable state for FullZip encoding, storing the decoded repetition index
3440#[derive(Debug)]
3441struct FullZipCacheableState {
3442    /// The raw repetition index buffer for future decoding
3443    rep_index_buffer: LanceBuffer,
3444}
3445
3446impl DeepSizeOf for FullZipCacheableState {
3447    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
3448        self.rep_index_buffer.len()
3449    }
3450}
3451
3452impl CachedPageData for FullZipCacheableState {
3453    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
3454        self
3455    }
3456}
3457
3458impl StructuralPageScheduler for FullZipScheduler {
3459    fn needs_initialization(&self) -> bool {
3460        self.enable_cache && self.rep_index.is_some()
3461    }
3462
3463    fn init_layout(&self) -> Result<PageInitialization> {
3464        let repetition_index_range = if self.enable_cache
3465            && let Some(rep_index) = self.rep_index
3466        {
3467            let total_size = (self.rows_in_page + 1) * rep_index.bytes_per_value;
3468            Some(rep_index.buf_position..(rep_index.buf_position + total_size))
3469        } else {
3470            None
3471        };
3472        Ok(PageInitialization::new([repetition_index_range]))
3473    }
3474
3475    fn init_from_buffers<'a>(
3476        &'a mut self,
3477        buffers: PageInitializationBuffers,
3478        _io: &Arc<dyn EncodingsIo>,
3479    ) -> BoxFuture<'a, Result<Arc<dyn CachedPageData>>> {
3480        let [rep_index_buffer] = match buffers.try_into_array("full-zip") {
3481            Ok(buffers) => buffers,
3482            Err(error) => return std::future::ready(Err(error)).boxed(),
3483        };
3484        if let Some(rep_index_buffer) = rep_index_buffer {
3485            let state = Arc::new(FullZipCacheableState {
3486                rep_index_buffer: LanceBuffer::from_bytes(rep_index_buffer, 1),
3487            });
3488            std::future::ready(Ok(state as Arc<dyn CachedPageData>)).boxed()
3489        } else {
3490            std::future::ready(Ok(Arc::new(NoCachedPageData) as Arc<dyn CachedPageData>)).boxed()
3491        }
3492    }
3493
3494    /// Loads previously cached repetition index data from the cache system.
3495    /// This method is called when a scheduler instance needs to use cached data
3496    /// that was initialized by another instance or in a previous operation.
3497    fn try_load(&mut self, cache: &Arc<dyn CachedPageData>) -> Result<()> {
3498        self.cached_state = Some(cache.clone().as_arc_any().downcast().map_err(|_| {
3499            Error::invalid_input_source("Cached full-zip page data has an unexpected type".into())
3500        })?);
3501        Ok(())
3502    }
3503
3504    fn schedule_ranges(
3505        &self,
3506        ranges: &[Range<u64>],
3507        io: &Arc<dyn EncodingsIo>,
3508    ) -> Result<Vec<PageLoadTask>> {
3509        if let Some(rep_index) = self.rep_index {
3510            self.schedule_ranges_rep(ranges, io, rep_index)
3511        } else {
3512            self.schedule_ranges_simple(ranges, io)
3513        }
3514    }
3515}
3516
3517/// A decoder for full-zip encoded data when the data has a fixed-width
3518///
3519/// Here we need to unzip the control words from the values themselves and
3520/// then decompress the requested values.
3521///
3522/// We use a PerValueDecompressor because we will only be decompressing the
3523/// requested data.  This decoder / scheduler does not do any read amplification.
3524#[derive(Debug)]
3525struct FixedFullZipDecoder {
3526    details: Arc<FullZipDecodeDetails>,
3527    data: VecDeque<LanceBuffer>,
3528    offset_in_current: usize,
3529    bytes_per_value: usize,
3530    total_bytes_per_value: usize,
3531    num_rows: u64,
3532}
3533
3534impl FixedFullZipDecoder {
3535    fn slice_next_task(&mut self, num_rows: u64) -> Result<FullZipDecodeTaskItem> {
3536        debug_assert!(num_rows > 0);
3537        let cur_buf = self.data.front_mut().ok_or_else(|| {
3538            Error::corrupt_file_named(
3539                "fixed_full_zip",
3540                format!("page data ran out with {num_rows} row(s) left to decode"),
3541            )
3542        })?;
3543        let start = self.offset_in_current;
3544        if self.details.ctrl_word_parser.has_rep() {
3545            // This is a slightly slower path.  In order to figure out where to split we need to
3546            // examine the rep index so we can convert num_lists to num_rows
3547            let mut rows_started = 0;
3548            // We always need at least one value.  Now loop through until we have passed num_rows
3549            // values
3550            let mut num_items = 0;
3551            let bytes_per_word = self.details.ctrl_word_parser.bytes_per_word();
3552            while self.offset_in_current < cur_buf.len() {
3553                // The item walk is driven by the control words themselves, so a page that
3554                // disagrees with its own layout can step past the end of the buffer.  Check
3555                // before every read and every advance so that a malformed page is reported
3556                // instead of panicking in `slice_with_length`.
3557                let bytes_avail = cur_buf.len() - self.offset_in_current;
3558                if bytes_avail < bytes_per_word {
3559                    return Err(Error::corrupt_file_named(
3560                        "fixed_full_zip",
3561                        format!(
3562                            "truncated control word: {bytes_avail} byte(s) remain in the page \
3563                             buffer but a control word requires {bytes_per_word}"
3564                        ),
3565                    ));
3566                }
3567                let control = self.details.ctrl_word_parser.parse_desc(
3568                    &cur_buf[self.offset_in_current..],
3569                    self.details.max_rep,
3570                    self.details.max_visible_def,
3571                );
3572                if control.is_new_row {
3573                    if rows_started == num_rows {
3574                        break;
3575                    }
3576                    rows_started += 1;
3577                }
3578                num_items += 1;
3579                let bytes_in_item = if control.is_visible {
3580                    self.total_bytes_per_value
3581                } else {
3582                    bytes_per_word
3583                };
3584                if bytes_in_item > bytes_avail {
3585                    return Err(Error::corrupt_file_named(
3586                        "fixed_full_zip",
3587                        format!(
3588                            "truncated item: {bytes_avail} byte(s) remain in the page buffer at \
3589                             offset {} but the item requires {bytes_in_item}",
3590                            self.offset_in_current
3591                        ),
3592                    ));
3593                }
3594                self.offset_in_current += bytes_in_item;
3595            }
3596
3597            let task_slice = cur_buf.slice_with_length(start, self.offset_in_current - start);
3598            if self.offset_in_current == cur_buf.len() {
3599                self.data.pop_front();
3600                self.offset_in_current = 0;
3601            }
3602
3603            Ok(FullZipDecodeTaskItem {
3604                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3605                    data: task_slice,
3606                    bits_per_value: self.bytes_per_value as u64 * 8,
3607                    num_values: num_items,
3608                    block_info: BlockInfo::new(),
3609                }),
3610                rows_in_buf: rows_started,
3611            })
3612        } else {
3613            // If there's no repetition we can calculate the slicing point by just multiplying
3614            // the number of rows by the total bytes per value
3615            let cur_buf = self.data.front_mut().unwrap();
3616            let bytes_avail = cur_buf.len() - self.offset_in_current;
3617            let offset_in_cur = self.offset_in_current;
3618
3619            let bytes_needed = num_rows as usize * self.total_bytes_per_value;
3620            let mut rows_taken = num_rows;
3621            let task_slice = if bytes_needed >= bytes_avail {
3622                self.offset_in_current = 0;
3623                rows_taken = bytes_avail as u64 / self.total_bytes_per_value as u64;
3624                self.data
3625                    .pop_front()
3626                    .unwrap()
3627                    .slice_with_length(offset_in_cur, bytes_avail)
3628            } else {
3629                self.offset_in_current += bytes_needed;
3630                cur_buf.slice_with_length(offset_in_cur, bytes_needed)
3631            };
3632            Ok(FullZipDecodeTaskItem {
3633                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3634                    data: task_slice,
3635                    bits_per_value: self.bytes_per_value as u64 * 8,
3636                    num_values: rows_taken,
3637                    block_info: BlockInfo::new(),
3638                }),
3639                rows_in_buf: rows_taken,
3640            })
3641        }
3642    }
3643}
3644
3645impl StructuralPageDecoder for FixedFullZipDecoder {
3646    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3647        if self.total_bytes_per_value == 0 {
3648            // No bytes per row: constant-null page with no ctrl-word bytes.
3649            // The decompressor (ConstantDecompressor) ignores its input and returns AllNull.
3650            return Ok(Box::new(FixedFullZipDecodeTask {
3651                details: self.details.clone(),
3652                data: vec![FullZipDecodeTaskItem {
3653                    data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
3654                        data: LanceBuffer::empty(),
3655                        bits_per_value: 0,
3656                        num_values: num_rows,
3657                        block_info: BlockInfo::new(),
3658                    }),
3659                    rows_in_buf: num_rows,
3660                }],
3661                bytes_per_value: 0,
3662                num_rows: num_rows as usize,
3663            }));
3664        }
3665        let mut task_data = Vec::with_capacity(self.data.len());
3666        let mut remaining = num_rows;
3667        while remaining > 0 {
3668            let task_item = self.slice_next_task(remaining)?;
3669            remaining -= task_item.rows_in_buf;
3670            task_data.push(task_item);
3671        }
3672        Ok(Box::new(FixedFullZipDecodeTask {
3673            details: self.details.clone(),
3674            data: task_data,
3675            bytes_per_value: self.bytes_per_value,
3676            num_rows: num_rows as usize,
3677        }))
3678    }
3679
3680    fn num_rows(&self) -> u64 {
3681        self.num_rows
3682    }
3683}
3684
3685/// A decoder for full-zip encoded data when the data has a variable-width
3686///
3687/// Here we need to unzip the control words AND lengths from the values and
3688/// then decompress the requested values.
3689#[derive(Debug)]
3690struct VariableFullZipDecoder {
3691    details: Arc<FullZipDecodeDetails>,
3692    decompressor: Arc<dyn VariablePerValueDecompressor>,
3693    data: LanceBuffer,
3694    offsets: LanceBuffer,
3695    rep: ScalarBuffer<u16>,
3696    def: ScalarBuffer<u16>,
3697    repdef_starts: Vec<usize>,
3698    data_starts: Vec<usize>,
3699    offset_starts: Vec<usize>,
3700    visible_item_counts: Vec<u64>,
3701    bits_per_offset: u8,
3702    current_idx: usize,
3703    num_rows: u64,
3704}
3705
3706impl VariableFullZipDecoder {
3707    fn new(
3708        details: Arc<FullZipDecodeDetails>,
3709        data: VecDeque<LanceBuffer>,
3710        num_rows: u64,
3711        in_bits_per_length: u8,
3712        out_bits_per_offset: u8,
3713    ) -> Result<Self> {
3714        let decompressor = match details.value_decompressor {
3715            PerValueDecompressor::Variable(ref d) => d.clone(),
3716            _ => unreachable!(),
3717        };
3718
3719        assert_eq!(in_bits_per_length % 8, 0);
3720        assert!(out_bits_per_offset == 32 || out_bits_per_offset == 64);
3721
3722        let mut decoder = Self {
3723            details,
3724            decompressor,
3725            data: LanceBuffer::empty(),
3726            offsets: LanceBuffer::empty(),
3727            rep: LanceBuffer::empty().borrow_to_typed_slice(),
3728            def: LanceBuffer::empty().borrow_to_typed_slice(),
3729            bits_per_offset: out_bits_per_offset,
3730            repdef_starts: Vec::with_capacity(num_rows as usize + 1),
3731            data_starts: Vec::with_capacity(num_rows as usize + 1),
3732            offset_starts: Vec::with_capacity(num_rows as usize + 1),
3733            visible_item_counts: Vec::with_capacity(num_rows as usize + 1),
3734            current_idx: 0,
3735            num_rows,
3736        };
3737
3738        // There's no great time to do this and this is the least worst time.  If we don't unzip then
3739        // we can't slice the data during the decode phase.  This is because we need the offsets to be
3740        // unpacked to know where the values start and end.
3741        //
3742        // We don't want to unzip on the decode thread because that is a single-threaded path
3743        // We don't want to unzip on the scheduling thread because that is a single-threaded path
3744        //
3745        // Fortunately, we know variable length data will always be read indirectly and so we can do it
3746        // here, which should be on the indirect thread.  The primary disadvantage to doing it here is that
3747        // we load all the data into memory and then throw it away only to load it all into memory again during
3748        // the decode.
3749        //
3750        // There are some alternatives to investigate:
3751        //   - Instead of just reading the beginning and end of the rep index we could read the entire
3752        //     range in between.  This will give us the break points that we need for slicing and won't increase
3753        //     the number of IOPs but it will mean we are doing more total I/O and we need to load the rep index
3754        //     even when doing a full scan.
3755        //   - We could force each decode task to do a full unzip of all the data.  Each decode task now
3756        //     has to do more work but the work is all fused.
3757        //   - We could just try doing this work on the decode thread and see if it is a problem.
3758        decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?;
3759
3760        Ok(decoder)
3761    }
3762
3763    fn slice_batch_data_and_rebase_offsets_typed<T>(
3764        data: &LanceBuffer,
3765        offsets: &LanceBuffer,
3766    ) -> Result<(LanceBuffer, LanceBuffer)>
3767    where
3768        T: arrow_buffer::ArrowNativeType
3769            + Copy
3770            + PartialOrd
3771            + std::ops::Sub<Output = T>
3772            + std::fmt::Display
3773            + TryInto<usize>,
3774    {
3775        let offsets_slice = offsets.borrow_to_typed_slice::<T>();
3776        let offsets_slice = offsets_slice.as_ref();
3777        if offsets_slice.is_empty() {
3778            return Err(Error::internal(
3779                "Variable offsets cannot be empty".to_string(),
3780            ));
3781        }
3782
3783        let base = offsets_slice[0];
3784        let end = *offsets_slice.last().unwrap();
3785        if end < base {
3786            return Err(Error::internal(format!(
3787                "Invalid variable offsets: end ({end}) is less than base ({base})"
3788            )));
3789        }
3790
3791        let data_start = base.try_into().map_err(|_| {
3792            Error::internal(format!("Variable offset ({base}) does not fit into usize"))
3793        })?;
3794        let data_end = end.try_into().map_err(|_| {
3795            Error::internal(format!("Variable offset ({end}) does not fit into usize"))
3796        })?;
3797        if data_end > data.len() {
3798            return Err(Error::internal(format!(
3799                "Invalid variable offsets: end ({data_end}) exceeds data len ({})",
3800                data.len()
3801            )));
3802        }
3803
3804        let mut rebased_offsets = Vec::with_capacity(offsets_slice.len());
3805        for &offset in offsets_slice {
3806            if offset < base {
3807                return Err(Error::internal(format!(
3808                    "Invalid variable offsets: offset ({offset}) is less than base ({base})"
3809                )));
3810            }
3811            rebased_offsets.push(offset - base);
3812        }
3813
3814        let sliced_data = data.slice_with_length(data_start, data_end - data_start);
3815        // Copy into a compact buffer so each output batch owns only what it references.
3816        let sliced_data = LanceBuffer::copy_slice(&sliced_data);
3817        let rebased_offsets = LanceBuffer::reinterpret_vec(rebased_offsets);
3818        Ok((sliced_data, rebased_offsets))
3819    }
3820
3821    fn slice_batch_data_and_rebase_offsets(
3822        data: &LanceBuffer,
3823        offsets: &LanceBuffer,
3824        bits_per_offset: u8,
3825    ) -> Result<(LanceBuffer, LanceBuffer)> {
3826        match bits_per_offset {
3827            32 => Self::slice_batch_data_and_rebase_offsets_typed::<u32>(data, offsets),
3828            64 => Self::slice_batch_data_and_rebase_offsets_typed::<u64>(data, offsets),
3829            _ => Err(Error::internal(format!(
3830                "Unsupported bits_per_offset={bits_per_offset}"
3831            ))),
3832        }
3833    }
3834
3835    /// Reads a single length prefix from the front of `data`.
3836    ///
3837    /// The bytes come from the file. A page whose item walk ends with a partial
3838    /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this
3839    /// is bounds checked and reports a corrupt file rather than reading past the
3840    /// end of the buffer.
3841    fn parse_length(data: &[u8], bits_per_offset: u8) -> Result<u64> {
3842        let width = bits_per_offset as usize / 8;
3843        if data.len() < width {
3844            return Err(Error::corrupt_file_named(
3845                "variable_full_zip",
3846                format!(
3847                    "truncated length prefix: {} byte(s) remain in the page buffer but a \
3848                     {}-bit length prefix requires {}",
3849                    data.len(),
3850                    bits_per_offset,
3851                    width
3852                ),
3853            ));
3854        }
3855        Ok(match bits_per_offset {
3856            8 => data[0] as u64,
3857            16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64,
3858            32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64,
3859            64 => u64::from_le_bytes(data[..8].try_into().unwrap()),
3860            _ => unreachable!(),
3861        })
3862    }
3863
3864    fn unzip(
3865        &mut self,
3866        data: VecDeque<LanceBuffer>,
3867        in_bits_per_length: u8,
3868        out_bits_per_offset: u8,
3869        num_rows: u64,
3870    ) -> Result<()> {
3871        // This undercounts if there are lists but, at this point, we don't really know how many items we have
3872        let mut rep = Vec::with_capacity(num_rows as usize);
3873        let mut def = Vec::with_capacity(num_rows as usize);
3874        let bytes_cw = self.details.ctrl_word_parser.bytes_per_word() * num_rows as usize;
3875
3876        // This undercounts if there are lists
3877        // It can also overcount if there are invisible items
3878        let bytes_per_offset = out_bits_per_offset as usize / 8;
3879        let bytes_offsets = bytes_per_offset * (num_rows as usize + 1);
3880        let mut offsets_data = Vec::with_capacity(bytes_offsets);
3881
3882        let bytes_per_length = in_bits_per_length as usize / 8;
3883        let bytes_lengths = bytes_per_length * num_rows as usize;
3884
3885        let bytes_data = data.iter().map(|d| d.len()).sum::<usize>();
3886        // This overcounts since bytes_lengths and bytes_cw are undercounts
3887        // It can also undercount if there are invisible items (hence the saturating_sub)
3888        let mut unzipped_data =
3889            Vec::with_capacity((bytes_data - bytes_cw).saturating_sub(bytes_lengths));
3890
3891        let mut current_offset = 0_u64;
3892        let mut visible_item_count = 0_u64;
3893        for databuf in data.into_iter() {
3894            let mut databuf = databuf.as_ref();
3895            while !databuf.is_empty() {
3896                let data_start = unzipped_data.len();
3897                let offset_start = offsets_data.len();
3898                // We might have only-rep or only-def, neither, or both.  They move at the same
3899                // speed though so we only need one index into it
3900                let repdef_start = rep.len().max(def.len());
3901                // TODO: Kind of inefficient we parse the control word twice here
3902                let ctrl_desc = self.details.ctrl_word_parser.parse_desc(
3903                    databuf,
3904                    self.details.max_rep,
3905                    self.details.max_visible_def,
3906                );
3907                self.details
3908                    .ctrl_word_parser
3909                    .parse(databuf, &mut rep, &mut def);
3910                databuf = &databuf[self.details.ctrl_word_parser.bytes_per_word()..];
3911
3912                if ctrl_desc.is_new_row {
3913                    self.repdef_starts.push(repdef_start);
3914                    self.data_starts.push(data_start);
3915                    self.offset_starts.push(offset_start);
3916                    self.visible_item_counts.push(visible_item_count);
3917                }
3918                if ctrl_desc.is_visible {
3919                    visible_item_count += 1;
3920                    if ctrl_desc.is_valid_item {
3921                        let length = Self::parse_length(databuf, in_bits_per_length)?;
3922                        match out_bits_per_offset {
3923                            32 => offsets_data
3924                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3925                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
3926                            _ => unreachable!(),
3927                        };
3928                        databuf = &databuf[bytes_per_offset..];
3929                        unzipped_data.extend_from_slice(&databuf[..length as usize]);
3930                        databuf = &databuf[length as usize..];
3931                        current_offset += length;
3932                    } else {
3933                        // Null items still get an offset
3934                        match out_bits_per_offset {
3935                            32 => offsets_data
3936                                .extend_from_slice(&(current_offset as u32).to_le_bytes()),
3937                            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
3938                            _ => unreachable!(),
3939                        }
3940                    }
3941                }
3942            }
3943        }
3944        self.repdef_starts.push(rep.len().max(def.len()));
3945        self.data_starts.push(unzipped_data.len());
3946        self.offset_starts.push(offsets_data.len());
3947        self.visible_item_counts.push(visible_item_count);
3948        match out_bits_per_offset {
3949            32 => offsets_data.extend_from_slice(&(current_offset as u32).to_le_bytes()),
3950            64 => offsets_data.extend_from_slice(&current_offset.to_le_bytes()),
3951            _ => unreachable!(),
3952        };
3953        self.rep = ScalarBuffer::from(rep);
3954        self.def = ScalarBuffer::from(def);
3955        self.data = LanceBuffer::from(unzipped_data);
3956        self.offsets = LanceBuffer::from(offsets_data);
3957        Ok(())
3958    }
3959}
3960
3961impl StructuralPageDecoder for VariableFullZipDecoder {
3962    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn DecodePageTask>> {
3963        let start = self.current_idx;
3964        let end = start + num_rows as usize;
3965
3966        let offset_start = self.offset_starts[start];
3967        let offset_end = self.offset_starts[end] + (self.bits_per_offset as usize / 8);
3968        let offsets = self
3969            .offsets
3970            .slice_with_length(offset_start, offset_end - offset_start);
3971        // Keep each batch's variable data buffer bounded to the selected rows.
3972        let (data, offsets) =
3973            Self::slice_batch_data_and_rebase_offsets(&self.data, &offsets, self.bits_per_offset)?;
3974
3975        let repdef_start = self.repdef_starts[start];
3976        let repdef_end = self.repdef_starts[end];
3977        let rep = if self.rep.is_empty() {
3978            self.rep.clone()
3979        } else {
3980            self.rep.slice(repdef_start, repdef_end - repdef_start)
3981        };
3982        let def = if self.def.is_empty() {
3983            self.def.clone()
3984        } else {
3985            self.def.slice(repdef_start, repdef_end - repdef_start)
3986        };
3987
3988        let visible_item_counts_start = self.visible_item_counts[start];
3989        let visible_item_counts_end = self.visible_item_counts[end];
3990        let num_visible_items = visible_item_counts_end - visible_item_counts_start;
3991
3992        self.current_idx += num_rows as usize;
3993
3994        Ok(Box::new(VariableFullZipDecodeTask {
3995            details: self.details.clone(),
3996            decompressor: self.decompressor.clone(),
3997            data,
3998            offsets,
3999            bits_per_offset: self.bits_per_offset,
4000            num_visible_items,
4001            rep,
4002            def,
4003        }))
4004    }
4005
4006    fn num_rows(&self) -> u64 {
4007        self.num_rows
4008    }
4009}
4010
4011#[derive(Debug)]
4012struct VariableFullZipDecodeTask {
4013    details: Arc<FullZipDecodeDetails>,
4014    decompressor: Arc<dyn VariablePerValueDecompressor>,
4015    data: LanceBuffer,
4016    offsets: LanceBuffer,
4017    bits_per_offset: u8,
4018    num_visible_items: u64,
4019    rep: ScalarBuffer<u16>,
4020    def: ScalarBuffer<u16>,
4021}
4022
4023impl DecodePageTask for VariableFullZipDecodeTask {
4024    fn decode(self: Box<Self>) -> Result<DecodedPage> {
4025        let block = VariableWidthBlock {
4026            data: self.data,
4027            offsets: self.offsets,
4028            bits_per_offset: self.bits_per_offset,
4029            num_values: self.num_visible_items,
4030            block_info: BlockInfo::new(),
4031        };
4032        let decomopressed = self.decompressor.decompress(block)?;
4033        let rep = if self.rep.is_empty() {
4034            None
4035        } else {
4036            Some(self.rep.to_vec())
4037        };
4038        let def = if self.def.is_empty() {
4039            None
4040        } else {
4041            Some(self.def.to_vec())
4042        };
4043        let unraveler = RepDefUnraveler::new(
4044            rep,
4045            def,
4046            self.details.def_meaning.clone(),
4047            self.num_visible_items,
4048        );
4049        Ok(DecodedPage {
4050            data: decomopressed,
4051            repdef: unraveler,
4052        })
4053    }
4054}
4055
4056#[derive(Debug)]
4057struct FullZipDecodeTaskItem {
4058    data: PerValueDataBlock,
4059    rows_in_buf: u64,
4060}
4061
4062/// A task to unzip and decompress full-zip encoded data when that data
4063/// has a fixed-width.
4064#[derive(Debug)]
4065struct FixedFullZipDecodeTask {
4066    details: Arc<FullZipDecodeDetails>,
4067    data: Vec<FullZipDecodeTaskItem>,
4068    num_rows: usize,
4069    bytes_per_value: usize,
4070}
4071
4072impl DecodePageTask for FixedFullZipDecodeTask {
4073    fn decode(self: Box<Self>) -> Result<DecodedPage> {
4074        let estimated_size_bytes = if self.details.ctrl_word_parser.bytes_per_word() == 0 {
4075            let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor else {
4076                return Err(Error::internal(
4077                    "FixedFullZipDecodeTask requires a fixed-width decompressor",
4078                ));
4079            };
4080            decompressor
4081                .decoded_size_bytes(self.num_rows as u64)
4082                .unwrap_or_else(|| {
4083                    self.data
4084                        .iter()
4085                        .map(|task_item| task_item.data.data_size())
4086                        .sum::<u64>()
4087                        * 2
4088                })
4089        } else {
4090            // Rep/def levels can suppress values, so the exact output size is not known
4091            // until they are decoded. Keep the existing conservative estimate.
4092            self.data
4093                .iter()
4094                .map(|task_item| task_item.data.data_size())
4095                .sum::<u64>()
4096                * 2
4097        };
4098        let mut data_builder = DataBlockBuilder::with_capacity_estimate(estimated_size_bytes);
4099
4100        if self.details.ctrl_word_parser.bytes_per_word() == 0 {
4101            // Fast path, no need to unzip because there is no rep/def
4102            //
4103            // We decompress each buffer and add it to our output buffer
4104            for task_item in self.data.into_iter() {
4105                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
4106                    unreachable!()
4107                };
4108                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
4109                else {
4110                    unreachable!()
4111                };
4112                debug_assert_eq!(fixed_data.num_values, task_item.rows_in_buf);
4113                let decompressed = decompressor.decompress(fixed_data, task_item.rows_in_buf)?;
4114                data_builder.append(&decompressed, 0..task_item.rows_in_buf)?;
4115            }
4116
4117            let unraveler = RepDefUnraveler::new(
4118                None,
4119                None,
4120                self.details.def_meaning.clone(),
4121                self.num_rows as u64,
4122            );
4123
4124            Ok(DecodedPage {
4125                data: data_builder.finish(),
4126                repdef: unraveler,
4127            })
4128        } else {
4129            // Slow path, unzipping needed
4130            let mut rep = Vec::with_capacity(self.num_rows);
4131            let mut def = Vec::with_capacity(self.num_rows);
4132
4133            for task_item in self.data.into_iter() {
4134                let PerValueDataBlock::Fixed(fixed_data) = task_item.data else {
4135                    unreachable!()
4136                };
4137                let mut buf_slice = fixed_data.data.as_ref();
4138                let num_values = fixed_data.num_values as usize;
4139                // We will be unzipping repdef in to `rep` and `def` and the
4140                // values into `values` (which contains the compressed values)
4141                let mut values = Vec::with_capacity(
4142                    fixed_data.data.len()
4143                        - (self.details.ctrl_word_parser.bytes_per_word() * num_values),
4144                );
4145                let mut visible_items = 0;
4146                for _ in 0..num_values {
4147                    // Extract rep/def
4148                    self.details
4149                        .ctrl_word_parser
4150                        .parse(buf_slice, &mut rep, &mut def);
4151                    buf_slice = &buf_slice[self.details.ctrl_word_parser.bytes_per_word()..];
4152
4153                    let is_visible = def
4154                        .last()
4155                        .map(|d| *d <= self.details.max_visible_def)
4156                        .unwrap_or(true);
4157                    if is_visible {
4158                        // Extract value
4159                        values.extend_from_slice(buf_slice[..self.bytes_per_value].as_ref());
4160                        buf_slice = &buf_slice[self.bytes_per_value..];
4161                        visible_items += 1;
4162                    }
4163                }
4164
4165                // Finally, we decompress the values and add them to our output buffer
4166                let values_buf = LanceBuffer::from(values);
4167                let fixed_data = FixedWidthDataBlock {
4168                    bits_per_value: self.bytes_per_value as u64 * 8,
4169                    block_info: BlockInfo::new(),
4170                    data: values_buf,
4171                    num_values: visible_items,
4172                };
4173                let PerValueDecompressor::Fixed(decompressor) = &self.details.value_decompressor
4174                else {
4175                    unreachable!()
4176                };
4177                let decompressed = decompressor.decompress(fixed_data, visible_items)?;
4178                data_builder.append(&decompressed, 0..visible_items)?;
4179            }
4180
4181            let repetition = if rep.is_empty() { None } else { Some(rep) };
4182            let definition = if def.is_empty() { None } else { Some(def) };
4183
4184            let unraveler = RepDefUnraveler::new(
4185                repetition,
4186                definition,
4187                self.details.def_meaning.clone(),
4188                self.num_rows as u64,
4189            );
4190            let data = data_builder.finish();
4191
4192            Ok(DecodedPage {
4193                data,
4194                repdef: unraveler,
4195            })
4196        }
4197    }
4198}
4199
4200#[derive(Debug)]
4201struct StructuralPrimitiveFieldSchedulingJob<'a> {
4202    scheduler: &'a StructuralPrimitiveFieldScheduler,
4203    mappings: std::vec::IntoIter<PageRangeMapping>,
4204}
4205
4206impl<'a> StructuralPrimitiveFieldSchedulingJob<'a> {
4207    pub fn try_new(
4208        scheduler: &'a StructuralPrimitiveFieldScheduler,
4209        ranges: Vec<Range<u64>>,
4210    ) -> Result<Self> {
4211        let mappings = map_ranges_to_pages(&scheduler.page_row_ends, &ranges)?.into_iter();
4212        Ok(Self {
4213            scheduler,
4214            mappings,
4215        })
4216    }
4217}
4218
4219impl StructuralSchedulingJob for StructuralPrimitiveFieldSchedulingJob<'_> {
4220    fn schedule_next(&mut self, context: &mut SchedulerContext) -> Result<Vec<ScheduledScanLine>> {
4221        let Some(first_mapping) = self.mappings.next() else {
4222            return Ok(Vec::new());
4223        };
4224        // Pages are scheduled in row order until their reads add up to a
4225        // batch, then submitted together so the I/O scheduler can coalesce
4226        // neighbouring pages.
4227        let batch = Arc::new(RequestBatch::new(context.io().clone()));
4228        let batch_io: Arc<dyn EncodingsIo> = batch.clone();
4229        let cur_path = context.current_path();
4230        let mut scan_lines = Vec::new();
4231        let mut mapping = Some(first_mapping);
4232        while let Some(cur_mapping) = mapping.take() {
4233            let cur_page = self.scheduler.built_page(cur_mapping.page_idx)?;
4234
4235            trace!(
4236                "Scheduling {} rows across {} ranges from page with {} rows (column_index={}, page_index={})",
4237                cur_mapping
4238                    .ranges_in_page
4239                    .iter()
4240                    .map(|r| r.end - r.start)
4241                    .sum::<u64>(),
4242                cur_mapping.ranges_in_page.len(),
4243                cur_page.row_range.end - cur_page.row_range.start,
4244                self.scheduler.column_index,
4245                cur_page.page_index,
4246            );
4247
4248            let page_decoders = cur_page
4249                .scheduler
4250                .schedule_ranges(&cur_mapping.ranges_in_page, &batch_io)?;
4251
4252            scan_lines.extend(page_decoders.into_iter().map(|page_load_task| {
4253                let cur_path = cur_path.clone();
4254                let page_decoder = page_load_task.decoder_fut;
4255                let unloaded_page = async move {
4256                    let page_decoder = page_decoder.await?;
4257                    Ok(LoadedPageShard {
4258                        decoder: page_decoder,
4259                        path: cur_path,
4260                    })
4261                }
4262                .boxed();
4263                ScheduledScanLine {
4264                    decoders: vec![MessageType::UnloadedPage(UnloadedPageShard(unloaded_page))],
4265                    rows_scheduled: page_load_task.num_rows,
4266                }
4267            }));
4268
4269            if batch.pending_bytes() < IO_REQUEST_BATCH_BYTES {
4270                mapping = self.mappings.next();
4271            }
4272        }
4273        batch.flush();
4274        Ok(scan_lines)
4275    }
4276}
4277
4278#[derive(Debug)]
4279struct PageInfoAndScheduler {
4280    page_index: usize,
4281    row_range: Range<u64>,
4282    scheduler: Box<dyn StructuralPageScheduler>,
4283}
4284
4285/// Everything needed to build a page scheduler from its `PageInfo` on demand.
4286///
4287/// Building a page scheduler clones buffer tables and instantiates
4288/// decompressors, so a column with hundreds of thousands of pages must not pay
4289/// for every page on every read. The field scheduler keeps this source and
4290/// builds only the pages `initialize` is asked to cover.
4291#[derive(Debug)]
4292struct PageSchedulerSource {
4293    column_info: Arc<ColumnInfo>,
4294    decompressors: Arc<dyn DecompressionStrategy>,
4295    cache_repetition_index: bool,
4296    target_field: Field,
4297}
4298
4299/// A scheduler for a leaf node
4300///
4301/// Here we look at the layout of the various pages and delegate scheduling to a scheduler
4302/// appropriate for the layout of the page.
4303#[derive(Debug)]
4304pub struct StructuralPrimitiveFieldScheduler {
4305    /// Exclusive end row of every page in the column, in page order, so page
4306    /// `i` covers `page_row_ends[i - 1]..page_row_ends[i]` (page 0 starts at 0).
4307    page_row_ends: Vec<u64>,
4308    /// Page schedulers built so far, ordered by `page_index`. After
4309    /// `initialize` this holds exactly the pages the requested ranges touch,
4310    /// or every page when no ranges were given.
4311    page_schedulers: Vec<PageInfoAndScheduler>,
4312    /// Metadata to build pages that are not in `page_schedulers` yet. `None`
4313    /// once every page has been built.
4314    page_source: Option<PageSchedulerSource>,
4315    column_index: u32,
4316    target_data_type_tag: Arc<str>,
4317}
4318
4319impl StructuralPrimitiveFieldScheduler {
4320    /// Creates the scheduler and builds every page scheduler up front.
4321    ///
4322    /// This pays for all pages of the column on every read. Prefer
4323    /// [`Self::try_new_lazy`], which builds only the pages a read touches.
4324    #[deprecated(
4325        since = "13.0.0",
4326        note = "builds every page scheduler eagerly; use try_new_lazy"
4327    )]
4328    pub fn try_new(
4329        column_info: &ColumnInfo,
4330        decompressors: &dyn DecompressionStrategy,
4331        cache_repetition_index: bool,
4332        target_field: &Field,
4333    ) -> Result<Self> {
4334        let page_row_ends = Self::page_row_ends(column_info)?;
4335        let mut page_schedulers = Vec::with_capacity(page_row_ends.len());
4336        let mut row_start = 0_u64;
4337        for (page_index, (page_info, &row_end)) in column_info
4338            .page_infos
4339            .iter()
4340            .zip(&page_row_ends)
4341            .enumerate()
4342        {
4343            page_schedulers.push(Self::page_info_to_scheduler(
4344                page_info,
4345                page_index,
4346                row_start..row_end,
4347                decompressors,
4348                cache_repetition_index,
4349                target_field,
4350            )?);
4351            row_start = row_end;
4352        }
4353        Ok(Self {
4354            page_row_ends,
4355            page_schedulers,
4356            page_source: None,
4357            column_index: column_info.index,
4358            target_data_type_tag: format!("{:?}", target_field.data_type()).into(),
4359        })
4360    }
4361
4362    /// Creates the scheduler without building any page schedulers.
4363    ///
4364    /// Only the per-page row boundaries are computed here. Page schedulers
4365    /// are built in [`StructuralFieldScheduler::initialize`] for the pages the
4366    /// requested ranges touch (or for every page when no ranges are given), so
4367    /// a small read from a column with many pages stays cheap.
4368    pub fn try_new_lazy(
4369        column_info: Arc<ColumnInfo>,
4370        decompressors: Arc<dyn DecompressionStrategy>,
4371        cache_repetition_index: bool,
4372        target_field: &Field,
4373    ) -> Result<Self> {
4374        Ok(Self {
4375            page_row_ends: Self::page_row_ends(&column_info)?,
4376            page_schedulers: Vec::new(),
4377            column_index: column_info.index,
4378            target_data_type_tag: format!("{:?}", target_field.data_type()).into(),
4379            page_source: Some(PageSchedulerSource {
4380                column_info,
4381                decompressors,
4382                cache_repetition_index,
4383                target_field: target_field.clone(),
4384            }),
4385        })
4386    }
4387
4388    /// Exclusive end row of every page of `column_info`, in page order.
4389    fn page_row_ends(column_info: &ColumnInfo) -> Result<Vec<u64>> {
4390        let mut page_row_ends = Vec::with_capacity(column_info.page_infos.len());
4391        let mut row_start = 0_u64;
4392        for (page_index, page_info) in column_info.page_infos.iter().enumerate() {
4393            let row_end = row_start.checked_add(page_info.num_rows).ok_or_else(|| {
4394                Error::invalid_input_source(
4395                    format!(
4396                        "Page row range overflowed for column {} at page {page_index}: start={row_start}, rows={}",
4397                        column_info.index, page_info.num_rows
4398                    )
4399                    .into(),
4400                )
4401            })?;
4402            page_row_ends.push(row_end);
4403            row_start = row_end;
4404        }
4405        Ok(page_row_ends)
4406    }
4407
4408    /// Wraps already-built page schedulers; nothing is built lazily.
4409    #[cfg(test)]
4410    fn from_page_schedulers(
4411        page_schedulers: Vec<PageInfoAndScheduler>,
4412        column_index: u32,
4413        target_data_type_tag: Arc<str>,
4414    ) -> Self {
4415        Self {
4416            page_row_ends: page_schedulers
4417                .iter()
4418                .map(|page| page.row_range.end)
4419                .collect(),
4420            page_schedulers,
4421            page_source: None,
4422            column_index,
4423            target_data_type_tag,
4424        }
4425    }
4426
4427    fn page_row_range(&self, page_index: usize) -> Range<u64> {
4428        let start = match page_index.checked_sub(1) {
4429            Some(previous) => self.page_row_ends[previous],
4430            None => 0,
4431        };
4432        start..self.page_row_ends[page_index]
4433    }
4434
4435    /// Position of `page_index` in `page_schedulers` if it has been built.
4436    fn built_page_position(&self, page_index: usize) -> Option<usize> {
4437        self.page_schedulers
4438            .binary_search_by_key(&page_index, |page| page.page_index)
4439            .ok()
4440    }
4441
4442    fn built_page(&self, page_index: usize) -> Result<&PageInfoAndScheduler> {
4443        let position = self.built_page_position(page_index).ok_or_else(|| {
4444            Error::internal(format!(
4445                "Page {page_index} of column {} was scheduled but not initialized; the scheduled ranges must be covered by the ranges passed to initialize",
4446                self.column_index
4447            ))
4448        })?;
4449        Ok(&self.page_schedulers[position])
4450    }
4451
4452    /// Builds the schedulers for `page_indices` (ascending) that do not exist
4453    /// yet and returns the positions of all of them in `page_schedulers`.
4454    fn build_pages(&mut self, page_indices: &[usize]) -> Result<Vec<usize>> {
4455        if let Some(source) = &self.page_source {
4456            let mut missing = Vec::new();
4457            for &page_index in page_indices {
4458                if self.built_page_position(page_index).is_some() {
4459                    continue;
4460                }
4461                let page_info = source
4462                    .column_info
4463                    .page_infos
4464                    .get(page_index)
4465                    .ok_or_else(|| {
4466                        Error::internal(format!(
4467                            "Page {page_index} is out of bounds for column {} with {} pages",
4468                            self.column_index,
4469                            source.column_info.page_infos.len()
4470                        ))
4471                    })?;
4472                missing.push(Self::page_info_to_scheduler(
4473                    page_info,
4474                    page_index,
4475                    self.page_row_range(page_index),
4476                    source.decompressors.as_ref(),
4477                    source.cache_repetition_index,
4478                    &source.target_field,
4479                )?);
4480            }
4481            if !missing.is_empty() {
4482                self.page_schedulers = std::mem::take(&mut self.page_schedulers)
4483                    .into_iter()
4484                    .merge_by(missing, |built, new| built.page_index < new.page_index)
4485                    .collect();
4486            }
4487            if self.page_schedulers.len() == self.page_row_ends.len() {
4488                self.page_source = None;
4489            }
4490        }
4491        page_indices
4492            .iter()
4493            .map(|&page_index| {
4494                self.built_page_position(page_index).ok_or_else(|| {
4495                    Error::internal(format!(
4496                        "Page {page_index} of column {} is missing after construction",
4497                        self.column_index
4498                    ))
4499                })
4500            })
4501            .collect()
4502    }
4503
4504    fn page_layout_to_scheduler(
4505        page_info: &PageInfo,
4506        page_layout: &PageLayout,
4507        decompressors: &dyn DecompressionStrategy,
4508        cache_repetition_index: bool,
4509        target_field: &Field,
4510    ) -> Result<Box<dyn StructuralPageScheduler>> {
4511        use pb21::page_layout::Layout;
4512        Ok(match page_layout.layout.as_ref().expect_ok()? {
4513            Layout::MiniBlockLayout(mini_block) => Box::new(MiniBlockScheduler::try_new(
4514                &page_info.buffer_offsets_and_sizes,
4515                page_info.priority,
4516                mini_block.num_items,
4517                mini_block,
4518                decompressors,
4519            )?),
4520            Layout::SparseLayout(sparse_layout) => {
4521                Box::new(sparse::SparseStructuralScheduler::try_new(
4522                    &page_info.buffer_offsets_and_sizes,
4523                    page_info.priority,
4524                    page_info.num_rows,
4525                    target_field.data_type(),
4526                    sparse_layout,
4527                    decompressors,
4528                )?)
4529            }
4530            Layout::FullZipLayout(full_zip) => {
4531                let mut scheduler = FullZipScheduler::try_new(
4532                    &page_info.buffer_offsets_and_sizes,
4533                    page_info.priority,
4534                    page_info.num_rows,
4535                    full_zip,
4536                    decompressors,
4537                )?;
4538                scheduler.enable_cache = cache_repetition_index;
4539                Box::new(scheduler)
4540            }
4541            Layout::ConstantLayout(constant_layout) => {
4542                let def_meaning = constant_layout
4543                    .layers
4544                    .iter()
4545                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4546                    .collect::<Vec<_>>();
4547                let has_scalar_value = constant_layout.inline_value.is_some()
4548                    || page_info.buffer_offsets_and_sizes.len() == 1
4549                    || page_info.buffer_offsets_and_sizes.len() == 3;
4550                if has_scalar_value {
4551                    Box::new(constant::ConstantPageScheduler::try_new(
4552                        page_info.buffer_offsets_and_sizes.clone(),
4553                        constant_layout.inline_value.clone(),
4554                        target_field.data_type(),
4555                        def_meaning.into(),
4556                    )?) as Box<dyn StructuralPageScheduler>
4557                } else if def_meaning.len() == 1
4558                    && def_meaning[0] == DefinitionInterpretation::NullableItem
4559                {
4560                    Box::new(SimpleAllNullScheduler::default()) as Box<dyn StructuralPageScheduler>
4561                } else {
4562                    // RLE levels select a validated cache representation; other
4563                    // block compressions keep flowing through the eager decompressor.
4564                    let rep_codec = LevelCodec::try_new(
4565                        constant_layout.rep_compression.as_ref(),
4566                        decompressors,
4567                    )?;
4568                    let def_codec = LevelCodec::try_new(
4569                        constant_layout.def_compression.as_ref(),
4570                        decompressors,
4571                    )?;
4572
4573                    Box::new(ComplexAllNullScheduler::new(
4574                        page_info.buffer_offsets_and_sizes.clone(),
4575                        def_meaning.into(),
4576                        rep_codec,
4577                        def_codec,
4578                        constant_layout.num_rep_values,
4579                        constant_layout.num_def_values,
4580                    )) as Box<dyn StructuralPageScheduler>
4581                }
4582            }
4583            Layout::BlobLayout(blob) => {
4584                let inner_scheduler = Self::page_layout_to_scheduler(
4585                    page_info,
4586                    blob.inner_layout.as_ref().expect_ok()?.as_ref(),
4587                    decompressors,
4588                    cache_repetition_index,
4589                    target_field,
4590                )?;
4591                let def_meaning = blob
4592                    .layers
4593                    .iter()
4594                    .map(|l| ProtobufUtils21::repdef_layer_to_def_interp(*l))
4595                    .collect::<Vec<_>>();
4596                if matches!(target_field.data_type(), DataType::Struct(_)) {
4597                    // User wants to decode blob into struct
4598                    Box::new(BlobDescriptionPageScheduler::new(
4599                        inner_scheduler,
4600                        def_meaning.into(),
4601                    ))
4602                } else {
4603                    // User wants to decode blob into binary data
4604                    Box::new(BlobPageScheduler::new(
4605                        inner_scheduler,
4606                        page_info.priority,
4607                        page_info.num_rows,
4608                        def_meaning.into(),
4609                    ))
4610                }
4611            }
4612        })
4613    }
4614
4615    fn page_info_to_scheduler(
4616        page_info: &PageInfo,
4617        page_index: usize,
4618        row_range: Range<u64>,
4619        decompressors: &dyn DecompressionStrategy,
4620        cache_repetition_index: bool,
4621        target_field: &Field,
4622    ) -> Result<PageInfoAndScheduler> {
4623        let page_layout = page_info.encoding.as_structural();
4624        let scheduler = Self::page_layout_to_scheduler(
4625            page_info,
4626            page_layout,
4627            decompressors,
4628            cache_repetition_index,
4629            target_field,
4630        )?;
4631        Ok(PageInfoAndScheduler {
4632            page_index,
4633            row_range,
4634            scheduler,
4635        })
4636    }
4637}
4638
4639pub trait CachedPageData: Any + Send + Sync + DeepSizeOf + 'static {
4640    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
4641}
4642
4643pub struct NoCachedPageData;
4644
4645impl DeepSizeOf for NoCachedPageData {
4646    fn deep_size_of_children(&self, _ctx: &mut Context) -> usize {
4647        0
4648    }
4649}
4650impl CachedPageData for NoCachedPageData {
4651    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
4652        self
4653    }
4654}
4655
4656#[doc = "Deprecated compatibility wrapper; new code uses per-page PageDataCacheKey entries."]
4657pub struct CachedFieldData {
4658    pages: Vec<Arc<dyn CachedPageData>>,
4659}
4660
4661impl DeepSizeOf for CachedFieldData {
4662    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
4663        self.pages.deep_size_of_children(ctx)
4664    }
4665}
4666
4667/// Compatibility key for the former whole-field page metadata cache.
4668#[derive(Debug, Clone)]
4669pub struct FieldDataCacheKey {
4670    pub column_index: u32,
4671    pub view_tag: String,
4672}
4673
4674impl CacheKey for FieldDataCacheKey {
4675    type ValueType = CachedFieldData;
4676
4677    fn key(&self) -> std::borrow::Cow<'_, str> {
4678        format!("{}:{}", self.column_index, self.view_tag).into()
4679    }
4680
4681    fn type_name() -> &'static str {
4682        "FieldData"
4683    }
4684
4685    fn schema() -> CacheKeySchema {
4686        CacheKeySchema::new("lance.encoding.logical.primitive.field-data-key", 1)
4687    }
4688
4689    fn write_key(&self, builder: &mut KeyBuilder) {
4690        builder.write_u32(self.column_index);
4691        builder.write_str(&self.view_tag);
4692    }
4693}
4694
4695/// Sized wrapper around a page's cached metadata so [`PageDataCacheKey`] can use
4696/// the sized [`CacheKey`] path; the concrete per-layout state still lives behind
4697/// `dyn CachedPageData`.
4698pub struct CachedPage {
4699    data: Arc<dyn CachedPageData>,
4700}
4701
4702impl DeepSizeOf for CachedPage {
4703    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
4704        self.data.deep_size_of_children(ctx)
4705    }
4706}
4707
4708// Cache key for a single page's initialized metadata, keyed per page (not per
4709// column) so a read populates and reuses only the pages it touches.
4710//
4711// `view` distinguishes decode shapes of one physical column: a blob column
4712// materializes as a `Struct<position, size>` descriptor in one variant and raw
4713// `LargeBinary` in another, each caching a different `CachedPageData` type.
4714// Without it, two readers at the same column/page but different shapes collide
4715// and crash on a downcast.
4716/// Decode shape used to isolate cached page metadata for one physical column.
4717#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4718#[non_exhaustive]
4719pub enum PageCacheView {
4720    /// The column's native non-blob representation.
4721    Native,
4722    /// Blob position and size descriptors materialized as a struct.
4723    BlobDescriptor,
4724    /// Blob contents materialized as binary values.
4725    BlobPayload,
4726}
4727
4728impl PageCacheView {
4729    fn cache_tag(self) -> &'static str {
4730        match self {
4731            Self::Native => "n",
4732            Self::BlobDescriptor => "bd",
4733            Self::BlobPayload => "bp",
4734        }
4735    }
4736}
4737
4738#[derive(Debug, Clone)]
4739pub struct PageDataCacheKey {
4740    pub column_index: u32,
4741    pub page_index: u32,
4742    pub view: PageCacheView,
4743    pub target_data_type_tag: Arc<str>,
4744}
4745
4746impl CacheKey for PageDataCacheKey {
4747    type ValueType = CachedPage;
4748
4749    fn key(&self) -> std::borrow::Cow<'_, str> {
4750        format!(
4751            "v2/{}/{}/{}/{}",
4752            self.column_index,
4753            self.page_index,
4754            self.view.cache_tag(),
4755            self.target_data_type_tag
4756        )
4757        .into()
4758    }
4759
4760    fn type_name() -> &'static str {
4761        "PageData"
4762    }
4763
4764    fn schema() -> CacheKeySchema {
4765        CacheKeySchema::new("lance.encoding.logical.primitive.page-data-key", 2)
4766    }
4767
4768    fn write_key(&self, builder: &mut KeyBuilder) {
4769        builder.write_u32(self.column_index);
4770        builder.write_u32(self.page_index);
4771        builder.write_variant(match self.view {
4772            PageCacheView::Native => 0,
4773            PageCacheView::BlobDescriptor => 1,
4774            PageCacheView::BlobPayload => 2,
4775        });
4776        builder.write_str(&self.target_data_type_tag);
4777    }
4778}
4779
4780/// One page's slice of a scheduling request: the page's index and the
4781/// page-local row sub-ranges (relative to the page's first row) of the request
4782/// that fall in it.
4783#[derive(Debug)]
4784struct PageRangeMapping {
4785    page_idx: usize,
4786    ranges_in_page: Vec<Range<u64>>,
4787}
4788
4789/// Visits every intersection between ordered row ranges and ordered pages.
4790///
4791/// `page_row_ends` holds the exclusive end row of each page (see
4792/// [`StructuralPrimitiveFieldScheduler::page_row_ends`]). Each range
4793/// binary-searches its first overlapping page and then visits only the pages
4794/// it touches. For `R` ranges, `P` pages, and `T` intersections this takes
4795/// `O(R log P + T)` time without scanning sparse gaps.
4796fn for_each_page_range(
4797    page_row_ends: &[u64],
4798    ranges: &[Range<u64>],
4799    mut visit: impl FnMut(usize, Range<u64>),
4800) -> Result<()> {
4801    for (range_idx, range) in ranges.iter().enumerate() {
4802        if range.start > range.end {
4803            return Err(Error::invalid_input_source(
4804                format!(
4805                    "Requested row range {range_idx} is reversed: {}..{}",
4806                    range.start, range.end
4807                )
4808                .into(),
4809            ));
4810        }
4811        if range_idx > 0 {
4812            let previous = &ranges[range_idx - 1];
4813            if previous.start >= range.start || previous.end > range.start {
4814                return Err(Error::invalid_input_source(
4815                    format!(
4816                        "Requested row ranges overlap or are not strictly ordered at index {range_idx}: {previous:?}, then {range:?}"
4817                    )
4818                    .into(),
4819                ));
4820            }
4821        }
4822    }
4823
4824    for range in ranges {
4825        let mut page_idx = page_row_ends.partition_point(|page_end| *page_end <= range.start);
4826        let mut page_start = match page_idx.checked_sub(1) {
4827            Some(previous) => page_row_ends[previous],
4828            None => 0,
4829        };
4830        while let Some(&page_end) = page_row_ends.get(page_idx) {
4831            if page_start >= range.end {
4832                break;
4833            }
4834            let start_in_page = range.start.max(page_start) - page_start;
4835            let end_in_page = range.end.min(page_end) - page_start;
4836            if start_in_page < end_in_page {
4837                visit(page_idx, start_in_page..end_in_page);
4838            }
4839            page_idx += 1;
4840            page_start = page_end;
4841        }
4842    }
4843    Ok(())
4844}
4845
4846/// Maps global row `ranges` onto a column's pages, yielding the page-local
4847/// sub-ranges for every page any range touches.
4848///
4849/// `ranges` must be ascending and non-overlapping (the precondition
4850/// `schedule_ranges` relies on).
4851fn map_ranges_to_pages(
4852    page_row_ends: &[u64],
4853    ranges: &[Range<u64>],
4854) -> Result<Vec<PageRangeMapping>> {
4855    let mut result: Vec<PageRangeMapping> =
4856        Vec::with_capacity(ranges.len().min(page_row_ends.len()));
4857    for_each_page_range(page_row_ends, ranges, |page_idx, range_in_page| {
4858        if let Some(mapping) = result.last_mut()
4859            && mapping.page_idx == page_idx
4860        {
4861            mapping.ranges_in_page.push(range_in_page);
4862        } else {
4863            result.push(PageRangeMapping {
4864                page_idx,
4865                ranges_in_page: vec![range_in_page],
4866            });
4867        }
4868    })?;
4869    Ok(result)
4870}
4871
4872fn pages_overlapping_ranges(page_row_ends: &[u64], ranges: &[Range<u64>]) -> Result<Vec<usize>> {
4873    let mut result = Vec::with_capacity(ranges.len().min(page_row_ends.len()));
4874    for_each_page_range(page_row_ends, ranges, |page_idx, _| {
4875        if result.last().copied() != Some(page_idx) {
4876            result.push(page_idx);
4877        }
4878    })?;
4879    Ok(result)
4880}
4881
4882#[derive(Debug)]
4883struct PageCacheMiss {
4884    /// Position of the page in `StructuralPrimitiveFieldScheduler::page_schedulers`,
4885    /// which holds only the built pages and so may differ from `page_index`.
4886    page_idx: usize,
4887    cache_key: PageDataCacheKey,
4888    initialization: PageInitialization,
4889}
4890
4891async fn read_page_initialization_buffers(
4892    io: &Arc<dyn EncodingsIo>,
4893    all_ranges: Vec<Range<u64>>,
4894) -> Result<Vec<Bytes>> {
4895    // `submit_request` coalesces (and later un-coalesces) in one forward
4896    // pass, so submit ranges in file order. Page layouts are not required
4897    // to declare their initialization ranges in that order, so retain each
4898    // range's original position and restore it before initializing pages.
4899    let mut sorted_ranges = all_ranges.into_iter().enumerate().collect::<Vec<_>>();
4900    sorted_ranges.sort_unstable_by_key(|(_, range)| (range.start, range.end));
4901    let sorted_requests = sorted_ranges
4902        .iter()
4903        .map(|(_, range)| range.clone())
4904        .collect::<Vec<_>>();
4905    let expected_buffer_count = sorted_requests.len();
4906    let sorted_buffers = if sorted_requests.is_empty() {
4907        Vec::new()
4908    } else {
4909        io.submit_request(sorted_requests, 0).await?
4910    };
4911    if sorted_buffers.len() != expected_buffer_count {
4912        return Err(Error::invalid_input_source(
4913            format!(
4914                "Page initialization I/O returned {} buffers for {} ranges",
4915                sorted_buffers.len(),
4916                expected_buffer_count
4917            )
4918            .into(),
4919        ));
4920    }
4921    let mut buffers = vec![None; expected_buffer_count];
4922    for ((original_idx, _), buffer) in sorted_ranges.into_iter().zip(sorted_buffers) {
4923        buffers[original_idx] = Some(buffer);
4924    }
4925    buffers
4926        .into_iter()
4927        .map(|buffer| {
4928            buffer.ok_or_else(|| Error::internal("Page initialization buffer was not assigned"))
4929        })
4930        .collect()
4931}
4932
4933async fn cache_initialized_page(
4934    cache: Arc<LanceCache>,
4935    cache_key: PageDataCacheKey,
4936    data: Arc<dyn CachedPageData>,
4937) {
4938    if data
4939        .clone()
4940        .as_arc_any()
4941        .downcast::<NoCachedPageData>()
4942        .is_err()
4943    {
4944        cache
4945            .insert_with_key(&cache_key, Arc::new(CachedPage { data }))
4946            .await;
4947    }
4948}
4949
4950impl StructuralPrimitiveFieldScheduler {
4951    /// Checks the built pages at `positions` (ascending indices into
4952    /// `page_schedulers`) against the cache and returns those still needing
4953    /// initialization I/O.
4954    async fn find_cache_misses(
4955        &mut self,
4956        positions: impl ExactSizeIterator<Item = usize>,
4957        cache: &Arc<LanceCache>,
4958    ) -> Result<Vec<PageCacheMiss>> {
4959        // Serve cache hits in place; the rest are misses we must read.
4960        let mut misses = Vec::with_capacity(positions.len());
4961        for page_idx in positions {
4962            let page = &self.page_schedulers[page_idx];
4963            if !page.scheduler.needs_initialization() {
4964                continue;
4965            }
4966            let page_index = u32::try_from(page.page_index).map_err(|_| {
4967                Error::invalid_input_source(
4968                    format!("Page index does not fit u32: {}", page.page_index).into(),
4969                )
4970            })?;
4971            let cache_key = PageDataCacheKey {
4972                column_index: self.column_index,
4973                page_index,
4974                view: page.scheduler.cache_view(),
4975                target_data_type_tag: self.target_data_type_tag.clone(),
4976            };
4977            if let Some(cached_page) = cache.get_with_key(&cache_key).await {
4978                match self.page_schedulers[page_idx]
4979                    .scheduler
4980                    .try_load(&cached_page.data)
4981                {
4982                    Ok(()) => continue,
4983                    Err(error) => {
4984                        debug!("Replacing invalid cached page data for {cache_key:?}: {error}");
4985                    }
4986                }
4987            }
4988            let initialization = self.page_schedulers[page_idx].scheduler.init_layout()?;
4989            misses.push(PageCacheMiss {
4990                page_idx,
4991                cache_key,
4992                initialization,
4993            });
4994        }
4995        Ok(misses)
4996    }
4997
4998    fn initialization_ranges(misses: &[PageCacheMiss]) -> Vec<Range<u64>> {
4999        let num_ranges = misses
5000            .iter()
5001            .map(|miss| miss.initialization.num_required_ranges())
5002            .sum();
5003        let mut all_ranges = Vec::with_capacity(num_ranges);
5004        for miss in misses {
5005            all_ranges.extend(miss.initialization.required_ranges().cloned());
5006        }
5007        all_ranges
5008    }
5009
5010    async fn initialize_cache_misses(
5011        &mut self,
5012        misses: Vec<PageCacheMiss>,
5013        buffers: Vec<Bytes>,
5014        cache: Arc<LanceCache>,
5015        io: &Arc<dyn EncodingsIo>,
5016    ) -> Result<()> {
5017        let mut buffers = buffers.into_iter();
5018
5019        // Initialize each missing page from its slice of the buffers, in file
5020        // order. `init_from_buffers` can perform further I/O (e.g. loading blob
5021        // descriptions), so run pages concurrently; peel disjoint `&mut` page
5022        // borrows off the slice to satisfy the borrow checker without a lookup
5023        // table.
5024        let mut tail = self.page_schedulers.as_mut_slice();
5025        let mut next_page = 0;
5026        let mut page_inits = FuturesOrdered::new();
5027        for miss in misses {
5028            let PageCacheMiss {
5029                page_idx,
5030                cache_key,
5031                initialization,
5032            } = miss;
5033            let relative_idx = page_idx
5034                .checked_sub(next_page)
5035                .ok_or_else(|| Error::internal("Page initialization indices are not ordered"))?;
5036            if relative_idx >= tail.len() {
5037                return Err(Error::internal(format!(
5038                    "Page initialization index {page_idx} is out of bounds"
5039                )));
5040            }
5041            let (_, rest) = tail.split_at_mut(relative_idx);
5042            let (page, rest) = rest.split_first_mut().ok_or_else(|| {
5043                Error::internal(format!(
5044                    "Page initialization index {page_idx} is out of bounds"
5045                ))
5046            })?;
5047            tail = rest;
5048            next_page = page_idx + 1;
5049
5050            let count = initialization.num_required_ranges();
5051            let page_buffers = buffers.by_ref().take(count).collect::<Vec<_>>();
5052            if page_buffers.len() != count {
5053                return Err(Error::invalid_input_source(
5054                    format!(
5055                        "Page {page_idx} received {} initialization buffers, expected {count}",
5056                        page_buffers.len()
5057                    )
5058                    .into(),
5059                ));
5060            }
5061            let page_buffers = initialization.with_buffers(page_buffers)?;
5062            let cache = cache.clone();
5063            let scheduler = &mut page.scheduler;
5064            page_inits.push_back(async move {
5065                let cached = scheduler.init_from_buffers(page_buffers, io).await?;
5066                scheduler.try_load(&cached)?;
5067                cache_initialized_page(cache, cache_key, cached).await;
5068                Ok::<(), Error>(())
5069            });
5070        }
5071        if buffers.next().is_some() {
5072            return Err(Error::invalid_input_source(
5073                "Page initialization returned unassigned buffers".into(),
5074            ));
5075        }
5076
5077        page_inits.try_collect::<Vec<()>>().await?;
5078        Ok(())
5079    }
5080
5081    async fn initialize_pages(
5082        &mut self,
5083        requested_ranges: Option<&[Range<u64>]>,
5084        context: &SchedulerContext,
5085    ) -> Result<()> {
5086        let cache = context.cache().clone();
5087        let page_indices = match requested_ranges {
5088            None => (0..self.page_row_ends.len()).collect::<Vec<_>>(),
5089            Some(ranges) => pages_overlapping_ranges(&self.page_row_ends, ranges)?,
5090        };
5091        let positions = self.build_pages(&page_indices)?;
5092        let misses = self
5093            .find_cache_misses(positions.into_iter(), &cache)
5094            .await?;
5095
5096        if misses.is_empty() {
5097            return Ok(());
5098        }
5099
5100        // Concatenate the misses' metadata ranges into one request so adjacent
5101        // ranges coalesce into shared GETs instead of one request per page.
5102        let all_ranges = Self::initialization_ranges(&misses);
5103        let buffers = read_page_initialization_buffers(context.io(), all_ranges).await?;
5104        self.initialize_cache_misses(misses, buffers, cache, context.io())
5105            .await
5106    }
5107}
5108
5109impl StructuralFieldScheduler for StructuralPrimitiveFieldScheduler {
5110    fn initialize<'a>(
5111        &'a mut self,
5112        requested_ranges: Option<&'a [Range<u64>]>,
5113        _filter: &'a FilterExpression,
5114        context: &'a SchedulerContext,
5115    ) -> BoxFuture<'a, Result<()>> {
5116        self.initialize_pages(requested_ranges, context).boxed()
5117    }
5118
5119    fn schedule_ranges<'a>(
5120        &'a self,
5121        ranges: &[Range<u64>],
5122        _filter: &FilterExpression,
5123    ) -> Result<Box<dyn StructuralSchedulingJob + 'a>> {
5124        let ranges = ranges.to_vec();
5125        Ok(Box::new(StructuralPrimitiveFieldSchedulingJob::try_new(
5126            self, ranges,
5127        )?))
5128    }
5129}
5130
5131/// Takes the output from several pages decoders and
5132/// concatenates them.
5133#[derive(Debug)]
5134pub struct StructuralCompositeDecodeArrayTask {
5135    tasks: Vec<Box<dyn DecodePageTask>>,
5136    should_validate: bool,
5137    data_type: DataType,
5138}
5139
5140impl StructuralCompositeDecodeArrayTask {
5141    fn restore_validity(
5142        array: Arc<dyn Array>,
5143        unraveler: &mut CompositeRepDefUnraveler,
5144    ) -> Result<Arc<dyn Array>> {
5145        let validity = unraveler.unravel_validity(array.len())?;
5146        let Some(validity) = validity else {
5147            return Ok(array);
5148        };
5149        if array.data_type() == &DataType::Null {
5150            // We unravel from a null array but we don't add the null buffer because arrow-rs doesn't like it
5151            return Ok(array);
5152        }
5153        if validity.len() != array.len() {
5154            return Err(Error::invalid_input_source(
5155                format!(
5156                    "Structural validity has {} entries for an array with {} values",
5157                    validity.len(),
5158                    array.len()
5159                )
5160                .into(),
5161            ));
5162        }
5163        // SAFETY: The array buffers have already been validated and the null buffer length
5164        // matches the array. We are only attaching the null buffer here.
5165        Ok(make_array(unsafe {
5166            array
5167                .to_data()
5168                .into_builder()
5169                .nulls(Some(validity))
5170                .build_unchecked()
5171        }))
5172    }
5173}
5174
5175impl StructuralDecodeArrayTask for StructuralCompositeDecodeArrayTask {
5176    fn decode(self: Box<Self>) -> Result<DecodedArray> {
5177        let mut arrays = Vec::with_capacity(self.tasks.len());
5178        let mut unravelers = Vec::with_capacity(self.tasks.len());
5179        let mut data_size = 0u64;
5180        for task in self.tasks {
5181            let decoded = task.decode()?;
5182            data_size += decoded.data.data_size();
5183            unravelers.push(decoded.repdef);
5184
5185            let array = make_array(
5186                decoded
5187                    .data
5188                    .into_arrow(self.data_type.clone(), self.should_validate)?,
5189            );
5190
5191            arrays.push(array);
5192        }
5193        let array_refs = arrays.iter().map(|arr| arr.as_ref()).collect::<Vec<_>>();
5194        let array = arrow_select::concat::concat(&array_refs)?;
5195        let mut repdef = CompositeRepDefUnraveler::new(unravelers);
5196
5197        let array = Self::restore_validity(array, &mut repdef)?;
5198
5199        Ok(DecodedArray {
5200            array,
5201            repdef,
5202            data_size,
5203        })
5204    }
5205}
5206
5207/// Returns the total bytes consumed by validity bitmaps of nested child fields
5208/// for `rows` logical rows of `data_type`.
5209///
5210/// Only `FixedSizeList` types have nested children that contribute additional
5211/// bitmaps; all other fixed-width leaf types return 0.
5212fn fsl_child_bitmap_bytes(data_type: &arrow_schema::DataType, rows: u64) -> u64 {
5213    match data_type {
5214        arrow_schema::DataType::FixedSizeList(child_field, dimension) => {
5215            let child_rows = rows * *dimension as u64;
5216            let own = if child_field.is_nullable() {
5217                child_rows.div_ceil(8)
5218            } else {
5219                0
5220            };
5221            own + fsl_child_bitmap_bytes(child_field.data_type(), child_rows)
5222        }
5223        _ => 0,
5224    }
5225}
5226
5227#[derive(Debug)]
5228pub struct StructuralPrimitiveFieldDecoder {
5229    field: Arc<ArrowField>,
5230    page_decoders: VecDeque<Box<dyn StructuralPageDecoder>>,
5231    should_validate: bool,
5232    rows_drained_in_current: u64,
5233}
5234
5235impl StructuralPrimitiveFieldDecoder {
5236    pub fn new(field: &Arc<ArrowField>, should_validate: bool) -> Self {
5237        Self {
5238            field: field.clone(),
5239            page_decoders: VecDeque::new(),
5240            should_validate,
5241            rows_drained_in_current: 0,
5242        }
5243    }
5244
5245    fn decoded_bytes_for_rows(&self, rows: u64) -> lance_core::Result<u64> {
5246        let mut remaining = rows;
5247        let mut total = 0u64;
5248        for (page_num, page) in self.page_decoders.iter().enumerate() {
5249            if remaining == 0 {
5250                break;
5251            }
5252            let available = if page_num == 0 {
5253                page.num_rows().saturating_sub(self.rows_drained_in_current)
5254            } else {
5255                page.num_rows()
5256            };
5257            let take = available.min(remaining);
5258            total += page.decoded_bytes(take)?;
5259            remaining -= take;
5260        }
5261        if remaining > 0 {
5262            return Err(lance_core::Error::not_supported(format!(
5263                "plan_decoded_bytes: queued pages cover only {} of {} requested rows",
5264                rows - remaining,
5265                rows,
5266            )));
5267        }
5268        Ok(total)
5269    }
5270}
5271
5272impl StructuralFieldDecoder for StructuralPrimitiveFieldDecoder {
5273    fn accept_page(&mut self, child: LoadedPageShard) -> Result<()> {
5274        assert!(child.path.is_empty());
5275        self.page_decoders.push_back(child.decoder);
5276        Ok(())
5277    }
5278
5279    fn drain(&mut self, num_rows: u64) -> Result<Box<dyn StructuralDecodeArrayTask>> {
5280        let mut remaining = num_rows;
5281        let mut tasks = Vec::new();
5282        while remaining > 0 {
5283            let queued_pages = self.page_decoders.len();
5284            let Some(cur_page) = self.page_decoders.front_mut() else {
5285                return Err(Error::internal(format!(
5286                    "Primitive decoder missing page decoder while draining field '{}' (data_type={:?}, requested_rows={}, remaining_rows={}, rows_drained_in_current={}, queued_pages={})",
5287                    self.field.name(),
5288                    self.field.data_type(),
5289                    num_rows,
5290                    remaining,
5291                    self.rows_drained_in_current,
5292                    queued_pages
5293                )));
5294            };
5295            let num_in_page = cur_page.num_rows() - self.rows_drained_in_current;
5296            let to_take = num_in_page.min(remaining);
5297
5298            let task = cur_page.drain(to_take)?;
5299            tasks.push(task);
5300
5301            if to_take == num_in_page {
5302                self.page_decoders.pop_front();
5303                self.rows_drained_in_current = 0;
5304            } else {
5305                self.rows_drained_in_current += to_take;
5306            }
5307
5308            remaining -= to_take;
5309        }
5310        Ok(Box::new(StructuralCompositeDecodeArrayTask {
5311            tasks,
5312            should_validate: self.should_validate,
5313            data_type: self.field.data_type().clone(),
5314        }))
5315    }
5316
5317    fn data_type(&self) -> &DataType {
5318        self.field.data_type()
5319    }
5320
5321    fn plan_decoded_bytes(&self, rows_remaining: u64) -> lance_core::Result<[u64; 8]> {
5322        use crate::decoder::CANDIDATE_BATCH_SIZES;
5323        let data_type = self.field.data_type();
5324        let is_nullable = self.field.is_nullable();
5325
5326        // Fixed-width: exact from type metadata, no page inspection needed.
5327        if let Some(byte_width) = data_type.byte_width_opt() {
5328            let mut out = [0u64; 8];
5329            for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5330                let rows = (c as u64).min(rows_remaining);
5331                out[i] = rows * byte_width as u64;
5332                if is_nullable {
5333                    out[i] += rows.div_ceil(8);
5334                }
5335                out[i] += fsl_child_bitmap_bytes(data_type, rows);
5336            }
5337            return Ok(out);
5338        }
5339
5340        // Boolean: bit-packed, 1 bit per value.
5341        if matches!(data_type, DataType::Boolean) {
5342            let mut out = [0u64; 8];
5343            for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5344                let rows = (c as u64).min(rows_remaining);
5345                out[i] = rows.div_ceil(8);
5346                if is_nullable {
5347                    out[i] += rows.div_ceil(8);
5348                }
5349            }
5350            return Ok(out);
5351        }
5352
5353        // Null type: no data bytes.
5354        if matches!(data_type, DataType::Null) {
5355            return Ok([0u64; 8]);
5356        }
5357
5358        // Variable-width: walk page decoders for exact byte counts.
5359        let mut out = [0u64; 8];
5360        for (i, &c) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
5361            let rows = (c as u64).min(rows_remaining);
5362            out[i] = self.decoded_bytes_for_rows(rows)?;
5363            if is_nullable {
5364                out[i] += rows.div_ceil(8);
5365            }
5366        }
5367        Ok(out)
5368    }
5369}
5370
5371/// The serialized representation of full-zip data
5372struct SerializedFullZip {
5373    /// The zipped values buffer
5374    values: LanceBuffer,
5375    /// The repetition index (only present if there is repetition)
5376    repetition_index: Option<LanceBuffer>,
5377}
5378
5379// We align and pad mini-blocks to 8 byte boundaries for two reasons.  First,
5380// to allow us to store a chunk size in 12 bits.
5381//
5382// If we directly record the size in bytes with 12 bits we would be limited to
5383// 4KiB which is too small.  Since we know each mini-block consists of 8 byte
5384// words we can store the # of words instead which gives us 32KiB.
5385//
5386// Second, each chunk in a mini-block is aligned to 8 bytes.  This allows multi-byte
5387// values like offsets to be stored in a mini-block and safely read back out.  It also
5388// helps ensure zero-copy reads in cases where zero-copy is possible (e.g. no decoding
5389// needed).
5390//
5391// Note: by "aligned to 8 bytes" we mean BOTH "aligned to 8 bytes from the start of
5392// the page" and "aligned to 8 bytes from the start of the file."
5393const MINIBLOCK_ALIGNMENT: usize = 8;
5394
5395/// An encoder for primitive (leaf) arrays
5396///
5397/// This encoder is fairly complicated and follows a number of paths depending
5398/// on the data.
5399///
5400/// First, we convert the validity & offsets information into repetition and
5401/// definition levels.  Then we compress the data itself into a single buffer.
5402///
5403/// If the data is narrow then we encode the data in small chunks (each chunk
5404/// should be a few disk sectors and contains a buffer of repetition, a buffer
5405/// of definition, and a buffer of value data).  This approach is called
5406/// "mini-block".  These mini-blocks are stored into a single data buffer.
5407///
5408/// If the data is wide then we zip together the repetition and definition value
5409/// with the value data into a single buffer.  This approach is called "zipped".
5410///
5411/// If there is any repetition information then we create a repetition index
5412///
5413/// In addition, the compression process may create zero or more metadata buffers.
5414/// For example, a dictionary compression will create dictionary metadata.  Any
5415/// mini-block approach has a metadata buffer of block sizes.  This metadata is
5416/// stored in a separate buffer on disk and read at initialization time.
5417///
5418/// TODO: We should concatenate metadata buffers from all pages into a single buffer
5419/// at (roughly) the end of the file so there is, at most, one read per column of
5420/// metadata per file.
5421#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5422pub(crate) enum MiniblockChunkSize {
5423    U16,
5424    U32,
5425}
5426
5427#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5428enum ComplexNullEncoding {
5429    RawLevels,
5430    CompressedLevels,
5431}
5432
5433#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5434enum FixedWidthDictionaryEncoding {
5435    Exclude64Bit,
5436    Include64Bit,
5437}
5438
5439trait PrimitivePageEncodingBehavior: Send + Sync + Debug {
5440    fn validate_field(&self, _field: &Field, _metadata: &HashMap<String, String>) -> Result<()> {
5441        Ok(())
5442    }
5443
5444    fn try_plan_pages(
5445        &self,
5446        _ctx: &PrimitivePlanContext<'_>,
5447        _arrays: &[ArrayRef],
5448        _normalized: &NormalizedStructuralPlan,
5449        _row_number: u64,
5450        _num_rows: u64,
5451        _num_values: u64,
5452    ) -> Result<Option<Vec<PrimitivePageData>>> {
5453        Ok(None)
5454    }
5455
5456    fn try_encode_page(
5457        &self,
5458        _ctx: &PrimitiveEncodeContext,
5459        page: PrimitivePageData,
5460    ) -> Result<PrimitiveEncodeAttempt> {
5461        Ok(PrimitiveEncodeAttempt::Unhandled(page))
5462    }
5463}
5464
5465/// One executable primitive-page behavior selected by an exact file
5466/// composition.
5467#[derive(Debug, Clone)]
5468pub struct PrimitivePageEncoding {
5469    behavior: Arc<dyn PrimitivePageEncodingBehavior>,
5470}
5471
5472impl PrimitivePageEncoding {
5473    /// Reject an explicit request for sparse structural encoding.
5474    pub fn reject_sparse() -> Self {
5475        Self {
5476            behavior: Arc::new(RejectSparsePrimitiveEncoding),
5477        }
5478    }
5479
5480    /// Encode constant non-null values as a constant page when applicable.
5481    pub fn constant() -> Self {
5482        Self {
5483            behavior: Arc::new(ConstantPrimitiveEncoding),
5484        }
5485    }
5486
5487    /// Plan and encode sparse structural pages when applicable.
5488    pub fn sparse(compression: Arc<dyn CompressionStrategy>) -> Self {
5489        Self {
5490            behavior: Arc::new(SparsePrimitiveEncoding { compression }),
5491        }
5492    }
5493
5494    /// Encode dense pages with the original u16 miniblock grammar.
5495    pub fn dense_u16(compression: Arc<dyn CompressionStrategy>) -> Self {
5496        Self {
5497            behavior: Arc::new(DenseU16PrimitiveEncoding { compression }),
5498        }
5499    }
5500
5501    /// Encode dense pages with the u32 miniblock grammar.
5502    pub fn dense_u32(compression: Arc<dyn CompressionStrategy>) -> Self {
5503        Self {
5504            behavior: Arc::new(DenseU32PrimitiveEncoding { compression }),
5505        }
5506    }
5507}
5508
5509#[derive(Debug)]
5510struct RejectSparsePrimitiveEncoding;
5511
5512#[derive(Debug)]
5513struct ConstantPrimitiveEncoding;
5514
5515#[derive(Debug)]
5516struct SparsePrimitiveEncoding {
5517    compression: Arc<dyn CompressionStrategy>,
5518}
5519
5520#[derive(Debug)]
5521struct DenseU16PrimitiveEncoding {
5522    compression: Arc<dyn CompressionStrategy>,
5523}
5524
5525#[derive(Debug)]
5526struct DenseU32PrimitiveEncoding {
5527    compression: Arc<dyn CompressionStrategy>,
5528}
5529
5530pub struct PrimitiveStructuralEncoder {
5531    // Accumulates arrays until we have enough data to justify a disk page
5532    accumulation_queue: AccumulationQueue,
5533
5534    keep_original_array: bool,
5535    accumulated_repdefs: Vec<RepDefBuilder>,
5536    page_encodings: Arc<[PrimitivePageEncoding]>,
5537    column_index: u32,
5538    field: Field,
5539    encoding_metadata: Arc<HashMap<String, String>>,
5540}
5541
5542struct CompressedLevelsChunk {
5543    data: LanceBuffer,
5544    num_levels: u16,
5545}
5546
5547struct CompressedLevels {
5548    data: Vec<CompressedLevelsChunk>,
5549    compression: CompressiveEncoding,
5550    rep_index: Option<LanceBuffer>,
5551}
5552
5553struct SerializedMiniBlockPage {
5554    num_buffers: u64,
5555    data: LanceBuffer,
5556    metadata: LanceBuffer,
5557}
5558
5559#[derive(Debug, Clone, Copy)]
5560struct DictEncodingBudget {
5561    max_dict_entries: u32,
5562    max_encoded_size: usize,
5563}
5564
5565enum PrimitivePageStructure {
5566    Dense {
5567        repdef: SerializedRepDefs,
5568        single_row_miniblock_repdef_levels: Option<u64>,
5569    },
5570    Sparse {
5571        plan: sparse::SparseStructuralPlan,
5572        prepared_values: Option<sparse::writer::PreparedSparseValues>,
5573    },
5574}
5575
5576// A primitive page after structural encoding selection and optional dense splitting.
5577struct PrimitivePageData {
5578    // Arrow leaf arrays that contain this page's visible values.
5579    arrays: Vec<ArrayRef>,
5580    // Structural representation aligned to this page.
5581    structure: PrimitivePageStructure,
5582    // Top-level row number of the first row in this page.
5583    row_number: u64,
5584    // Number of top-level rows in this page.
5585    num_rows: u64,
5586}
5587
5588struct PrimitivePlanContext<'a> {
5589    column_idx: u32,
5590    field: &'a Field,
5591    encoding_metadata: &'a HashMap<String, String>,
5592}
5593
5594enum PrimitiveEncodeAttempt {
5595    Encoded(EncodedPage),
5596    Unhandled(PrimitivePageData),
5597}
5598
5599// Immutable encoder state shared by per-page encode tasks.
5600//
5601// Cloning this only clones Arc-backed configuration and field metadata.  Page data
5602// stays in PrimitivePageData and is moved into exactly one task.
5603#[derive(Clone)]
5604struct PrimitiveEncodeContext {
5605    // Column being encoded.
5606    column_idx: u32,
5607    field: Field,
5608    encoding_metadata: Arc<HashMap<String, String>>,
5609    is_simple_validity: bool,
5610    has_repdef_info: bool,
5611}
5612
5613impl PrimitiveStructuralEncoder {
5614    pub fn try_new(
5615        options: &EncodingOptions,
5616        page_encodings: Arc<[PrimitivePageEncoding]>,
5617        column_index: u32,
5618        field: Field,
5619        encoding_metadata: Arc<HashMap<String, String>>,
5620    ) -> Result<Self> {
5621        for page_encoding in page_encodings.iter() {
5622            page_encoding
5623                .behavior
5624                .validate_field(&field, &encoding_metadata)?;
5625        }
5626        Ok(Self {
5627            accumulation_queue: AccumulationQueue::new(
5628                options.cache_bytes_per_column,
5629                column_index,
5630                options.keep_original_array,
5631            ),
5632            keep_original_array: options.keep_original_array,
5633            accumulated_repdefs: Vec::new(),
5634            column_index,
5635            page_encodings,
5636            field,
5637            encoding_metadata,
5638        })
5639    }
5640
5641    fn encode_page(
5642        page_encodings: &[PrimitivePageEncoding],
5643        ctx: &PrimitiveEncodeContext,
5644        mut page: PrimitivePageData,
5645    ) -> Result<EncodedPage> {
5646        for page_encoding in page_encodings {
5647            match page_encoding.behavior.try_encode_page(ctx, page)? {
5648                PrimitiveEncodeAttempt::Encoded(page) => return Ok(page),
5649                PrimitiveEncodeAttempt::Unhandled(unhandled) => page = unhandled,
5650            }
5651        }
5652        Err(Error::invalid_input_source(
5653            format!(
5654                "No primitive page encoding atom supports field '{}'",
5655                ctx.field.name
5656            )
5657            .into(),
5658        ))
5659    }
5660
5661    // TODO: This is a heuristic we may need to tune at some point
5662    //
5663    // As data gets narrow then the "zipping" process gets too expensive
5664    //   and we prefer mini-block
5665    // As data gets wide then the # of values per block shrinks (very wide)
5666    //   data doesn't even fit in a mini-block and the block overhead gets
5667    //   too large and we prefer zipped.
5668    fn is_narrow(data_block: &DataBlock) -> bool {
5669        const MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE: u64 = 256;
5670
5671        if let Some(max_len_array) = data_block.get_stat(Stat::MaxLength) {
5672            let max_len_array = max_len_array
5673                .as_any()
5674                .downcast_ref::<PrimitiveArray<UInt64Type>>()
5675                .unwrap();
5676            if max_len_array.value(0) < MINIBLOCK_MAX_BYTE_LENGTH_PER_VALUE {
5677                return true;
5678            }
5679        }
5680        false
5681    }
5682
5683    fn prefers_miniblock(
5684        data_block: &DataBlock,
5685        encoding_metadata: &HashMap<String, String>,
5686    ) -> bool {
5687        // If the user specifically requested miniblock then use it
5688        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
5689            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_MINIBLOCK;
5690        }
5691        // Otherwise only use miniblock if it is narrow
5692        Self::is_narrow(data_block)
5693    }
5694
5695    fn prefers_fullzip(encoding_metadata: &HashMap<String, String>) -> bool {
5696        // Fullzip is the backup option so the only reason we wouldn't use it is if the
5697        // user specifically requested not to use it (in which case we're probably going
5698        // to emit an error)
5699        if let Some(user_requested) = encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY) {
5700            return user_requested.to_lowercase() == STRUCTURAL_ENCODING_FULLZIP;
5701        }
5702        true
5703    }
5704
5705    // Converts value data, repetition levels, and definition levels into a single
5706    // buffer of mini-blocks.  In addition, creates a buffer of mini-block metadata
5707    // which tells us the size of each block.  Finally, if repetition is present then
5708    // we also create a buffer for the repetition index.
5709    //
5710    // Each chunk is serialized as:
5711    // | num_bufs (1 byte) | buf_lens (2 bytes per buffer) | P | buf0 | P | buf1 | ... | bufN | P |
5712    //
5713    // P - Padding inserted to ensure each buffer is 8-byte aligned and the buffer size is a multiple
5714    //     of 8 bytes (so that the next chunk is 8-byte aligned).
5715    //
5716    // Each block has a u16 word of metadata.  The upper 12 bits contain the
5717    // # of 8-byte words in the block (if the block does not fill the final word
5718    // then up to 7 bytes of padding are added).  The lower 4 bits describe the log_2
5719    // number of values (e.g. if there are 1024 then the lower 4 bits will be
5720    // 0xA)  All blocks except the last must have power-of-two number of values.
5721    // This not only makes metadata smaller but it makes decoding easier since
5722    // batch sizes are typically a power of 2.  4 bits would allow us to express
5723    // up to 32Ki values.
5724    //
5725    // This means blocks can have 1 to 32Ki values and 8 - 32Ki bytes.
5726    //
5727    // All metadata words are serialized (as little endian) into a single buffer
5728    // of metadata values.
5729    //
5730    // If there is repetition then we also create a repetition index.  This is a
5731    // single buffer of integer vectors (stored in row major order).  There is one
5732    // entry for each chunk.  The size of the vector is based on the depth of random
5733    // access we want to support.
5734    //
5735    // A vector of size 2 is the minimum and will support row-based random access (e.g.
5736    // "take the 57th row").  A vector of size 3 will support 1 level of nested access
5737    // (e.g. "take the 3rd item in the 57th row").  A vector of size 4 will support 2
5738    // levels of nested access and so on.
5739    //
5740    // The first number in the vector is the number of top-level rows that complete in
5741    // the chunk.  The second number is the number of second-level rows that complete
5742    // after the final top-level row completed (or beginning of the chunk if no top-level
5743    // row completes in the chunk).  And so on.  The final number in the vector is always
5744    // the number of leftover items not covered by earlier entries in the vector.
5745    //
5746    // Currently we are limited to 0 levels of nested access but that will change in the
5747    // future.
5748    //
5749    // The repetition index and the chunk metadata are read at initialization time and
5750    // cached in memory.
5751    fn serialize_miniblocks(
5752        miniblocks: MiniBlockCompressed,
5753        rep: Option<Vec<CompressedLevelsChunk>>,
5754        def: Option<Vec<CompressedLevelsChunk>>,
5755        miniblock_chunk_size: MiniblockChunkSize,
5756    ) -> Result<SerializedMiniBlockPage> {
5757        let bytes_rep = rep
5758            .as_ref()
5759            .map(|rep| rep.iter().map(|r| r.data.len()).sum::<usize>())
5760            .unwrap_or(0);
5761        let bytes_def = def
5762            .as_ref()
5763            .map(|def| def.iter().map(|d| d.data.len()).sum::<usize>())
5764            .unwrap_or(0);
5765        let bytes_data = miniblocks.data.iter().map(|d| d.len()).sum::<usize>();
5766        let mut num_buffers = miniblocks.data.len();
5767        if rep.is_some() {
5768            num_buffers += 1;
5769        }
5770        if def.is_some() {
5771            num_buffers += 1;
5772        }
5773        // 2 bytes for the length of each buffer and up to 7 bytes of padding per buffer
5774        let max_extra = 9 * num_buffers;
5775        let mut data_buffer = Vec::with_capacity(bytes_rep + bytes_def + bytes_data + max_extra);
5776        let chunk_size_bytes = match miniblock_chunk_size {
5777            MiniblockChunkSize::U16 => 2,
5778            MiniblockChunkSize::U32 => 4,
5779        };
5780        let mut meta_buffer = Vec::with_capacity(miniblocks.chunks.len() * chunk_size_bytes);
5781
5782        let mut rep_iter = rep.map(|r| r.into_iter());
5783        let mut def_iter = def.map(|d| d.into_iter());
5784
5785        let mut buffer_offsets = vec![0; miniblocks.data.len()];
5786        for chunk in miniblocks.chunks {
5787            let start_pos = data_buffer.len();
5788            // Start of chunk should be aligned
5789            debug_assert_eq!(start_pos % MINIBLOCK_ALIGNMENT, 0);
5790
5791            let rep = rep_iter.as_mut().map(|r| r.next().unwrap());
5792            let def = def_iter.as_mut().map(|d| d.next().unwrap());
5793
5794            // Write the number of levels, or 0 if there is no rep/def
5795            let num_levels = rep
5796                .as_ref()
5797                .map(|r| r.num_levels)
5798                .unwrap_or(def.as_ref().map(|d| d.num_levels).unwrap_or(0));
5799            data_buffer.extend_from_slice(&num_levels.to_le_bytes());
5800
5801            // Write the buffer lengths
5802            if let Some(rep) = rep.as_ref() {
5803                let bytes_rep = u16::try_from(rep.data.len()).map_err(|_| {
5804                    Error::internal(format!(
5805                        "Repetition buffer size ({} bytes) too large",
5806                        rep.data.len()
5807                    ))
5808                })?;
5809                data_buffer.extend_from_slice(&bytes_rep.to_le_bytes());
5810            }
5811            if let Some(def) = def.as_ref() {
5812                let bytes_def = u16::try_from(def.data.len()).map_err(|_| {
5813                    Error::internal(format!(
5814                        "Definition buffer size ({} bytes) too large",
5815                        def.data.len()
5816                    ))
5817                })?;
5818                data_buffer.extend_from_slice(&bytes_def.to_le_bytes());
5819            }
5820
5821            if miniblock_chunk_size == MiniblockChunkSize::U32 {
5822                for &buffer_size in &chunk.buffer_sizes {
5823                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
5824                }
5825            } else {
5826                for &buffer_size in &chunk.buffer_sizes {
5827                    let buffer_size = u16::try_from(buffer_size).map_err(|_| {
5828                        Error::internal(format!(
5829                            "Mini-block buffer size ({} bytes) too large for 16-bit metadata",
5830                            buffer_size
5831                        ))
5832                    })?;
5833                    data_buffer.extend_from_slice(&buffer_size.to_le_bytes());
5834                }
5835            }
5836
5837            // Pad
5838            let add_padding = |data_buffer: &mut Vec<u8>| {
5839                let pad = pad_bytes::<MINIBLOCK_ALIGNMENT>(data_buffer.len());
5840                data_buffer.extend(iter::repeat_n(FILL_BYTE, pad));
5841            };
5842            add_padding(&mut data_buffer);
5843
5844            // Write the buffers themselves
5845            if let Some(rep) = rep.as_ref() {
5846                data_buffer.extend_from_slice(&rep.data);
5847                add_padding(&mut data_buffer);
5848            }
5849            if let Some(def) = def.as_ref() {
5850                data_buffer.extend_from_slice(&def.data);
5851                add_padding(&mut data_buffer);
5852            }
5853            for (buffer_size, (buffer, buffer_offset)) in chunk
5854                .buffer_sizes
5855                .iter()
5856                .zip(miniblocks.data.iter().zip(buffer_offsets.iter_mut()))
5857            {
5858                let start = *buffer_offset;
5859                let end = start + *buffer_size as usize;
5860                *buffer_offset += *buffer_size as usize;
5861                data_buffer.extend_from_slice(&buffer[start..end]);
5862                add_padding(&mut data_buffer);
5863            }
5864
5865            let chunk_bytes = data_buffer.len() - start_pos;
5866            let max_chunk_size = match miniblock_chunk_size {
5867                MiniblockChunkSize::U16 => 32 * 1024,
5868                MiniblockChunkSize::U32 => 1_u64 << 31,
5869            };
5870            if chunk_bytes == 0 || chunk_bytes as u64 > max_chunk_size {
5871                return Err(Error::internal(format!(
5872                    "Mini-block chunk size {} bytes exceeds the {} byte metadata limit",
5873                    chunk_bytes, max_chunk_size
5874                )));
5875            }
5876            if chunk_bytes % MINIBLOCK_ALIGNMENT != 0 {
5877                return Err(Error::internal(format!(
5878                    "Mini-block chunk size {} bytes is not aligned to {} bytes",
5879                    chunk_bytes, MINIBLOCK_ALIGNMENT
5880                )));
5881            }
5882            if chunk.log_num_values > 15 {
5883                return Err(Error::internal(format!(
5884                    "Mini-block log_num_values {} exceeds the 4-bit metadata limit",
5885                    chunk.log_num_values
5886                )));
5887            }
5888            // We subtract 1 here from chunk_bytes because we want to be able to express
5889            // a size of 32KiB and not (32Ki - 8)B which is what we'd get otherwise with
5890            // 0xFFF
5891            let divided_bytes = chunk_bytes / MINIBLOCK_ALIGNMENT;
5892            let divided_bytes_minus_one = (divided_bytes - 1) as u64;
5893
5894            let metadata = (divided_bytes_minus_one << 4) | chunk.log_num_values as u64;
5895            if miniblock_chunk_size == MiniblockChunkSize::U32 {
5896                meta_buffer.extend_from_slice(&(metadata as u32).to_le_bytes());
5897            } else {
5898                meta_buffer.extend_from_slice(&(metadata as u16).to_le_bytes());
5899            }
5900        }
5901
5902        let data_buffer = LanceBuffer::from(data_buffer);
5903        let metadata_buffer = LanceBuffer::from(meta_buffer);
5904
5905        Ok(SerializedMiniBlockPage {
5906            num_buffers: miniblocks.data.len() as u64,
5907            data: data_buffer,
5908            metadata: metadata_buffer,
5909        })
5910    }
5911
5912    /// Compresses a buffer of levels into chunks
5913    ///
5914    /// If these are repetition levels then we also calculate the repetition index here (that
5915    /// is the third return value)
5916    fn compress_levels(
5917        mut levels: RepDefSlicer<'_>,
5918        num_elements: u64,
5919        compression_strategy: &dyn CompressionStrategy,
5920        chunks: &[MiniBlockChunk],
5921        // This will be 0 if we are compressing def levels
5922        max_rep: u16,
5923    ) -> Result<CompressedLevels> {
5924        let mut rep_index = if max_rep > 0 {
5925            Vec::with_capacity(chunks.len())
5926        } else {
5927            vec![]
5928        };
5929        // Make the levels into a FixedWidth data block
5930        let num_levels = levels.num_levels() as u64;
5931        let levels_buf = levels.all_levels().clone();
5932
5933        let mut fixed_width_block = FixedWidthDataBlock {
5934            data: levels_buf,
5935            bits_per_value: 16,
5936            num_values: num_levels,
5937            block_info: BlockInfo::new(),
5938        };
5939        // Compute statistics to enable optimal compression for rep/def levels
5940        fixed_width_block.compute_stat();
5941
5942        let levels_block = DataBlock::FixedWidth(fixed_width_block);
5943        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
5944        // Pick a block compressor
5945        let compressor =
5946            compression_strategy.create_block_compressor(&levels_field, &levels_block)?;
5947        let mut compressor_desc = None;
5948        // Compress blocks of levels (sized according to the chunks)
5949        let mut level_chunks = Vec::with_capacity(chunks.len());
5950        let mut values_counter = 0;
5951        for (chunk_idx, chunk) in chunks.iter().enumerate() {
5952            let chunk_num_values = chunk.num_values(values_counter, num_elements);
5953            debug_assert!(chunk_num_values > 0);
5954            values_counter += chunk_num_values;
5955            let chunk_levels = if chunk_idx < chunks.len() - 1 {
5956                levels.slice_next(chunk_num_values as usize)
5957            } else {
5958                levels.slice_rest()
5959            };
5960            let num_chunk_levels = (chunk_levels.len() / 2) as u64;
5961            if max_rep > 0 {
5962                // If max_rep > 0 then we are working with rep levels and we need
5963                // to calculate the repetition index.  The repetition index for a
5964                // chunk is currently 2 values (in the future it may be more).
5965                //
5966                // The first value is the number of rows that _finish_ in the
5967                // chunk.
5968                //
5969                // The second value is the number of "leftovers" after the last
5970                // finished row in the chunk.
5971                let rep_values = chunk_levels.borrow_to_typed_slice::<u16>();
5972                let rep_values = rep_values.as_ref();
5973
5974                // We skip 1 here because a max_rep at spot 0 doesn't count as a finished list (we
5975                // will count it in the previous chunk)
5976                let mut num_rows = rep_values.iter().skip(1).filter(|v| **v == max_rep).count();
5977                let num_leftovers = if chunk_idx < chunks.len() - 1 {
5978                    rep_values
5979                        .iter()
5980                        .rev()
5981                        .position(|v| *v == max_rep)
5982                        // # of leftovers includes the max_rep spot
5983                        .map(|pos| pos + 1)
5984                        .unwrap_or(rep_values.len())
5985                } else {
5986                    // Last chunk can't have leftovers
5987                    0
5988                };
5989
5990                if chunk_idx != 0 && rep_values.first() == Some(&max_rep) {
5991                    // This chunk starts with a new row and so, if we thought we had leftovers
5992                    // in the previous chunk, we were mistaken
5993                    // TODO: Can use unchecked here
5994                    let rep_len = rep_index.len();
5995                    if rep_index[rep_len - 1] != 0 {
5996                        // We thought we had leftovers but that was actually a full row
5997                        rep_index[rep_len - 2] += 1;
5998                        rep_index[rep_len - 1] = 0;
5999                    }
6000                }
6001
6002                if chunk_idx == chunks.len() - 1 {
6003                    // The final list
6004                    num_rows += 1;
6005                }
6006                rep_index.push(num_rows as u64);
6007                rep_index.push(num_leftovers as u64);
6008            }
6009            let mut chunk_fixed_width = FixedWidthDataBlock {
6010                data: chunk_levels,
6011                bits_per_value: 16,
6012                num_values: num_chunk_levels,
6013                block_info: BlockInfo::new(),
6014            };
6015            chunk_fixed_width.compute_stat();
6016            let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width);
6017            let (compressed_levels, chunk_compressor_desc) =
6018                compressor.compress(chunk_levels_block)?;
6019            if let Some(compressor_desc) = compressor_desc.as_ref() {
6020                if compressor_desc != &chunk_compressor_desc {
6021                    return Err(Error::internal(
6022                        "Rep/def block compressor changed encoding between chunks".to_string(),
6023                    ));
6024                }
6025            } else {
6026                compressor_desc = Some(chunk_compressor_desc);
6027            }
6028            let compressed_levels = compressed_levels.ok_or_else(|| {
6029                Error::internal(
6030                    "Rep/def block compressor selected a metadata-only codec".to_string(),
6031                )
6032            })?;
6033            let num_levels = u16::try_from(num_chunk_levels).map_err(|_| {
6034                Error::invalid_input_source(
6035                    format!(
6036                        "Mini-block cannot encode {} rep/def levels in one chunk. \
6037                         This usually means a top-level row contains too much nested structure \
6038                         for the current layout.",
6039                        num_chunk_levels
6040                    )
6041                    .into(),
6042                )
6043            })?;
6044            level_chunks.push(CompressedLevelsChunk {
6045                data: compressed_levels,
6046                num_levels,
6047            });
6048        }
6049        debug_assert_eq!(levels.num_levels_remaining(), 0);
6050        let rep_index = if rep_index.is_empty() {
6051            None
6052        } else {
6053            Some(LanceBuffer::reinterpret_vec(rep_index))
6054        };
6055        Ok(CompressedLevels {
6056            data: level_chunks,
6057            compression: compressor_desc.ok_or_else(|| {
6058                Error::internal("Rep/def compression produced no chunks".to_string())
6059            })?,
6060            rep_index,
6061        })
6062    }
6063
6064    fn encode_simple_all_null(
6065        column_idx: u32,
6066        num_rows: u64,
6067        row_number: u64,
6068    ) -> Result<EncodedPage> {
6069        let description =
6070            ProtobufUtils21::constant_layout(&[DefinitionInterpretation::NullableItem], None);
6071        Ok(EncodedPage {
6072            column_idx,
6073            data: vec![],
6074            description: PageEncoding::Structural(description),
6075            num_rows,
6076            row_number,
6077        })
6078    }
6079
6080    fn encode_complex_all_null_vals(
6081        data: &Arc<[u16]>,
6082        compression_strategy: &dyn CompressionStrategy,
6083    ) -> Result<(LanceBuffer, pb21::CompressiveEncoding)> {
6084        let buffer = LanceBuffer::reinterpret_slice(data.clone());
6085        let mut fixed_width_block = FixedWidthDataBlock {
6086            data: buffer,
6087            bits_per_value: 16,
6088            num_values: data.len() as u64,
6089            block_info: BlockInfo::new(),
6090        };
6091        fixed_width_block.compute_stat();
6092
6093        let levels_block = DataBlock::FixedWidth(fixed_width_block);
6094        let levels_field = Field::new_arrow("", DataType::UInt16, false)?;
6095        let (compressed_buffer, encoding) =
6096            compress_required_block(compression_strategy, &levels_field, levels_block)?;
6097        Ok((compressed_buffer, encoding))
6098    }
6099
6100    // Encodes a page where all values are null but we have rep/def
6101    // information that we need to store (e.g. to distinguish between
6102    // different kinds of null)
6103    fn encode_complex_all_null(
6104        column_idx: u32,
6105        repdef: crate::repdef::SerializedRepDefs,
6106        row_number: u64,
6107        num_rows: u64,
6108        complex_null_encoding: ComplexNullEncoding,
6109        compression_strategy: &dyn CompressionStrategy,
6110    ) -> Result<EncodedPage> {
6111        if complex_null_encoding == ComplexNullEncoding::RawLevels {
6112            let rep_bytes = if let Some(rep) = repdef.repetition_levels.as_ref() {
6113                LanceBuffer::reinterpret_slice(rep.clone())
6114            } else {
6115                LanceBuffer::empty()
6116            };
6117
6118            let def_bytes = if let Some(def) = repdef.definition_levels.as_ref() {
6119                LanceBuffer::reinterpret_slice(def.clone())
6120            } else {
6121                LanceBuffer::empty()
6122            };
6123
6124            let description = ProtobufUtils21::constant_layout(&repdef.def_meaning, None);
6125            return Ok(EncodedPage {
6126                column_idx,
6127                data: vec![rep_bytes, def_bytes],
6128                description: PageEncoding::Structural(description),
6129                num_rows,
6130                row_number,
6131            });
6132        }
6133
6134        let (rep_bytes, rep_encoding, num_rep_values) = if let Some(rep) =
6135            repdef.repetition_levels.as_ref()
6136        {
6137            let num_values = rep.len() as u64;
6138            let (buffer, encoding) = Self::encode_complex_all_null_vals(rep, compression_strategy)?;
6139            (buffer, Some(encoding), num_values)
6140        } else {
6141            (LanceBuffer::empty(), None, 0)
6142        };
6143
6144        let (def_bytes, def_encoding, num_def_values) = if let Some(def) =
6145            repdef.definition_levels.as_ref()
6146        {
6147            let num_values = def.len() as u64;
6148            let (buffer, encoding) = Self::encode_complex_all_null_vals(def, compression_strategy)?;
6149            (buffer, Some(encoding), num_values)
6150        } else {
6151            (LanceBuffer::empty(), None, 0)
6152        };
6153
6154        let description = ProtobufUtils21::compressed_all_null_constant_layout(
6155            &repdef.def_meaning,
6156            rep_encoding,
6157            def_encoding,
6158            num_rep_values,
6159            num_def_values,
6160        );
6161        Ok(EncodedPage {
6162            column_idx,
6163            data: vec![rep_bytes, def_bytes],
6164            description: PageEncoding::Structural(description),
6165            num_rows,
6166            row_number,
6167        })
6168    }
6169
6170    fn leaf_validity(
6171        repdef: &crate::repdef::SerializedRepDefs,
6172        num_values: usize,
6173    ) -> Result<Option<BooleanBuffer>> {
6174        let rep = repdef
6175            .repetition_levels
6176            .as_ref()
6177            .map(|rep| rep.as_ref().to_vec());
6178        let def = repdef
6179            .definition_levels
6180            .as_ref()
6181            .map(|def| def.as_ref().to_vec());
6182        let mut unraveler = RepDefUnraveler::new(
6183            rep,
6184            def,
6185            repdef.def_meaning.clone().into(),
6186            num_values as u64,
6187        );
6188        if unraveler.is_all_valid() {
6189            return Ok(None);
6190        }
6191        let mut validity = BooleanBufferBuilder::new(num_values);
6192        unraveler.unravel_validity(&mut validity)?;
6193        Ok(Some(validity.finish()))
6194    }
6195
6196    fn is_constant_values(
6197        arrays: &[ArrayRef],
6198        scalar: &ArrayRef,
6199        validity: Option<&BooleanBuffer>,
6200    ) -> Result<bool> {
6201        debug_assert_eq!(scalar.len(), 1);
6202        debug_assert_eq!(scalar.null_count(), 0);
6203
6204        match scalar.data_type() {
6205            DataType::Boolean => {
6206                let mut global_idx = 0usize;
6207                let scalar_val = scalar.as_boolean().value(0);
6208                for arr in arrays {
6209                    let bool_arr = arr.as_boolean();
6210                    for i in 0..arr.len() {
6211                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6212                        global_idx += 1;
6213                        if !is_valid {
6214                            continue;
6215                        }
6216                        if bool_arr.value(i) != scalar_val {
6217                            return Ok(false);
6218                        }
6219                    }
6220                }
6221                Ok(true)
6222            }
6223            DataType::Utf8 => Self::is_constant_utf8::<i32>(arrays, scalar, validity),
6224            DataType::LargeUtf8 => Self::is_constant_utf8::<i64>(arrays, scalar, validity),
6225            DataType::Binary => Self::is_constant_binary::<i32>(arrays, scalar, validity),
6226            DataType::LargeBinary => Self::is_constant_binary::<i64>(arrays, scalar, validity),
6227            data_type => {
6228                let mut global_idx = 0usize;
6229                let Some(byte_width) = data_type.byte_width_opt() else {
6230                    return Ok(false);
6231                };
6232                let scalar_data = scalar.to_data();
6233                if scalar_data.buffers().len() != 1 || !scalar_data.child_data().is_empty() {
6234                    return Ok(false);
6235                }
6236                let scalar_bytes = scalar_data.buffers()[0].as_slice();
6237                if scalar_bytes.len() != byte_width {
6238                    return Ok(false);
6239                }
6240
6241                for arr in arrays {
6242                    let data = arr.to_data();
6243                    if data.buffers().is_empty() {
6244                        return Ok(false);
6245                    }
6246                    let buf = data.buffers()[0].as_slice();
6247                    let base = data.offset();
6248                    for i in 0..arr.len() {
6249                        let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6250                        global_idx += 1;
6251                        if !is_valid {
6252                            continue;
6253                        }
6254                        let start = (base + i) * byte_width;
6255                        if buf[start..start + byte_width] != scalar_bytes[..] {
6256                            return Ok(false);
6257                        }
6258                    }
6259                }
6260                Ok(true)
6261            }
6262        }
6263    }
6264
6265    fn is_constant_utf8<O: arrow_array::OffsetSizeTrait>(
6266        arrays: &[ArrayRef],
6267        scalar: &ArrayRef,
6268        validity: Option<&BooleanBuffer>,
6269    ) -> Result<bool> {
6270        debug_assert_eq!(scalar.len(), 1);
6271        let scalar_val = scalar.as_string::<O>().value(0).as_bytes();
6272        let mut global_idx = 0usize;
6273        for arr in arrays {
6274            let str_arr = arr.as_string::<O>();
6275            for i in 0..arr.len() {
6276                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6277                global_idx += 1;
6278                if !is_valid {
6279                    continue;
6280                }
6281                if str_arr.value(i).as_bytes() != scalar_val {
6282                    return Ok(false);
6283                }
6284            }
6285        }
6286        Ok(true)
6287    }
6288
6289    fn is_constant_binary<O: arrow_array::OffsetSizeTrait>(
6290        arrays: &[ArrayRef],
6291        scalar: &ArrayRef,
6292        validity: Option<&BooleanBuffer>,
6293    ) -> Result<bool> {
6294        debug_assert_eq!(scalar.len(), 1);
6295        let scalar_val = scalar.as_binary::<O>().value(0);
6296        let mut global_idx = 0usize;
6297        for arr in arrays {
6298            let bin_arr = arr.as_binary::<O>();
6299            for i in 0..arr.len() {
6300                let is_valid = validity.map(|v| v.value(global_idx)).unwrap_or(true);
6301                global_idx += 1;
6302                if !is_valid {
6303                    continue;
6304                }
6305                if bin_arr.value(i) != scalar_val {
6306                    return Ok(false);
6307                }
6308            }
6309        }
6310        Ok(true)
6311    }
6312
6313    fn find_constant_scalar(
6314        arrays: &[ArrayRef],
6315        validity: Option<&BooleanBuffer>,
6316    ) -> Result<Option<ArrayRef>> {
6317        if arrays.is_empty() {
6318            return Ok(None);
6319        }
6320
6321        let global_scalar_idx = if let Some(validity) = validity {
6322            let Some(idx) = (0..validity.len()).find(|&i| validity.value(i)) else {
6323                return Ok(None);
6324            };
6325            idx
6326        } else {
6327            0
6328        };
6329
6330        let mut idx_remaining = global_scalar_idx;
6331        let mut scalar_arr_idx = 0usize;
6332        while scalar_arr_idx < arrays.len() {
6333            let len = arrays[scalar_arr_idx].len();
6334            if idx_remaining < len {
6335                break;
6336            }
6337            idx_remaining -= len;
6338            scalar_arr_idx += 1;
6339        }
6340
6341        if scalar_arr_idx >= arrays.len() {
6342            return Ok(None);
6343        }
6344
6345        let scalar =
6346            lance_arrow::scalar::extract_scalar_value(&arrays[scalar_arr_idx], idx_remaining)?;
6347        if scalar.null_count() != 0 {
6348            return Ok(None);
6349        }
6350        if !Self::is_constant_values(arrays, &scalar, validity)? {
6351            return Ok(None);
6352        }
6353        Ok(Some(scalar))
6354    }
6355
6356    fn resolve_dict_values_compression_metadata(
6357        field_metadata: &HashMap<String, String>,
6358        env_compression: Option<String>,
6359        env_compression_level: Option<String>,
6360    ) -> HashMap<String, String> {
6361        let mut metadata = HashMap::new();
6362
6363        let compression = field_metadata
6364            .get(DICT_VALUES_COMPRESSION_META_KEY)
6365            .cloned()
6366            .or(env_compression)
6367            .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string());
6368        metadata.insert(COMPRESSION_META_KEY.to_string(), compression);
6369
6370        if let Some(compression_level) = field_metadata
6371            .get(DICT_VALUES_COMPRESSION_LEVEL_META_KEY)
6372            .cloned()
6373            .or(env_compression_level)
6374        {
6375            metadata.insert(COMPRESSION_LEVEL_META_KEY.to_string(), compression_level);
6376        }
6377
6378        metadata
6379    }
6380
6381    fn build_dict_values_compressor_field(field: &Field) -> Result<Field> {
6382        // This is an internal synthetic field used only to feed metadata into
6383        // `create_block_compressor` for dictionary values. The concrete type/name here
6384        // are not semantically meaningful; we rely on explicit metadata below to control
6385        // general compression selection for dictionary values.
6386        let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?;
6387        dict_values_field.metadata = Self::resolve_dict_values_compression_metadata(
6388            &field.metadata,
6389            env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(),
6390            env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(),
6391        );
6392        Ok(dict_values_field)
6393    }
6394
6395    #[allow(clippy::too_many_arguments)]
6396    fn encode_miniblock(
6397        column_idx: u32,
6398        field: &Field,
6399        compression_strategy: &dyn CompressionStrategy,
6400        data: DataBlock,
6401        repdef: crate::repdef::SerializedRepDefs,
6402        row_number: u64,
6403        dictionary_data: Option<DataBlock>,
6404        num_rows: u64,
6405        miniblock_chunk_size: MiniblockChunkSize,
6406    ) -> Result<EncodedPage> {
6407        if let DataBlock::AllNull(_null_block) = data {
6408            // We should not be using mini-block for all-null.  There are other structural
6409            // encodings for that.
6410            unreachable!()
6411        }
6412
6413        let num_items = data.num_values();
6414
6415        let compressor = compression_strategy.create_miniblock_compressor(field, &data)?;
6416        let common_chunk_buffers =
6417            u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some());
6418        let support_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
6419        let compression_context =
6420            MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true);
6421        let (compressed_data, value_encoding) = compressor.compress(compression_context, data)?;
6422
6423        let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16;
6424
6425        let mut compressed_rep = repdef
6426            .rep_slicer()
6427            .map(|rep_slicer| {
6428                Self::compress_levels(
6429                    rep_slicer,
6430                    num_items,
6431                    compression_strategy,
6432                    &compressed_data.chunks,
6433                    max_rep,
6434                )
6435            })
6436            .transpose()?;
6437
6438        let (rep_index, rep_index_depth) =
6439            match compressed_rep.as_mut().and_then(|cr| cr.rep_index.as_mut()) {
6440                Some(rep_index) => (Some(rep_index.clone()), 1),
6441                None => (None, 0),
6442            };
6443
6444        let mut compressed_def = repdef
6445            .def_slicer()
6446            .map(|def_slicer| {
6447                Self::compress_levels(
6448                    def_slicer,
6449                    num_items,
6450                    compression_strategy,
6451                    &compressed_data.chunks,
6452                    /*max_rep=*/ 0,
6453                )
6454            })
6455            .transpose()?;
6456
6457        // TODO: Parquet sparsely encodes values here.  We could do the same but
6458        // then we won't have log2 values per chunk.  This means more metadata
6459        // and potentially more decoder asymmetry.  However, it may be worth
6460        // investigating at some point
6461
6462        let rep_data = compressed_rep
6463            .as_mut()
6464            .map(|cr| std::mem::take(&mut cr.data));
6465        let def_data = compressed_def
6466            .as_mut()
6467            .map(|cd| std::mem::take(&mut cd.data));
6468
6469        let serialized =
6470            Self::serialize_miniblocks(compressed_data, rep_data, def_data, miniblock_chunk_size)?;
6471        let has_large_chunk = miniblock_chunk_size == MiniblockChunkSize::U32;
6472
6473        // Metadata, Data, Dictionary, (maybe) Repetition Index
6474        let mut data = Vec::with_capacity(4);
6475        data.push(serialized.metadata);
6476        data.push(serialized.data);
6477
6478        if let Some(dictionary_data) = dictionary_data {
6479            let num_dictionary_items = dictionary_data.num_values();
6480            let dict_values_field = Self::build_dict_values_compressor_field(field)?;
6481
6482            let (dictionary_buffer, dictionary_encoding) =
6483                compress_required_block(compression_strategy, &dict_values_field, dictionary_data)?;
6484
6485            data.push(dictionary_buffer);
6486            if let Some(rep_index) = rep_index {
6487                data.push(rep_index);
6488            }
6489
6490            let description = ProtobufUtils21::miniblock_layout(
6491                compressed_rep.map(|cr| cr.compression),
6492                compressed_def.map(|cd| cd.compression),
6493                value_encoding,
6494                rep_index_depth,
6495                serialized.num_buffers,
6496                Some((dictionary_encoding, num_dictionary_items)),
6497                &repdef.def_meaning,
6498                num_items,
6499                has_large_chunk,
6500            );
6501            Ok(EncodedPage {
6502                num_rows,
6503                column_idx,
6504                data,
6505                description: PageEncoding::Structural(description),
6506                row_number,
6507            })
6508        } else {
6509            let description = ProtobufUtils21::miniblock_layout(
6510                compressed_rep.map(|cr| cr.compression),
6511                compressed_def.map(|cd| cd.compression),
6512                value_encoding,
6513                rep_index_depth,
6514                serialized.num_buffers,
6515                None,
6516                &repdef.def_meaning,
6517                num_items,
6518                has_large_chunk,
6519            );
6520
6521            if let Some(rep_index) = rep_index {
6522                let view = rep_index.borrow_to_typed_slice::<u64>();
6523                let total = view.as_chunks::<2>().0.iter().map(|c| c[0]).sum::<u64>();
6524                debug_assert_eq!(total, num_rows);
6525
6526                data.push(rep_index);
6527            }
6528
6529            Ok(EncodedPage {
6530                num_rows,
6531                column_idx,
6532                data,
6533                description: PageEncoding::Structural(description),
6534                row_number,
6535            })
6536        }
6537    }
6538
6539    // For fixed-size data we encode < control word | data > for each value
6540    fn serialize_full_zip_fixed(
6541        fixed: FixedWidthDataBlock,
6542        mut repdef: ControlWordIterator,
6543        num_values: u64,
6544    ) -> Result<SerializedFullZip> {
6545        if !fixed.bits_per_value.is_multiple_of(8) {
6546            return Err(Error::invalid_input_source(
6547                format!(
6548                    "Full-zip fixed-width values must be byte aligned, got {} bits per value",
6549                    fixed.bits_per_value
6550                )
6551                .into(),
6552            ));
6553        }
6554
6555        let len = fixed.data.len() + repdef.bytes_per_word() * num_values as usize;
6556        let mut zipped_data = Vec::with_capacity(len);
6557
6558        let max_rep_index_val = if repdef.has_repetition() {
6559            len as u64
6560        } else {
6561            // Setting this to 0 means we won't write a repetition index
6562            0
6563        };
6564        let mut rep_index_builder =
6565            BytepackedIntegerEncoder::with_capacity(num_values as usize + 1, max_rep_index_val);
6566
6567        let bytes_per_value = fixed.bits_per_value as usize / 8;
6568        let mut offset = 0;
6569
6570        if bytes_per_value == 0 {
6571            // No data, just dump the repdef into the buffer
6572            while let Some(control) = repdef.append_next(&mut zipped_data) {
6573                if control.is_new_row {
6574                    // We have finished a row
6575                    debug_assert!(offset <= len);
6576                    rep_index_builder.append_trusted(offset as u64);
6577                }
6578                offset = zipped_data.len();
6579            }
6580        } else {
6581            // We have data, zip it with the repdef
6582            let mut data_iter = fixed.data.chunks_exact(bytes_per_value);
6583            while let Some(control) = repdef.append_next(&mut zipped_data) {
6584                if control.is_new_row {
6585                    // We have finished a row
6586                    debug_assert!(offset <= len);
6587                    rep_index_builder.append_trusted(offset as u64);
6588                }
6589                if control.is_visible {
6590                    let value = data_iter.next().unwrap();
6591                    zipped_data.extend_from_slice(value);
6592                }
6593                offset = zipped_data.len();
6594            }
6595        }
6596
6597        debug_assert_eq!(zipped_data.len(), len);
6598        // Put the final value in the rep index
6599        rep_index_builder.append_trusted(zipped_data.len() as u64);
6600
6601        let zipped_data = LanceBuffer::from(zipped_data);
6602        let rep_index = rep_index_builder.into_data();
6603        let rep_index = if rep_index.is_empty() {
6604            None
6605        } else {
6606            Some(LanceBuffer::from(rep_index))
6607        };
6608        Ok(SerializedFullZip {
6609            values: zipped_data,
6610            repetition_index: rep_index,
6611        })
6612    }
6613
6614    // For variable-size data we encode < control word | length | data > for each value
6615    //
6616    // In addition, we create a second buffer, the repetition index
6617    fn serialize_full_zip_variable(
6618        variable: VariableWidthBlock,
6619        mut repdef: ControlWordIterator,
6620        num_items: u64,
6621    ) -> Result<SerializedFullZip> {
6622        let bytes_per_offset = variable.bits_per_offset as usize / 8;
6623        if !variable.bits_per_offset.is_multiple_of(8) {
6624            return Err(Error::invalid_input_source(
6625                format!(
6626                    "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
6627                    variable.bits_per_offset
6628                )
6629                .into(),
6630            ));
6631        }
6632        let len = variable.data.len()
6633            + repdef.bytes_per_word() * num_items as usize
6634            + bytes_per_offset * variable.num_values as usize;
6635        let mut buf = Vec::with_capacity(len);
6636
6637        let max_rep_index_val = len as u64;
6638        let mut rep_index_builder =
6639            BytepackedIntegerEncoder::with_capacity(num_items as usize + 1, max_rep_index_val);
6640
6641        // TODO: byte pack the item lengths with varint encoding
6642        match bytes_per_offset {
6643            4 => {
6644                let offs = variable.offsets.borrow_to_typed_slice::<u32>();
6645                let mut rep_offset = 0;
6646                let mut windows_iter = offs.as_ref().windows(2);
6647                while let Some(control) = repdef.append_next(&mut buf) {
6648                    if control.is_new_row {
6649                        // We have finished a row
6650                        debug_assert!(rep_offset <= len);
6651                        rep_index_builder.append_trusted(rep_offset as u64);
6652                    }
6653                    if control.is_visible {
6654                        let window = windows_iter.next().unwrap();
6655                        if control.is_valid_item {
6656                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
6657                            buf.extend_from_slice(
6658                                &variable.data[window[0] as usize..window[1] as usize],
6659                            );
6660                        }
6661                    }
6662                    rep_offset = buf.len();
6663                }
6664            }
6665            8 => {
6666                let offs = variable.offsets.borrow_to_typed_slice::<u64>();
6667                let mut rep_offset = 0;
6668                let mut windows_iter = offs.as_ref().windows(2);
6669                while let Some(control) = repdef.append_next(&mut buf) {
6670                    if control.is_new_row {
6671                        // We have finished a row
6672                        debug_assert!(rep_offset <= len);
6673                        rep_index_builder.append_trusted(rep_offset as u64);
6674                    }
6675                    if control.is_visible {
6676                        let window = windows_iter.next().unwrap();
6677                        if control.is_valid_item {
6678                            buf.extend_from_slice(&(window[1] - window[0]).to_le_bytes());
6679                            buf.extend_from_slice(
6680                                &variable.data[window[0] as usize..window[1] as usize],
6681                            );
6682                        }
6683                    }
6684                    rep_offset = buf.len();
6685                }
6686            }
6687            _ => {
6688                return Err(Error::invalid_input_source(
6689                    format!(
6690                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
6691                        variable.bits_per_offset
6692                    )
6693                    .into(),
6694                ));
6695            }
6696        }
6697
6698        // We might have saved a few bytes by not copying lengths when the length was zero.  However,
6699        // if we are over `len` then we have a bug.
6700        debug_assert!(buf.len() <= len);
6701        // Put the final value in the rep index
6702        rep_index_builder.append_trusted(buf.len() as u64);
6703
6704        let zipped_data = LanceBuffer::from(buf);
6705        let rep_index = rep_index_builder.into_data();
6706        debug_assert!(!rep_index.is_empty());
6707        let rep_index = Some(LanceBuffer::from(rep_index));
6708        Ok(SerializedFullZip {
6709            values: zipped_data,
6710            repetition_index: rep_index,
6711        })
6712    }
6713
6714    /// Serializes data into a single buffer according to the full-zip format which zips
6715    /// together the repetition, definition, and value data into a single buffer.
6716    fn serialize_full_zip(
6717        compressed_data: PerValueDataBlock,
6718        repdef: ControlWordIterator,
6719        num_items: u64,
6720    ) -> Result<SerializedFullZip> {
6721        match compressed_data {
6722            PerValueDataBlock::Fixed(fixed) => {
6723                Self::serialize_full_zip_fixed(fixed, repdef, num_items)
6724            }
6725            PerValueDataBlock::Variable(var) => {
6726                Self::serialize_full_zip_variable(var, repdef, num_items)
6727            }
6728        }
6729    }
6730
6731    fn expand_boolean_to_bytes(fixed: FixedWidthDataBlock) -> FixedWidthDataBlock {
6732        debug_assert_eq!(fixed.bits_per_value, 1);
6733        let num_values = fixed.num_values as usize;
6734        let bool_buf = BooleanBuffer::new(fixed.data.into_buffer(), 0, num_values);
6735        let expanded: Vec<u8> = (0..num_values).map(|i| bool_buf.value(i) as u8).collect();
6736        FixedWidthDataBlock {
6737            data: LanceBuffer::from(expanded),
6738            bits_per_value: 8,
6739            num_values: fixed.num_values,
6740            block_info: BlockInfo::new(),
6741        }
6742    }
6743
6744    fn encode_full_zip(
6745        column_idx: u32,
6746        field: &Field,
6747        compression_strategy: &dyn CompressionStrategy,
6748        data: DataBlock,
6749        repdef: crate::repdef::SerializedRepDefs,
6750        row_number: u64,
6751        num_lists: u64,
6752    ) -> Result<EncodedPage> {
6753        let max_rep = repdef
6754            .repetition_levels
6755            .as_ref()
6756            .map_or(0, |r| r.iter().max().copied().unwrap_or(0));
6757        let max_def = repdef
6758            .definition_levels
6759            .as_ref()
6760            .map_or(0, |d| d.iter().max().copied().unwrap_or(0));
6761
6762        // To handle FSL we just flatten
6763        // let data = data.flatten();
6764
6765        let (num_items, num_visible_items) =
6766            if let Some(rep_levels) = repdef.repetition_levels.as_ref() {
6767                // If there are rep levels there may be "invisible" items and we need to encode
6768                // rep_levels.len() things which might be larger than data.num_values()
6769                (rep_levels.len() as u64, data.num_values())
6770            } else {
6771                // If there are no rep levels then we encode data.num_values() things
6772                (data.num_values(), data.num_values())
6773            };
6774
6775        let max_visible_def = repdef.max_visible_level.unwrap_or(u16::MAX);
6776
6777        let repdef_iter = build_control_word_iterator(
6778            repdef.repetition_levels.as_deref(),
6779            max_rep,
6780            repdef.definition_levels.as_deref(),
6781            max_def,
6782            max_visible_def,
6783            num_items as usize,
6784        );
6785        let bits_rep = repdef_iter.bits_rep();
6786        let bits_def = repdef_iter.bits_def();
6787
6788        // Full-zip requires byte-aligned values; expand 1-bit booleans to 1 byte each.
6789        let data = match data {
6790            DataBlock::FixedWidth(fixed) if fixed.bits_per_value == 1 => {
6791                DataBlock::FixedWidth(Self::expand_boolean_to_bytes(fixed))
6792            }
6793            other => other,
6794        };
6795
6796        let compressor = compression_strategy.create_per_value(field, &data)?;
6797        let (compressed_data, value_encoding) = compressor.compress(data)?;
6798
6799        let description = match &compressed_data {
6800            PerValueDataBlock::Fixed(fixed) => ProtobufUtils21::fixed_full_zip_layout(
6801                bits_rep,
6802                bits_def,
6803                fixed.bits_per_value as u32,
6804                value_encoding,
6805                &repdef.def_meaning,
6806                num_items as u32,
6807                num_visible_items as u32,
6808            ),
6809            PerValueDataBlock::Variable(variable) => ProtobufUtils21::variable_full_zip_layout(
6810                bits_rep,
6811                bits_def,
6812                variable.bits_per_offset as u32,
6813                value_encoding,
6814                &repdef.def_meaning,
6815                num_items as u32,
6816                num_visible_items as u32,
6817            ),
6818        };
6819
6820        let zipped = Self::serialize_full_zip(compressed_data, repdef_iter, num_items)?;
6821
6822        let data = if let Some(repindex) = zipped.repetition_index {
6823            vec![zipped.values, repindex]
6824        } else {
6825            vec![zipped.values]
6826        };
6827
6828        Ok(EncodedPage {
6829            num_rows: num_lists,
6830            column_idx,
6831            data,
6832            description: PageEncoding::Structural(description),
6833            row_number,
6834        })
6835    }
6836
6837    fn should_dictionary_encode(
6838        data_block: &DataBlock,
6839        field: &Field,
6840        fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
6841    ) -> Option<DictEncodingBudget> {
6842        const DEFAULT_SAMPLE_SIZE: usize = 4096;
6843        const DEFAULT_SAMPLE_UNIQUE_RATIO: f64 = 0.98;
6844
6845        // Since we only dictionary encode FixedWidth and VariableWidth blocks for now, we skip
6846        // estimating the size for other types.
6847        match data_block {
6848            DataBlock::FixedWidth(fixed) => {
6849                if fixed.bits_per_value == 64
6850                    && fixed_width_dictionary_encoding == FixedWidthDictionaryEncoding::Exclude64Bit
6851                {
6852                    return None;
6853                }
6854                if fixed.bits_per_value != 64 && fixed.bits_per_value != 128 {
6855                    return None;
6856                }
6857                if fixed.bits_per_value % 8 != 0 {
6858                    return None;
6859                }
6860            }
6861            DataBlock::VariableWidth(var) => {
6862                if var.bits_per_offset != 32 && var.bits_per_offset != 64 {
6863                    return None;
6864                }
6865            }
6866            _ => return None,
6867        }
6868
6869        // Don't dictionary encode tiny arrays.
6870        let too_small = env::var("LANCE_ENCODING_DICT_TOO_SMALL")
6871            .ok()
6872            .and_then(|val| val.parse().ok())
6873            .unwrap_or(100);
6874        if data_block.num_values() < too_small {
6875            return None;
6876        }
6877
6878        let num_values = data_block.num_values();
6879
6880        // Apply divisor threshold and cap. This is intentionally conservative: the goal is to
6881        // avoid spending too much CPU trying to estimate very high cardinalities.
6882        let divisor: u64 = field
6883            .metadata
6884            .get(DICT_DIVISOR_META_KEY)
6885            .and_then(|val| val.parse().ok())
6886            .or_else(|| {
6887                env::var("LANCE_ENCODING_DICT_DIVISOR")
6888                    .ok()
6889                    .and_then(|val| val.parse().ok())
6890            })
6891            .unwrap_or(DEFAULT_DICT_DIVISOR);
6892
6893        let max_cardinality: u64 = env::var("LANCE_ENCODING_DICT_MAX_CARDINALITY")
6894            .ok()
6895            .and_then(|val| val.parse().ok())
6896            .unwrap_or(DEFAULT_DICT_MAX_CARDINALITY);
6897
6898        let threshold_cardinality = num_values
6899            .checked_div(divisor.max(1))
6900            .unwrap_or(0)
6901            .min(max_cardinality);
6902        if threshold_cardinality == 0 {
6903            return None;
6904        }
6905
6906        // Get size ratio from metadata or env var.
6907        let threshold_ratio = field
6908            .metadata
6909            .get(DICT_SIZE_RATIO_META_KEY)
6910            .and_then(|val| val.parse::<f64>().ok())
6911            .or_else(|| {
6912                env::var("LANCE_ENCODING_DICT_SIZE_RATIO")
6913                    .ok()
6914                    .and_then(|val| val.parse().ok())
6915            })
6916            .unwrap_or(DEFAULT_DICT_SIZE_RATIO);
6917
6918        if threshold_ratio <= 0.0 || threshold_ratio > 1.0 {
6919            panic!(
6920                "Invalid parameter: dict-size-ratio is {} which is not in the range (0, 1].",
6921                threshold_ratio
6922            );
6923        }
6924
6925        let data_size = data_block.data_size();
6926        if data_size == 0 {
6927            return None;
6928        }
6929
6930        let max_encoded_size = (data_size as f64 * threshold_ratio) as u64;
6931        let max_encoded_size = usize::try_from(max_encoded_size).ok()?;
6932
6933        // Avoid probing dictionary encoding on data that appears to be near-unique
6934        // or likely to exceed the dictionary budget.
6935        if let Some(sample_unique_ratio) =
6936            Self::sample_unique_ratio(data_block, DEFAULT_SAMPLE_SIZE)?
6937        {
6938            if sample_unique_ratio >= DEFAULT_SAMPLE_UNIQUE_RATIO {
6939                return None;
6940            }
6941
6942            let projected_cardinality = (sample_unique_ratio * num_values as f64).ceil() as u64;
6943            if projected_cardinality > threshold_cardinality {
6944                return None;
6945            }
6946        }
6947
6948        let max_dict_entries = u32::try_from(threshold_cardinality.min(i32::MAX as u64)).ok()?;
6949        Some(DictEncodingBudget {
6950            max_dict_entries,
6951            max_encoded_size,
6952        })
6953    }
6954
6955    /// Samples whether a page looks near-unique before attempting dictionary encoding.
6956    ///
6957    /// The probe uses deterministic block sampling (not RNG sampling), which keeps
6958    /// the check cheap and reproducible across runs. The result is only a gate for
6959    /// whether we try dictionary encoding, not a cardinality statistic.
6960    /// Returns `Some(None)` when there are too few reliable samples or the block type does not
6961    /// support dictionary encoding. Returns `None` for malformed data.
6962    fn sample_unique_ratio(data_block: &DataBlock, max_samples: usize) -> Option<Option<f64>> {
6963        use std::collections::HashSet;
6964
6965        const NUM_SAMPLE_BLOCKS: usize = 32;
6966        const MIN_RELIABLE_SAMPLES: usize = 1024;
6967
6968        let num_values = usize::try_from(data_block.num_values()).ok()?;
6969        if num_values == 0 {
6970            return Some(None);
6971        }
6972
6973        let sample_count = num_values.min(max_samples).max(1);
6974        if sample_count < MIN_RELIABLE_SAMPLES {
6975            return Some(None);
6976        }
6977
6978        let block_count = NUM_SAMPLE_BLOCKS.min(sample_count).min(num_values).max(1);
6979        let samples_per_block = (sample_count / block_count).max(1);
6980        let mut indices = Vec::with_capacity(sample_count);
6981        for block_idx in 0..block_count {
6982            let block_start = block_idx * num_values / block_count;
6983            let next_block_start = ((block_idx + 1) * num_values / block_count).min(num_values);
6984            let block_len = next_block_start.saturating_sub(block_start);
6985            let samples_in_block = samples_per_block.min(block_len);
6986            indices.extend((0..samples_in_block).map(|offset| block_start + offset));
6987        }
6988
6989        if indices.len() < MIN_RELIABLE_SAMPLES {
6990            return Some(None);
6991        }
6992
6993        let ratio = match data_block {
6994            DataBlock::FixedWidth(fixed) => match fixed.bits_per_value {
6995                64 => {
6996                    let values = fixed.data.borrow_to_typed_slice::<u64>();
6997                    let values = values.as_ref();
6998                    let mut unique: HashSet<u64> =
6999                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7000                    for idx in indices.iter().copied() {
7001                        unique.insert(values.get(idx).copied()?);
7002                    }
7003                    unique.len() as f64 / indices.len() as f64
7004                }
7005                128 => {
7006                    let values = fixed.data.borrow_to_typed_slice::<u128>();
7007                    let values = values.as_ref();
7008                    let mut unique: HashSet<u128> =
7009                        HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7010                    for idx in indices.iter().copied() {
7011                        unique.insert(values.get(idx).copied()?);
7012                    }
7013                    unique.len() as f64 / indices.len() as f64
7014                }
7015                _ => return Some(None),
7016            },
7017            DataBlock::VariableWidth(var) => {
7018                use xxhash_rust::xxh3::xxh3_64;
7019
7020                // Hash variable-width slices instead of storing borrowed slice keys.
7021                let mut unique: HashSet<u64> =
7022                    HashSet::with_capacity(indices.len().min(MIN_RELIABLE_SAMPLES));
7023                match var.bits_per_offset {
7024                    32 => {
7025                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u32>();
7026                        let offsets: &[u32] = offsets_ref.as_ref();
7027                        for i in indices.iter().copied() {
7028                            let start = usize::try_from(*offsets.get(i)?).ok()?;
7029                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
7030                            if start > end || end > var.data.len() {
7031                                return None;
7032                            }
7033                            unique.insert(xxh3_64(&var.data[start..end]));
7034                        }
7035                    }
7036                    64 => {
7037                        let offsets_ref = var.offsets.borrow_to_typed_slice::<u64>();
7038                        let offsets: &[u64] = offsets_ref.as_ref();
7039                        for i in indices.iter().copied() {
7040                            let start = usize::try_from(*offsets.get(i)?).ok()?;
7041                            let end = usize::try_from(*offsets.get(i + 1)?).ok()?;
7042                            if start > end || end > var.data.len() {
7043                                return None;
7044                            }
7045                            unique.insert(xxh3_64(&var.data[start..end]));
7046                        }
7047                    }
7048                    _ => return Some(None),
7049                }
7050                unique.len() as f64 / indices.len() as f64
7051            }
7052            _ => return Some(None),
7053        };
7054
7055        Some(Some(ratio))
7056    }
7057
7058    fn slice_repdef(repdef: &SerializedRepDefs, range: Range<usize>) -> SerializedRepDefs {
7059        let repetition_levels = repdef
7060            .repetition_levels
7061            .as_ref()
7062            .map(|levels| levels[range.clone()].to_vec());
7063        let definition_levels = repdef
7064            .definition_levels
7065            .as_ref()
7066            .map(|levels| levels[range].to_vec());
7067        SerializedRepDefs::new_with_fixed_size_list_levels(
7068            repetition_levels,
7069            definition_levels,
7070            repdef.def_meaning.clone(),
7071            repdef.has_fixed_size_list_levels(),
7072        )
7073    }
7074
7075    fn slice_arrays(
7076        arrays: &[ArrayRef],
7077        value_start: u64,
7078        num_values: u64,
7079    ) -> Result<Vec<ArrayRef>> {
7080        if num_values == 0 {
7081            return Ok(Vec::new());
7082        }
7083
7084        let mut values_to_skip = usize::try_from(value_start).map_err(|_| {
7085            Error::invalid_input(format!("Value start {} is too large", value_start))
7086        })?;
7087        let mut values_remaining = usize::try_from(num_values).map_err(|_| {
7088            Error::invalid_input(format!("Value count {} is too large", num_values))
7089        })?;
7090        let mut sliced = Vec::new();
7091
7092        for array in arrays {
7093            if values_to_skip >= array.len() {
7094                values_to_skip -= array.len();
7095                continue;
7096            }
7097
7098            let offset = values_to_skip;
7099            let len = (array.len() - offset).min(values_remaining);
7100            sliced.push(array.slice(offset, len));
7101            values_remaining -= len;
7102            values_to_skip = 0;
7103
7104            if values_remaining == 0 {
7105                break;
7106            }
7107        }
7108
7109        if values_remaining != 0 {
7110            return Err(Error::internal(format!(
7111                "Page split requested {} values starting at {}, but the page did not contain enough values",
7112                num_values, value_start
7113            )));
7114        }
7115
7116        Ok(sliced)
7117    }
7118
7119    fn split_pages_for_miniblock_repdef_budget(
7120        arrays: Vec<ArrayRef>,
7121        repdef: SerializedRepDefs,
7122        budget: MiniBlockRepDefBudget,
7123        row_number: u64,
7124        num_rows: u64,
7125    ) -> Result<Vec<PrimitivePageData>> {
7126        if budget == MiniBlockRepDefBudget::WithinBudget {
7127            return Ok(vec![PrimitivePageData {
7128                arrays,
7129                structure: PrimitivePageStructure::Dense {
7130                    repdef,
7131                    single_row_miniblock_repdef_levels: None,
7132                },
7133                row_number,
7134                num_rows,
7135            }]);
7136        }
7137        if let MiniBlockRepDefBudget::SingleRowOverBudget(num_levels) = budget {
7138            return Ok(vec![PrimitivePageData {
7139                arrays,
7140                structure: PrimitivePageStructure::Dense {
7141                    repdef,
7142                    single_row_miniblock_repdef_levels: Some(num_levels),
7143                },
7144                row_number,
7145                num_rows,
7146            }]);
7147        }
7148
7149        let MiniBlockRepDefBudget::RequiresPageSplit(splits) = budget else {
7150            unreachable!();
7151        };
7152
7153        let mut pages = Vec::with_capacity(splits.len());
7154        for split in splits {
7155            let arrays = Self::slice_arrays(&arrays, split.value_start, split.num_values)?;
7156            let repdef = Self::slice_repdef(&repdef, split.level_range);
7157            pages.push(PrimitivePageData {
7158                arrays,
7159                structure: PrimitivePageStructure::Dense {
7160                    repdef,
7161                    single_row_miniblock_repdef_levels: None,
7162                },
7163                row_number: row_number + split.row_start,
7164                num_rows: split.num_rows,
7165            });
7166        }
7167        Ok(pages)
7168    }
7169
7170    fn encode_dense_page(
7171        ctx: PrimitiveEncodeContext,
7172        page: PrimitivePageData,
7173        compression_strategy: Arc<dyn CompressionStrategy>,
7174        miniblock_chunk_size: MiniblockChunkSize,
7175        complex_null_encoding: ComplexNullEncoding,
7176        fixed_width_dictionary_encoding: FixedWidthDictionaryEncoding,
7177    ) -> Result<EncodedPage> {
7178        let PrimitiveEncodeContext {
7179            column_idx,
7180            field,
7181            encoding_metadata,
7182            is_simple_validity,
7183            has_repdef_info,
7184        } = ctx;
7185        let PrimitivePageData {
7186            arrays,
7187            structure,
7188            row_number,
7189            num_rows,
7190        } = page;
7191        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
7192
7193        let (repdef, single_row_miniblock_repdef_levels) = match structure {
7194            PrimitivePageStructure::Dense {
7195                repdef,
7196                single_row_miniblock_repdef_levels,
7197            } => (repdef, single_row_miniblock_repdef_levels),
7198            PrimitivePageStructure::Sparse { .. } => {
7199                unreachable!("dense atom received sparse page")
7200            }
7201        };
7202
7203        if num_values == 0 {
7204            // This page contains only structural events, such as empty/null list rows.
7205            // The existing complex-null layout stores the rep/def stream without value buffers.
7206            log::debug!(
7207                "Encoding column {} with {} items ({} rows) using complex-null layout",
7208                column_idx,
7209                num_values,
7210                num_rows
7211            );
7212            return Self::encode_complex_all_null(
7213                column_idx,
7214                repdef,
7215                row_number,
7216                num_rows,
7217                complex_null_encoding,
7218                compression_strategy.as_ref(),
7219            );
7220        }
7221
7222        let leaf_validity = Self::leaf_validity(&repdef, num_values as usize)?;
7223        let all_null = leaf_validity
7224            .as_ref()
7225            .map(|validity| validity.count_set_bits() == 0)
7226            .unwrap_or(false);
7227
7228        if all_null {
7229            return if is_simple_validity {
7230                log::debug!(
7231                    "Encoding column {} with {} items ({} rows) using simple-null layout",
7232                    column_idx,
7233                    num_values,
7234                    num_rows
7235                );
7236                Self::encode_simple_all_null(column_idx, num_values, row_number)
7237            } else {
7238                log::debug!(
7239                    "Encoding column {} with {} items ({} rows) using complex-null layout",
7240                    column_idx,
7241                    num_values,
7242                    num_rows
7243                );
7244                Self::encode_complex_all_null(
7245                    column_idx,
7246                    repdef,
7247                    row_number,
7248                    num_rows,
7249                    complex_null_encoding,
7250                    compression_strategy.as_ref(),
7251                )
7252            };
7253        }
7254
7255        if let DataType::Struct(fields) = &field.data_type()
7256            && fields.is_empty()
7257        {
7258            if has_repdef_info {
7259                return Err(Error::invalid_input_source(format!("Empty structs with rep/def information are not yet supported.  The field {} is an empty struct that either has nulls or is in a list.", field.name).into()));
7260            }
7261            // This is maybe a little confusing but the reader should never look at this anyways and it
7262            // seems like overkill to invent a new layout just for "empty structs".
7263            return Self::encode_simple_all_null(column_idx, num_values, row_number);
7264        }
7265
7266        let data_block = DataBlock::from_arrays(&arrays, num_values);
7267
7268        if let Some(num_levels) = single_row_miniblock_repdef_levels {
7269            let requested_encoding = encoding_metadata
7270                .get(STRUCTURAL_ENCODING_META_KEY)
7271                .map(|requested| requested.to_lowercase());
7272            let fullzip_error = match &data_block {
7273                // 1-bit booleans are widened to bytes inside `encode_full_zip`.
7274                DataBlock::FixedWidth(fixed)
7275                    if fixed.bits_per_value != 1 && !fixed.bits_per_value.is_multiple_of(8) =>
7276                {
7277                    Some(format!(
7278                        "Full-zip fixed-width values must be byte aligned, got {} bits per value",
7279                        fixed.bits_per_value
7280                    ))
7281                }
7282                DataBlock::VariableWidth(variable)
7283                    if !variable.bits_per_offset.is_multiple_of(8) =>
7284                {
7285                    Some(format!(
7286                        "Full-zip variable-width offsets must be byte aligned, got {} bits per offset",
7287                        variable.bits_per_offset
7288                    ))
7289                }
7290                DataBlock::VariableWidth(variable)
7291                    if variable.bits_per_offset != 32 && variable.bits_per_offset != 64 =>
7292                {
7293                    Some(format!(
7294                        "Full-zip variable-width offsets must be 32 or 64 bits, got {} bits",
7295                        variable.bits_per_offset
7296                    ))
7297                }
7298                DataBlock::Dictionary(_) => {
7299                    Some("Full-zip does not encode dictionary data blocks directly".to_string())
7300                }
7301                DataBlock::FixedSizeList(fsl) => match fsl.clone().try_into_flat() {
7302                    Some(flat) if flat.bits_per_value.is_multiple_of(8) => None,
7303                    Some(flat) => Some(format!(
7304                        "Full-zip fixed-size-list values must be byte aligned after flattening, got {} bits per value",
7305                        flat.bits_per_value
7306                    )),
7307                    None => Some(
7308                        "Full-zip fixed-size-list capability requires a flat fixed-width child"
7309                            .to_string(),
7310                    ),
7311                },
7312                DataBlock::FixedWidth(_) | DataBlock::VariableWidth(_) | DataBlock::Struct(_) => {
7313                    None
7314                }
7315                other => Some(format!(
7316                    "Full-zip does not support value block type {}",
7317                    other.name()
7318                )),
7319            };
7320            match requested_encoding.as_deref() {
7321                Some(STRUCTURAL_ENCODING_FULLZIP) => {
7322                    if let Some(reason) = fullzip_error {
7323                        return Err(Error::invalid_input_source(reason.into()));
7324                    }
7325                    return Self::encode_full_zip(
7326                        column_idx,
7327                        &field,
7328                        compression_strategy.as_ref(),
7329                        data_block,
7330                        repdef,
7331                        row_number,
7332                        num_rows,
7333                    );
7334                }
7335                Some(STRUCTURAL_ENCODING_MINIBLOCK) | None => {
7336                    if requested_encoding.is_none() && fullzip_error.is_none() {
7337                        log::debug!(
7338                            "Encoding column {} with {} items using full-zip layout because mini-block cannot split the structural page",
7339                            column_idx,
7340                            num_values
7341                        );
7342                        return Self::encode_full_zip(
7343                            column_idx,
7344                            &field,
7345                            compression_strategy.as_ref(),
7346                            data_block,
7347                            repdef,
7348                            row_number,
7349                            num_rows,
7350                        );
7351                    }
7352                    return Err(Error::invalid_input_source(
7353                        format!(
7354                            "Mini-block cannot encode {} rep/def levels in one top-level row. \
7355                             This usually means the row contains too much nested structure \
7356                             for the current layout.",
7357                            num_levels
7358                        )
7359                        .into(),
7360                    ));
7361                }
7362                _ => {}
7363            }
7364        }
7365
7366        let requires_full_zip_packed_struct =
7367            if let DataBlock::Struct(ref struct_data_block) = data_block {
7368                struct_data_block.has_variable_width_child()
7369            } else {
7370                false
7371            };
7372
7373        if requires_full_zip_packed_struct {
7374            log::debug!(
7375                "Encoding column {} with {} items using full-zip packed struct layout",
7376                column_idx,
7377                num_values
7378            );
7379            return Self::encode_full_zip(
7380                column_idx,
7381                &field,
7382                compression_strategy.as_ref(),
7383                data_block,
7384                repdef,
7385                row_number,
7386                num_rows,
7387            );
7388        }
7389
7390        if let DataBlock::Dictionary(dict) = data_block {
7391            log::debug!(
7392                "Encoding column {} with {} items using dictionary encoding (already dictionary encoded)",
7393                column_idx,
7394                num_values
7395            );
7396            let (mut indices_data_block, dictionary_data_block) = dict.into_parts();
7397            // TODO: https://github.com/lancedb/lance/issues/4809
7398            // If we compute stats on dictionary_data_block => panic.
7399            // If we don't compute stats on indices_data_block => panic.
7400            // This is messy.  Don't make me call compute_stat ever.
7401            indices_data_block.compute_stat();
7402            return Self::encode_miniblock(
7403                column_idx,
7404                &field,
7405                compression_strategy.as_ref(),
7406                indices_data_block,
7407                repdef,
7408                row_number,
7409                Some(dictionary_data_block),
7410                num_rows,
7411                miniblock_chunk_size,
7412            );
7413        }
7414
7415        // Try dictionary encoding first if applicable. If encoding aborts, fall back to the
7416        // preferred structural encoding.
7417        let dict_result = Self::should_dictionary_encode(
7418            &data_block,
7419            &field,
7420            fixed_width_dictionary_encoding,
7421        )
7422        .and_then(|budget| {
7423            log::debug!(
7424                "Encoding column {} with {} items using dictionary encoding (mini-block layout)",
7425                column_idx,
7426                num_values
7427            );
7428            dict::dictionary_encode(
7429                &data_block,
7430                budget.max_dict_entries,
7431                budget.max_encoded_size,
7432            )
7433        });
7434
7435        if let Some((indices_data_block, dictionary_data_block)) = dict_result {
7436            Self::encode_miniblock(
7437                column_idx,
7438                &field,
7439                compression_strategy.as_ref(),
7440                indices_data_block,
7441                repdef,
7442                row_number,
7443                Some(dictionary_data_block),
7444                num_rows,
7445                miniblock_chunk_size,
7446            )
7447        } else if Self::prefers_miniblock(&data_block, encoding_metadata.as_ref()) {
7448            log::debug!(
7449                "Encoding column {} with {} items using mini-block layout",
7450                column_idx,
7451                num_values
7452            );
7453            Self::encode_miniblock(
7454                column_idx,
7455                &field,
7456                compression_strategy.as_ref(),
7457                data_block,
7458                repdef,
7459                row_number,
7460                None,
7461                num_rows,
7462                miniblock_chunk_size,
7463            )
7464        } else if Self::prefers_fullzip(encoding_metadata.as_ref()) {
7465            log::debug!(
7466                "Encoding column {} with {} items using full-zip layout",
7467                column_idx,
7468                num_values
7469            );
7470            Self::encode_full_zip(
7471                column_idx,
7472                &field,
7473                compression_strategy.as_ref(),
7474                data_block,
7475                repdef,
7476                row_number,
7477                num_rows,
7478            )
7479        } else {
7480            Err(Error::invalid_input_source(format!("Cannot determine structural encoding for field {}.  This typically indicates an invalid value of the field metadata key {}", field.name, STRUCTURAL_ENCODING_META_KEY).into()))
7481        }
7482    }
7483
7484    /// Rebuilds dictionary chunks whose values array is empty.
7485    ///
7486    /// Such chunks are entirely null and retain their key validity until rep/def has been
7487    /// recorded.  At flush time their keys are zeroed and attached to a neighboring non-empty
7488    /// dictionary.  If every chunk is empty, one non-null placeholder value makes key zero valid.
7489    /// Rep/def retains the logical nullness, so these replacement keys are never exposed.
7490    fn rebuild_empty_dictionary_chunks(arrays: Vec<ArrayRef>) -> Result<Vec<ArrayRef>> {
7491        if !arrays.iter().any(|array| {
7492            array
7493                .as_any_dictionary_opt()
7494                .is_some_and(|dictionary| dictionary.values().is_empty())
7495        }) {
7496            return Ok(arrays);
7497        }
7498
7499        let zeroed = |data_type: &DataType, len: usize| {
7500            new_null_array(data_type, len)
7501                .to_data()
7502                .into_builder()
7503                .nulls(None)
7504                .build()
7505                .map(make_array)
7506        };
7507        let rebuild =
7508            |keys: Vec<ArrayRef>, data_type: &DataType, values: ArrayRef| -> Result<ArrayRef> {
7509                let keys = keys.iter().map(|keys| keys.as_ref()).collect::<Vec<_>>();
7510                let keys = arrow_select::concat::concat(&keys)?;
7511                let data = keys
7512                    .to_data()
7513                    .into_builder()
7514                    .data_type(data_type.clone())
7515                    .child_data(vec![values.to_data()])
7516                    .build()?;
7517                Ok(make_array(data))
7518            };
7519
7520        let mut rebuilt = Vec::with_capacity(arrays.len());
7521        let mut pending_keys = Vec::with_capacity(arrays.len());
7522        let mut empty_data_type = None;
7523        for array in arrays {
7524            let Some(dictionary) = array.as_any_dictionary_opt() else {
7525                return Err(Error::invalid_input_source(
7526                    "Cannot mix dictionary and non-dictionary chunks".into(),
7527                ));
7528            };
7529            if dictionary.values().is_empty() {
7530                pending_keys.push(zeroed(dictionary.keys().data_type(), array.len())?);
7531                empty_data_type.get_or_insert_with(|| array.data_type().clone());
7532            } else if pending_keys.is_empty() {
7533                rebuilt.push(array);
7534            } else {
7535                pending_keys.push(make_array(dictionary.keys().to_data()));
7536                rebuilt.push(rebuild(
7537                    std::mem::take(&mut pending_keys),
7538                    array.data_type(),
7539                    dictionary.values().clone(),
7540                )?);
7541            }
7542        }
7543
7544        if pending_keys.is_empty() {
7545            return Ok(rebuilt);
7546        }
7547        if let Some(array) = rebuilt.pop() {
7548            let Some(dictionary) = array.as_any_dictionary_opt() else {
7549                return Err(Error::invalid_input_source(
7550                    "Cannot mix dictionary and non-dictionary chunks".into(),
7551                ));
7552            };
7553            let mut keys = Vec::with_capacity(pending_keys.len() + 1);
7554            keys.push(make_array(dictionary.keys().to_data()));
7555            keys.append(&mut pending_keys);
7556            rebuilt.push(rebuild(
7557                keys,
7558                array.data_type(),
7559                dictionary.values().clone(),
7560            )?);
7561        } else {
7562            let data_type = empty_data_type.ok_or_else(|| {
7563                Error::internal("Missing data type for an empty dictionary chunk")
7564            })?;
7565            let DataType::Dictionary(_, value_type) = &data_type else {
7566                return Err(Error::internal(format!(
7567                    "Expected dictionary data type, got {data_type}"
7568                )));
7569            };
7570            rebuilt.push(rebuild(pending_keys, &data_type, zeroed(value_type, 1)?)?);
7571        }
7572        Ok(rebuilt)
7573    }
7574
7575    // Creates encode tasks, consuming all buffered data
7576    fn do_flush(
7577        &mut self,
7578        arrays: Vec<ArrayRef>,
7579        repdefs: Vec<RepDefBuilder>,
7580        row_number: u64,
7581        num_rows: u64,
7582    ) -> Result<Vec<EncodeTask>> {
7583        let arrays = Self::rebuild_empty_dictionary_chunks(arrays)?;
7584        DataBlock::validate_arrays(&arrays, &self.field.name)?;
7585        let num_values = arrays.iter().map(|arr| arr.len() as u64).sum();
7586        let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity());
7587        let has_repdef_info = repdefs.iter().any(|rd| !rd.is_empty());
7588        let normalized = RepDefBuilder::normalize(repdefs);
7589        let plan_ctx = PrimitivePlanContext {
7590            column_idx: self.column_index,
7591            field: &self.field,
7592            encoding_metadata: &self.encoding_metadata,
7593        };
7594        let mut pages = None;
7595        for page_encoding in self.page_encodings.iter() {
7596            if let Some(planned) = page_encoding.behavior.try_plan_pages(
7597                &plan_ctx,
7598                &arrays,
7599                &normalized,
7600                row_number,
7601                num_rows,
7602                num_values,
7603            )? {
7604                pages = Some(planned);
7605                break;
7606            }
7607        }
7608        let pages = pages.ok_or_else(|| {
7609            Error::invalid_input_source(
7610                format!(
7611                    "No primitive page planner supports field '{}'",
7612                    self.field.name
7613                )
7614                .into(),
7615            )
7616        })?;
7617
7618        let mut tasks = Vec::with_capacity(pages.len());
7619        let ctx = PrimitiveEncodeContext {
7620            column_idx: self.column_index,
7621            field: self.field.clone(),
7622            encoding_metadata: self.encoding_metadata.clone(),
7623            is_simple_validity,
7624            has_repdef_info,
7625        };
7626        for page in pages {
7627            let ctx = ctx.clone();
7628            let page_encodings = self.page_encodings.clone();
7629            let task =
7630                spawn_cpu(move || Self::encode_page(page_encodings.as_ref(), &ctx, page)).boxed();
7631            tasks.push(task);
7632        }
7633        Ok(tasks)
7634    }
7635
7636    fn extract_validity_buf(
7637        array: Arc<dyn Array>,
7638        repdef: &mut RepDefBuilder,
7639        keep_original_array: bool,
7640    ) -> Result<Arc<dyn Array>> {
7641        if let Some(validity) = array.nulls() {
7642            if keep_original_array {
7643                repdef.add_validity_bitmap(validity.clone());
7644            } else {
7645                repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap());
7646            }
7647            // Empty dictionaries have no valid key payload.  Keep the validity until rep/def is
7648            // grouped with the buffered page; `do_flush` then rebuilds the keys against either a
7649            // neighboring values array or an all-empty-page placeholder.
7650            if let Some(dictionary) = array.as_any_dictionary_opt()
7651                && dictionary.values().is_empty()
7652            {
7653                return Ok(array);
7654            }
7655            let data_no_nulls = array.to_data().into_builder().nulls(None).build()?;
7656            Ok(make_array(data_no_nulls))
7657        } else {
7658            repdef.add_no_null(array.len());
7659            Ok(array)
7660        }
7661    }
7662
7663    fn extract_validity(
7664        mut array: Arc<dyn Array>,
7665        repdef: &mut RepDefBuilder,
7666        keep_original_array: bool,
7667    ) -> Result<Arc<dyn Array>> {
7668        match array.data_type() {
7669            DataType::Null => {
7670                repdef.add_validity_bitmap(NullBuffer::new(BooleanBuffer::new_unset(array.len())));
7671                Ok(array)
7672            }
7673            DataType::Dictionary(_, _) => {
7674                array = dict::normalize_dict_nulls(array)?;
7675                array = dict::clear_out_of_range_null_keys(array)?;
7676                Self::extract_validity_buf(array, repdef, keep_original_array)
7677            }
7678            // Extract our validity buf but NOT any child validity bufs. (they will be encoded in
7679            // as part of the values).  Note: for FSL we do not use repdef.add_fsl because we do
7680            // NOT want to increase the repdef depth.
7681            //
7682            // This would be quite catasrophic for something like vector embeddings.  Imagine we
7683            // had thousands of vectors and some were null but no vector contained null items.  If
7684            // we treated the vectors (primitive FSL) like we treat structural FSL we would end up
7685            // with a rep/def value for every single item in the vector.
7686            _ => Self::extract_validity_buf(array, repdef, keep_original_array),
7687        }
7688    }
7689}
7690
7691impl PrimitivePageEncodingBehavior for RejectSparsePrimitiveEncoding {
7692    fn validate_field(&self, field: &Field, metadata: &HashMap<String, String>) -> Result<()> {
7693        if metadata
7694            .get(STRUCTURAL_ENCODING_META_KEY)
7695            .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE))
7696        {
7697            return Err(Error::invalid_input_source(
7698                format!(
7699                    "Field '{}' requests sparse structural encoding, which is not enabled by the selected file format",
7700                    field.name
7701                )
7702                .into(),
7703            ));
7704        }
7705        Ok(())
7706    }
7707}
7708
7709fn plan_dense_primitive_pages(
7710    arrays: &[ArrayRef],
7711    normalized: &NormalizedStructuralPlan,
7712    row_number: u64,
7713    num_rows: u64,
7714    num_values: u64,
7715) -> Result<Vec<PrimitivePageData>> {
7716    let (repdef, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
7717        miniblock::max_repdef_levels_per_chunk,
7718        num_rows,
7719        num_values,
7720    )?;
7721    PrimitiveStructuralEncoder::split_pages_for_miniblock_repdef_budget(
7722        arrays.to_vec(),
7723        repdef,
7724        miniblock_repdef_budget,
7725        row_number,
7726        num_rows,
7727    )
7728}
7729
7730impl PrimitivePageEncodingBehavior for DenseU16PrimitiveEncoding {
7731    fn try_plan_pages(
7732        &self,
7733        _ctx: &PrimitivePlanContext<'_>,
7734        arrays: &[ArrayRef],
7735        normalized: &NormalizedStructuralPlan,
7736        row_number: u64,
7737        num_rows: u64,
7738        num_values: u64,
7739    ) -> Result<Option<Vec<PrimitivePageData>>> {
7740        Ok(Some(plan_dense_primitive_pages(
7741            arrays, normalized, row_number, num_rows, num_values,
7742        )?))
7743    }
7744
7745    fn try_encode_page(
7746        &self,
7747        ctx: &PrimitiveEncodeContext,
7748        page: PrimitivePageData,
7749    ) -> Result<PrimitiveEncodeAttempt> {
7750        if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
7751            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7752        }
7753        Ok(PrimitiveEncodeAttempt::Encoded(
7754            PrimitiveStructuralEncoder::encode_dense_page(
7755                ctx.clone(),
7756                page,
7757                self.compression.clone(),
7758                MiniblockChunkSize::U16,
7759                ComplexNullEncoding::RawLevels,
7760                FixedWidthDictionaryEncoding::Exclude64Bit,
7761            )?,
7762        ))
7763    }
7764}
7765
7766impl PrimitivePageEncodingBehavior for DenseU32PrimitiveEncoding {
7767    fn try_plan_pages(
7768        &self,
7769        _ctx: &PrimitivePlanContext<'_>,
7770        arrays: &[ArrayRef],
7771        normalized: &NormalizedStructuralPlan,
7772        row_number: u64,
7773        num_rows: u64,
7774        num_values: u64,
7775    ) -> Result<Option<Vec<PrimitivePageData>>> {
7776        Ok(Some(plan_dense_primitive_pages(
7777            arrays, normalized, row_number, num_rows, num_values,
7778        )?))
7779    }
7780
7781    fn try_encode_page(
7782        &self,
7783        ctx: &PrimitiveEncodeContext,
7784        page: PrimitivePageData,
7785    ) -> Result<PrimitiveEncodeAttempt> {
7786        if !matches!(&page.structure, PrimitivePageStructure::Dense { .. }) {
7787            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7788        }
7789        Ok(PrimitiveEncodeAttempt::Encoded(
7790            PrimitiveStructuralEncoder::encode_dense_page(
7791                ctx.clone(),
7792                page,
7793                self.compression.clone(),
7794                MiniblockChunkSize::U32,
7795                ComplexNullEncoding::CompressedLevels,
7796                FixedWidthDictionaryEncoding::Include64Bit,
7797            )?,
7798        ))
7799    }
7800}
7801
7802impl PrimitivePageEncodingBehavior for SparsePrimitiveEncoding {
7803    fn try_plan_pages(
7804        &self,
7805        ctx: &PrimitivePlanContext<'_>,
7806        arrays: &[ArrayRef],
7807        normalized: &NormalizedStructuralPlan,
7808        row_number: u64,
7809        num_rows: u64,
7810        num_values: u64,
7811    ) -> Result<Option<Vec<PrimitivePageData>>> {
7812        let requested_encoding = ctx.encoding_metadata.get(STRUCTURAL_ENCODING_META_KEY);
7813        let requests_sparse = requested_encoding
7814            .is_some_and(|requested| requested.eq_ignore_ascii_case(STRUCTURAL_ENCODING_SPARSE));
7815        if requests_sparse {
7816            let plan = sparse::writer::plan(normalized, num_values)?;
7817            if sparse::writer::uses_constant_layout(&plan, ctx.field) {
7818                return Ok(None);
7819            }
7820            return Ok(Some(vec![PrimitivePageData {
7821                arrays: arrays.to_vec(),
7822                structure: PrimitivePageStructure::Sparse {
7823                    plan,
7824                    prepared_values: None,
7825                },
7826                row_number,
7827                num_rows,
7828            }]));
7829        }
7830
7831        let (_, miniblock_repdef_budget) = normalized.serialize_with_miniblock_repdef_budget(
7832            miniblock::max_repdef_levels_per_chunk,
7833            num_rows,
7834            num_values,
7835        )?;
7836        let automatic_sparse = layout::select_automatic_sparse(
7837            requested_encoding.map(String::as_str),
7838            &miniblock_repdef_budget,
7839            || {
7840                let data = DataBlock::from_arrays(arrays, num_values);
7841                if !sparse::writer::supports_value_block(&data) {
7842                    return Ok(None);
7843                }
7844                let prepared_values = match sparse::writer::prepare_values(
7845                    ctx.field,
7846                    self.compression.as_ref(),
7847                    data,
7848                    MiniblockChunkSize::U32,
7849                ) {
7850                    Ok(prepared_values) => prepared_values,
7851                    Err(error) => {
7852                        debug!(
7853                            "Keeping column {} on its dense structural path because sparse value preparation is unavailable: {}",
7854                            ctx.column_idx, error
7855                        );
7856                        return Ok(None);
7857                    }
7858                };
7859                let plan = sparse::writer::plan(normalized, num_values)?;
7860                if sparse::writer::uses_constant_layout(&plan, ctx.field) {
7861                    return Ok(None);
7862                }
7863                Ok(Some((plan, prepared_values)))
7864            },
7865        )?;
7866        Ok(automatic_sparse.map(|(plan, prepared_values)| {
7867            vec![PrimitivePageData {
7868                arrays: arrays.to_vec(),
7869                structure: PrimitivePageStructure::Sparse {
7870                    plan,
7871                    prepared_values: Some(prepared_values),
7872                },
7873                row_number,
7874                num_rows,
7875            }]
7876        }))
7877    }
7878
7879    fn try_encode_page(
7880        &self,
7881        ctx: &PrimitiveEncodeContext,
7882        page: PrimitivePageData,
7883    ) -> Result<PrimitiveEncodeAttempt> {
7884        if !matches!(&page.structure, PrimitivePageStructure::Sparse { .. }) {
7885            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7886        }
7887        let PrimitivePageData {
7888            arrays,
7889            structure:
7890                PrimitivePageStructure::Sparse {
7891                    plan,
7892                    prepared_values,
7893                },
7894            row_number,
7895            num_rows,
7896        } = page
7897        else {
7898            unreachable!()
7899        };
7900        let num_values = arrays.iter().map(|array| array.len() as u64).sum();
7901        log::debug!(
7902            "Encoding column {} with {} visible items ({} rows) using sparse layout",
7903            ctx.column_idx,
7904            num_values,
7905            num_rows
7906        );
7907        Ok(PrimitiveEncodeAttempt::Encoded(
7908            sparse::writer::encode_page(
7909                ctx.column_idx,
7910                &ctx.field,
7911                self.compression.as_ref(),
7912                prepared_values.map_or_else(
7913                    || {
7914                        sparse::writer::SparseValueInput::Unprepared(DataBlock::from_arrays(
7915                            &arrays, num_values,
7916                        ))
7917                    },
7918                    sparse::writer::SparseValueInput::Prepared,
7919                ),
7920                plan,
7921                row_number,
7922                num_rows,
7923                MiniblockChunkSize::U32,
7924            )?,
7925        ))
7926    }
7927}
7928
7929impl PrimitivePageEncodingBehavior for ConstantPrimitiveEncoding {
7930    fn try_encode_page(
7931        &self,
7932        ctx: &PrimitiveEncodeContext,
7933        page: PrimitivePageData,
7934    ) -> Result<PrimitiveEncodeAttempt> {
7935        let PrimitivePageStructure::Dense { repdef, .. } = &page.structure else {
7936            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7937        };
7938        let num_values: u64 = page.arrays.iter().map(|array| array.len() as u64).sum();
7939        if num_values == 0 {
7940            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7941        }
7942        let leaf_validity = PrimitiveStructuralEncoder::leaf_validity(repdef, num_values as usize)?;
7943        if leaf_validity
7944            .as_ref()
7945            .is_some_and(|validity| validity.count_set_bits() == 0)
7946            || matches!(ctx.field.data_type(), DataType::Struct(fields) if fields.is_empty())
7947        {
7948            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7949        }
7950        let Some(scalar) =
7951            PrimitiveStructuralEncoder::find_constant_scalar(&page.arrays, leaf_validity.as_ref())?
7952        else {
7953            return Ok(PrimitiveEncodeAttempt::Unhandled(page));
7954        };
7955        let PrimitivePageData {
7956            structure: PrimitivePageStructure::Dense { repdef, .. },
7957            row_number,
7958            num_rows,
7959            ..
7960        } = page
7961        else {
7962            unreachable!()
7963        };
7964        log::debug!(
7965            "Encoding column {} with {} items ({} rows) using constant layout",
7966            ctx.column_idx,
7967            num_values,
7968            num_rows
7969        );
7970        Ok(PrimitiveEncodeAttempt::Encoded(
7971            constant::encode_constant_page(ctx.column_idx, scalar, repdef, row_number, num_rows)?,
7972        ))
7973    }
7974}
7975
7976impl FieldEncoder for PrimitiveStructuralEncoder {
7977    // Buffers data, if there is enough to write a page then we create an encode task
7978    fn maybe_encode(
7979        &mut self,
7980        array: ArrayRef,
7981        _external_buffers: &mut OutOfLineBuffers,
7982        mut repdef: RepDefBuilder,
7983        row_number: u64,
7984        num_rows: u64,
7985    ) -> Result<Vec<EncodeTask>> {
7986        let array = Self::extract_validity(array, &mut repdef, self.keep_original_array)?;
7987        self.accumulated_repdefs.push(repdef);
7988
7989        if let Some((arrays, row_number, num_rows)) =
7990            self.accumulation_queue.insert(array, row_number, num_rows)
7991        {
7992            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
7993            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
7994        } else {
7995            Ok(vec![])
7996        }
7997    }
7998
7999    // If there is any data left in the buffer then create an encode task from it
8000    fn flush(&mut self, _external_buffers: &mut OutOfLineBuffers) -> Result<Vec<EncodeTask>> {
8001        if let Some((arrays, row_number, num_rows)) = self.accumulation_queue.flush() {
8002            let accumulated_repdefs = std::mem::take(&mut self.accumulated_repdefs);
8003            Ok(self.do_flush(arrays, accumulated_repdefs, row_number, num_rows)?)
8004        } else {
8005            Ok(vec![])
8006        }
8007    }
8008
8009    fn num_columns(&self) -> u32 {
8010        1
8011    }
8012
8013    fn finish(
8014        &mut self,
8015        _external_buffers: &mut OutOfLineBuffers,
8016    ) -> BoxFuture<'_, Result<Vec<crate::encoder::EncodedColumn>>> {
8017        std::future::ready(Ok(vec![EncodedColumn::default()])).boxed()
8018    }
8019}
8020
8021#[cfg(test)]
8022#[allow(clippy::single_range_in_vec_init)]
8023mod tests {
8024    use super::{
8025        CachedPage, CachedPageData, ChunkInstructions, DataBlock, DecodeMiniBlockTask,
8026        DecodePageTask, FixedFullZipDecodeTask, FixedPerValueDecompressor, FixedWidthDataBlock,
8027        FixedWidthDictionaryEncoding, FullZipCacheableState, FullZipDecodeDetails,
8028        FullZipDecodeTaskItem, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler,
8029        LazyLevels, LevelCodec, LevelCursor, LevelPlan, MiniBlockChunk, MiniBlockChunkIndex,
8030        MiniBlockCompressed, MiniblockChunkSize, PageCacheView, PageDataCacheKey,
8031        PageInfoAndScheduler, PageInitialization, PerValueDataBlock, PerValueDecompressor,
8032        PreambleAction, RunEndsBuilder, RunPosition, RunStorage, SimpleAllNullScheduler,
8033        StructuralPageScheduler, StructuralPrimitiveFieldScheduler, VariableFullZipDecoder,
8034        dense_levels_from_block, map_ranges_to_pages, pages_overlapping_ranges,
8035        validate_complex_all_null_levels,
8036    };
8037    use crate::EncodingsIo;
8038    use crate::buffer::LanceBuffer;
8039    use crate::compression::{
8040        BlockCompressor, DefaultDecompressionStrategy, MiniBlockDecompressor,
8041    };
8042    use crate::constants::{
8043        COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY,
8044        DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY,
8045        STRUCTURAL_ENCODING_MINIBLOCK,
8046    };
8047    use crate::data::BlockInfo;
8048    use crate::decoder::{
8049        FilterExpression, PageEncoding, RequestBatch, SchedulerContext, StructuralFieldDecoder,
8050        StructuralFieldScheduler,
8051    };
8052    use crate::encodings::logical::primitive::fullzip::PerValueCompressor;
8053    use crate::encodings::logical::primitive::{
8054        ChunkDrainInstructions, LoadedChunk, PrimitiveStructuralEncoder,
8055        StructuralPrimitiveFieldDecoder,
8056    };
8057    use crate::encodings::physical::rle::{RleDecompressor, RleEncoder, RleRuns, RunLengthWidth};
8058    use crate::encodings::physical::value::{ValueDecompressor, ValueEncoder};
8059    use crate::format::ProtobufUtils21;
8060    use crate::format::pb21;
8061    use crate::format::pb21::compressive_encoding::Compression;
8062    use crate::repdef::build_control_word_iterator;
8063    use crate::testing::TestEncoding;
8064    use crate::testing::{SimulatedScheduler, TestCases, check_round_trip_encoding_of_data};
8065    use arrow_array::{
8066        Array, ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int8Array,
8067        PrimitiveArray, StringArray, UInt8Array, make_array, new_null_array, types::Int32Type,
8068    };
8069    use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer};
8070    use arrow_schema::{DataType, Field as ArrowField};
8071    use std::collections::HashMap;
8072    use std::{collections::VecDeque, ops::Range, sync::Arc};
8073
8074    #[derive(Debug)]
8075    struct CacheMarker(u8);
8076
8077    impl lance_core::cache::DeepSizeOf for CacheMarker {
8078        fn deep_size_of_children(&self, _context: &mut lance_core::cache::Context) -> usize {
8079            0
8080        }
8081    }
8082
8083    impl CachedPageData for CacheMarker {
8084        fn as_arc_any(self: Arc<Self>) -> Arc<dyn std::any::Any + Send + Sync + 'static> {
8085            self
8086        }
8087    }
8088
8089    #[test]
8090    fn test_is_narrow() {
8091        let int8_array = Int8Array::from(vec![1, 2, 3]);
8092        let array_ref: ArrayRef = Arc::new(int8_array);
8093        let block = DataBlock::from_array(array_ref);
8094
8095        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
8096
8097        let string_array = StringArray::from(vec![Some("hello"), Some("world")]);
8098        let block = DataBlock::from_array(string_array);
8099        assert!(PrimitiveStructuralEncoder::is_narrow(&block));
8100
8101        let string_array = StringArray::from(vec![
8102            Some("hello world".repeat(100)),
8103            Some("world".to_string()),
8104        ]);
8105        let block = DataBlock::from_array(string_array);
8106        assert!((!PrimitiveStructuralEncoder::is_narrow(&block)));
8107    }
8108
8109    fn valued_dictionary() -> ArrayRef {
8110        Arc::new(
8111            DictionaryArray::<Int32Type>::try_new(
8112                PrimitiveArray::<Int32Type>::from(vec![0]),
8113                Arc::new(StringArray::from(vec!["a"])),
8114            )
8115            .unwrap(),
8116        )
8117    }
8118
8119    fn empty_dictionary() -> ArrayRef {
8120        new_null_array(
8121            &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8122            1,
8123        )
8124    }
8125
8126    fn null_valued_dictionary() -> ArrayRef {
8127        Arc::new(
8128            DictionaryArray::<Int32Type>::try_new(
8129                PrimitiveArray::<Int32Type>::from(vec![0]),
8130                Arc::new(StringArray::from(vec![None::<&str>])),
8131            )
8132            .unwrap(),
8133        )
8134    }
8135
8136    fn sliced_empty_dictionary_with_hidden_key_payload() -> ArrayRef {
8137        let keys = PrimitiveArray::<Int32Type>::new(
8138            ScalarBuffer::from(vec![5, 7, 9]),
8139            Some(NullBuffer::new(BooleanBuffer::new_unset(3))),
8140        );
8141        let values = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
8142        let dictionary = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
8143        Arc::new(dictionary.slice(1, 1))
8144    }
8145
8146    #[rstest::rstest]
8147    #[case::empty_after_value(vec![valued_dictionary(), empty_dictionary()])]
8148    #[case::empty_before_value(vec![empty_dictionary(), valued_dictionary()])]
8149    #[case::null_value_after_value(vec![valued_dictionary(), null_valued_dictionary()])]
8150    #[case::hidden_sliced_after_value(vec![
8151        valued_dictionary(),
8152        sliced_empty_dictionary_with_hidden_key_payload(),
8153    ])]
8154    #[tokio::test]
8155    async fn test_mixed_valued_and_all_null_dictionary_chunks(#[case] dictionaries: Vec<ArrayRef>) {
8156        check_round_trip_encoding_of_data(
8157            dictionaries,
8158            &TestCases::default()
8159                .with_structural_encodings()
8160                .with_page_sizes(vec![1])
8161                .with_range(0..1)
8162                .with_range(1..2)
8163                .with_indices(vec![0])
8164                .with_indices(vec![1]),
8165            HashMap::new(),
8166        )
8167        .await;
8168    }
8169
8170    #[tokio::test]
8171    async fn test_sliced_empty_dictionary_with_hidden_key_payload() {
8172        check_round_trip_encoding_of_data(
8173            vec![sliced_empty_dictionary_with_hidden_key_payload()],
8174            &TestCases::default()
8175                .with_structural_encodings()
8176                .with_page_sizes(vec![4096]),
8177            HashMap::new(),
8178        )
8179        .await;
8180    }
8181
8182    #[test]
8183    fn test_primitive_decoder_empty_page_queue_returns_error() {
8184        let field = Arc::new(ArrowField::new("vector", DataType::Float32, true));
8185        let mut decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8186
8187        let err = decoder.drain(1).unwrap_err();
8188        assert!(
8189            matches!(&err, lance_core::Error::Internal { .. }),
8190            "expected internal error, got: {err:?}"
8191        );
8192        let message = err.to_string();
8193        for expected in [
8194            "Primitive decoder missing page decoder",
8195            "field 'vector'",
8196            "data_type=Float32",
8197            "requested_rows=1",
8198            "remaining_rows=1",
8199            "rows_drained_in_current=0",
8200            "queued_pages=0",
8201        ] {
8202            assert!(
8203                message.contains(expected),
8204                "expected error to contain {expected:?}, got: {message}"
8205            );
8206        }
8207    }
8208
8209    #[test]
8210    fn test_plan_decoded_bytes_i32() {
8211        use crate::decoder::CANDIDATE_BATCH_SIZES;
8212
8213        let field = Arc::new(ArrowField::new("x", DataType::Int32, false));
8214        let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8215
8216        let rows_remaining = 100u64;
8217        let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap();
8218
8219        for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
8220            let rows = (candidate as u64).min(rows_remaining);
8221            assert_eq!(bytes[i], rows * 4, "candidate batch size {candidate}");
8222        }
8223    }
8224
8225    #[test]
8226    fn test_plan_decoded_bytes_i32_with_nulls() {
8227        use arrow_array::Int32Array;
8228
8229        // Use N == CANDIDATE_BATCH_SIZES[5] == 1024.  At this size both the
8230        // values buffer (1024*4 = 4096 bytes) and the validity bitmap
8231        // (ceil(1024/8) = 128 bytes) are naturally aligned to Arrow's 64-byte
8232        // allocation boundary, so plan_decoded_bytes matches
8233        // get_buffer_memory_size exactly without needing alignment arithmetic.
8234        const N: i32 = 1024;
8235        let array = Int32Array::from_iter((0..N).map(|i| if i % 2 == 0 { Some(i) } else { None }));
8236        let expected = array.get_buffer_memory_size() as u64;
8237
8238        let field = Arc::new(ArrowField::new("x", DataType::Int32, true));
8239        let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8240
8241        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
8242        // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N
8243        assert_eq!(
8244            bytes[5], expected,
8245            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
8246            bytes[5],
8247        );
8248    }
8249
8250    #[test]
8251    fn test_plan_decoded_bytes_nullable_fsl_with_nullable_items() {
8252        // FixedSizeList<nullable i32, DIM> where the FSL itself is also nullable.
8253        // Arrow allocates three buffers:
8254        //   1. FSL null bitmap:       ceil(rows / 8)
8255        //   2. Child i32 null bitmap: ceil(rows * DIM / 8)
8256        //   3. Child i32 values:      rows * DIM * 4
8257        //
8258        // N=1024, DIM=4 keeps all three naturally 64-byte aligned so the
8259        // comparison against get_buffer_memory_size() is exact.
8260        use arrow_array::{FixedSizeListArray, Int32Array};
8261        use arrow_buffer::NullBuffer;
8262
8263        const N: usize = 1024;
8264        const DIM: i32 = 4;
8265
8266        let child_values = Int32Array::from_iter(
8267            (0..(N as i32 * DIM)).map(|i| if i % 2 == 0 { Some(i) } else { None }),
8268        );
8269        let item_field = Arc::new(ArrowField::new("item", DataType::Int32, true));
8270        let fsl_nulls = NullBuffer::from((0..N).map(|i| i % 3 != 0).collect::<Vec<_>>());
8271        let fsl_array = FixedSizeListArray::new(
8272            item_field.clone(),
8273            DIM,
8274            Arc::new(child_values),
8275            Some(fsl_nulls),
8276        );
8277        let expected = fsl_array.get_buffer_memory_size() as u64;
8278
8279        let fsl_type = DataType::FixedSizeList(item_field, DIM);
8280        let field = Arc::new(ArrowField::new("v", fsl_type, true));
8281        let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8282
8283        let bytes = decoder.plan_decoded_bytes(N as u64).unwrap();
8284        // bytes[5] corresponds to CANDIDATE_BATCH_SIZES[5] == 1024 == N
8285        assert_eq!(
8286            bytes[5], expected,
8287            "plan_decoded_bytes({N}) = {} but get_buffer_memory_size = {expected}",
8288            bytes[5],
8289        );
8290    }
8291
8292    #[test]
8293    fn test_plan_decoded_bytes_fixed_size_list_of_f32() {
8294        use crate::decoder::CANDIDATE_BATCH_SIZES;
8295
8296        const DIMENSION: i32 = 8;
8297        // Items are non-nullable, which is the typical case for vector embeddings.
8298        let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false));
8299        let fsl_type = DataType::FixedSizeList(item_field, DIMENSION);
8300        let field = Arc::new(ArrowField::new("vector", fsl_type, true));
8301        let decoder = StructuralPrimitiveFieldDecoder::new(&field, false);
8302
8303        let rows_remaining = 100u64;
8304        let bytes = decoder.plan_decoded_bytes(rows_remaining).unwrap();
8305
8306        // Each FSL row is DIMENSION f32 values = DIMENSION * 4 bytes, plus
8307        // 1 bit of validity bitmap (the field is nullable).
8308        let bytes_per_row = DIMENSION as u64 * 4;
8309        for (i, &candidate) in CANDIDATE_BATCH_SIZES.iter().enumerate() {
8310            let rows = (candidate as u64).min(rows_remaining);
8311            assert_eq!(
8312                bytes[i],
8313                rows * bytes_per_row + rows.div_ceil(8),
8314                "candidate batch size {candidate}"
8315            );
8316        }
8317    }
8318
8319    #[test]
8320    fn test_fullzip_fixed_rejects_non_byte_aligned_values() {
8321        let fixed = FixedWidthDataBlock {
8322            data: LanceBuffer::from(vec![0_u8]),
8323            bits_per_value: 1,
8324            num_values: 8,
8325            block_info: BlockInfo::new(),
8326        };
8327        let repdef = build_control_word_iterator(None, 0, None, 0, u16::MAX, 8);
8328
8329        let Err(err) = PrimitiveStructuralEncoder::serialize_full_zip_fixed(fixed, repdef, 8)
8330        else {
8331            panic!("expected full-zip to reject 1-bit fixed-width values");
8332        };
8333        assert!(
8334            err.to_string().contains("byte aligned"),
8335            "unexpected error: {err}"
8336        );
8337    }
8338
8339    fn decode_fixed_fullzip_no_levels(
8340        decompressor: Arc<dyn FixedPerValueDecompressor>,
8341        data: Vec<FullZipDecodeTaskItem>,
8342        num_rows: usize,
8343        bytes_per_value: usize,
8344    ) -> DataBlock {
8345        Box::new(FixedFullZipDecodeTask {
8346            details: Arc::new(FullZipDecodeDetails {
8347                value_decompressor: PerValueDecompressor::Fixed(decompressor),
8348                def_meaning: Arc::from([]),
8349                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 0),
8350                max_rep: 0,
8351                max_visible_def: u16::MAX,
8352            }),
8353            data,
8354            num_rows,
8355            bytes_per_value,
8356        })
8357        .decode()
8358        .unwrap()
8359        .data
8360    }
8361
8362    #[test]
8363    fn test_fixed_fullzip_decode_preallocates_exact_output_size() {
8364        #[derive(Debug)]
8365        struct IdentityFixedDecompressor;
8366
8367        impl FixedPerValueDecompressor for IdentityFixedDecompressor {
8368            fn decompress(
8369                &self,
8370                data: FixedWidthDataBlock,
8371                num_rows: u64,
8372            ) -> crate::Result<DataBlock> {
8373                assert_eq!(data.num_values, num_rows);
8374                Ok(DataBlock::FixedWidth(data))
8375            }
8376
8377            fn bits_per_value(&self) -> u64 {
8378                32
8379            }
8380
8381            fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8382                num_values.checked_mul(4)
8383            }
8384        }
8385
8386        let make_item = |num_rows: u64| FullZipDecodeTaskItem {
8387            data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
8388                data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]),
8389                bits_per_value: 32,
8390                num_values: num_rows,
8391                block_info: BlockInfo::new(),
8392            }),
8393            rows_in_buf: num_rows,
8394        };
8395
8396        let num_rows = 512;
8397        let decoded = decode_fixed_fullzip_no_levels(
8398            Arc::new(IdentityFixedDecompressor),
8399            vec![make_item(128), make_item(384)],
8400            num_rows,
8401            4,
8402        );
8403        let values = decoded.as_fixed_width_ref().unwrap();
8404        let expected_size = num_rows * 4;
8405        assert_eq!(values.data.len(), expected_size);
8406        assert_eq!(values.data.clone().into_buffer().capacity(), expected_size);
8407    }
8408
8409    #[test]
8410    fn test_fixed_fullzip_decode_falls_back_when_output_size_is_not_exact() {
8411        #[derive(Debug)]
8412        struct FallbackFixedDecompressor;
8413
8414        impl FixedPerValueDecompressor for FallbackFixedDecompressor {
8415            fn decompress(
8416                &self,
8417                data: FixedWidthDataBlock,
8418                num_rows: u64,
8419            ) -> crate::Result<DataBlock> {
8420                assert_eq!(data.num_values, num_rows);
8421                Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
8422                    data: LanceBuffer::from(vec![7_u8; num_rows as usize * 4]),
8423                    bits_per_value: 32,
8424                    num_values: num_rows,
8425                    block_info: BlockInfo::new(),
8426                }))
8427            }
8428
8429            fn bits_per_value(&self) -> u64 {
8430                // This deliberately cannot be multiplied by num_rows. If FullZip treats
8431                // bits_per_value as an exact decoded-size estimate, decoding will fail.
8432                u64::MAX - 7
8433            }
8434        }
8435
8436        let num_rows = 2;
8437        let decoded = decode_fixed_fullzip_no_levels(
8438            Arc::new(FallbackFixedDecompressor),
8439            vec![FullZipDecodeTaskItem {
8440                data: PerValueDataBlock::Fixed(FixedWidthDataBlock {
8441                    data: LanceBuffer::from(vec![0_u8; num_rows * 4]),
8442                    bits_per_value: 32,
8443                    num_values: num_rows as u64,
8444                    block_info: BlockInfo::new(),
8445                }),
8446                rows_in_buf: num_rows as u64,
8447            }],
8448            num_rows,
8449            4,
8450        );
8451        let values = decoded.as_fixed_width_ref().unwrap();
8452        assert_eq!(values.num_values, num_rows as u64);
8453        assert_eq!(values.data.len(), num_rows * 4);
8454    }
8455
8456    #[test]
8457    fn test_fixed_fullzip_real_fsl_preallocates_exact_output_size() {
8458        let num_rows = 64;
8459        let dimension = 32;
8460        let items = Arc::new(Float32Array::from_iter_values(
8461            (0..num_rows * dimension).map(|value| value as f32),
8462        ));
8463        let item_field = Arc::new(ArrowField::new("item", DataType::Float32, false));
8464        let sample = FixedSizeListArray::new(item_field, dimension as i32, items, None);
8465
8466        let (data, compression) = PerValueCompressor::compress(
8467            &ValueEncoder::default(),
8468            DataBlock::from_array(sample.clone()),
8469        )
8470        .unwrap();
8471        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
8472            panic!("expected fixed-size-list compression");
8473        };
8474        let decompressor = ValueDecompressor::from_fsl(fsl.as_ref()).unwrap();
8475        let expected_size = num_rows * dimension * size_of::<f32>();
8476        assert_eq!(
8477            FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64),
8478            Some(expected_size as u64)
8479        );
8480
8481        let decoded = decode_fixed_fullzip_no_levels(
8482            Arc::new(decompressor),
8483            vec![FullZipDecodeTaskItem {
8484                data,
8485                rows_in_buf: num_rows as u64,
8486            }],
8487            num_rows,
8488            dimension * size_of::<f32>(),
8489        );
8490        let fsl = decoded.as_fixed_size_list_ref().unwrap();
8491        let values = fsl.child.as_fixed_width_ref().unwrap();
8492        assert_eq!(values.data.len(), expected_size);
8493        assert_eq!(values.data.clone().into_buffer().capacity(), expected_size);
8494
8495        let decoded_array = make_array(
8496            decoded
8497                .into_arrow(sample.data_type().clone(), true)
8498                .unwrap(),
8499        );
8500        assert_eq!(decoded_array.as_ref(), &sample);
8501    }
8502
8503    #[test]
8504    fn test_fixed_fullzip_nullable_fsl_uses_fallback_end_to_end() {
8505        #[derive(Debug)]
8506        struct NullableFslDecompressor {
8507            inner: ValueDecompressor,
8508        }
8509
8510        impl FixedPerValueDecompressor for NullableFslDecompressor {
8511            fn decompress(
8512                &self,
8513                data: FixedWidthDataBlock,
8514                num_rows: u64,
8515            ) -> crate::Result<DataBlock> {
8516                FixedPerValueDecompressor::decompress(&self.inner, data, num_rows)
8517            }
8518
8519            fn bits_per_value(&self) -> u64 {
8520                // FullZip must not use this physical row width as an exact decoded-size
8521                // estimate for the nullable, multi-buffer Arrow output.
8522                u64::MAX - 7
8523            }
8524
8525            fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8526                FixedPerValueDecompressor::decoded_size_bytes(&self.inner, num_values)
8527            }
8528        }
8529
8530        let num_rows = 64;
8531        let items = Arc::new(UInt8Array::from_iter(
8532            (0..num_rows).map(|value| (value % 3 != 0).then_some(value as u8)),
8533        ));
8534        let item_field = Arc::new(ArrowField::new("item", DataType::UInt8, true));
8535        let sample = FixedSizeListArray::new(item_field, 1, items, None);
8536
8537        let (data, compression) = PerValueCompressor::compress(
8538            &ValueEncoder::default(),
8539            DataBlock::from_array(sample.clone()),
8540        )
8541        .unwrap();
8542        let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else {
8543            panic!("expected fixed-size-list compression");
8544        };
8545        let decompressor = NullableFslDecompressor {
8546            inner: ValueDecompressor::from_fsl(fsl.as_ref()).unwrap(),
8547        };
8548        assert_eq!(
8549            FixedPerValueDecompressor::decoded_size_bytes(&decompressor, num_rows as u64),
8550            None
8551        );
8552
8553        let decoded = decode_fixed_fullzip_no_levels(
8554            Arc::new(decompressor),
8555            vec![FullZipDecodeTaskItem {
8556                data,
8557                rows_in_buf: num_rows as u64,
8558            }],
8559            num_rows,
8560            2,
8561        );
8562        let decoded_array = make_array(
8563            decoded
8564                .into_arrow(sample.data_type().clone(), true)
8565                .unwrap(),
8566        );
8567        assert_eq!(decoded_array.as_ref(), &sample);
8568    }
8569
8570    #[test]
8571    fn test_miniblock_decode_uses_exact_fixed_width_output_size() {
8572        #[derive(Debug)]
8573        struct FixedWidthMiniBlockDecompressor;
8574
8575        impl MiniBlockDecompressor for FixedWidthMiniBlockDecompressor {
8576            fn decompress(
8577                &self,
8578                data: Vec<LanceBuffer>,
8579                num_values: u64,
8580            ) -> crate::Result<DataBlock> {
8581                assert_eq!(data.len(), 1);
8582                Ok(DataBlock::FixedWidth(FixedWidthDataBlock {
8583                    data: data.into_iter().next().unwrap(),
8584                    bits_per_value: 32,
8585                    num_values,
8586                    block_info: BlockInfo::new(),
8587                }))
8588            }
8589
8590            fn decoded_size_bytes(&self, num_values: u64) -> Option<u64> {
8591                num_values.checked_mul(4)
8592            }
8593        }
8594
8595        let num_rows = 512;
8596        let expected_size = num_rows * 4;
8597        let mut chunk_data = Vec::new();
8598        chunk_data.extend_from_slice(&0_u16.to_le_bytes());
8599        chunk_data.extend_from_slice(&(expected_size as u16).to_le_bytes());
8600        let header_padding =
8601            lance_core::utils::bit::pad_bytes::<{ super::MINIBLOCK_ALIGNMENT }>(chunk_data.len());
8602        chunk_data.resize(chunk_data.len() + header_padding, 0);
8603        chunk_data.resize(chunk_data.len() + expected_size as usize, 7);
8604
8605        let task = DecodeMiniBlockTask {
8606            rep_decompressor: None,
8607            def_decompressor: None,
8608            value_decompressor: Arc::new(FixedWidthMiniBlockDecompressor),
8609            dictionary_data: None,
8610            def_meaning: Arc::from([]),
8611            num_buffers: 1,
8612            max_visible_level: 0,
8613            instructions: vec![(
8614                ChunkDrainInstructions {
8615                    chunk_instructions: ChunkInstructions {
8616                        chunk_idx: 0,
8617                        preamble: PreambleAction::Absent,
8618                        rows_to_skip: 0,
8619                        rows_to_take: num_rows,
8620                        take_trailer: false,
8621                    },
8622                    rows_to_skip: 0,
8623                    rows_to_take: num_rows,
8624                    preamble_action: PreambleAction::Absent,
8625                },
8626                LoadedChunk {
8627                    byte_range: 0..chunk_data.len() as u64,
8628                    data: LanceBuffer::from(chunk_data),
8629                    items_in_chunk: num_rows,
8630                    chunk_idx: 0,
8631                },
8632            )],
8633            has_large_chunk: false,
8634        };
8635
8636        let decoded = Box::new(task).decode().unwrap();
8637        let values = decoded.data.as_fixed_width_ref().unwrap();
8638        assert_eq!(values.data.len(), expected_size as usize);
8639        assert_eq!(
8640            values.data.clone().into_buffer().capacity(),
8641            expected_size as usize
8642        );
8643    }
8644
8645    #[test]
8646    fn test_map_range() {
8647        // Null in the middle
8648        // [[A, B, C], [D, E], NULL, [F, G, H]]
8649        let rep = Some(vec![1, 0, 0, 1, 0, 1, 1, 0, 0]);
8650        let def = Some(vec![0, 0, 0, 0, 0, 1, 0, 0, 0]);
8651        let max_visible_def = 0;
8652        let total_items = 8;
8653        let max_rep = 1;
8654
8655        let check = |range, expected_item_range, expected_level_range| {
8656            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8657                range,
8658                rep.as_ref(),
8659                def.as_ref(),
8660                max_rep,
8661                max_visible_def,
8662                total_items,
8663                PreambleAction::Absent,
8664            );
8665            assert_eq!(item_range, expected_item_range);
8666            assert_eq!(level_range, expected_level_range);
8667        };
8668
8669        check(0..1, 0..3, 0..3);
8670        check(1..2, 3..5, 3..5);
8671        check(2..3, 5..5, 5..6);
8672        check(3..4, 5..8, 6..9);
8673        check(0..2, 0..5, 0..5);
8674        check(1..3, 3..5, 3..6);
8675        check(2..4, 5..8, 5..9);
8676        check(0..3, 0..5, 0..6);
8677        check(1..4, 3..8, 3..9);
8678        check(0..4, 0..8, 0..9);
8679
8680        // Null at start
8681        // [NULL, [A, B], [C]]
8682        let rep = Some(vec![1, 1, 0, 1]);
8683        let def = Some(vec![1, 0, 0, 0]);
8684        let max_visible_def = 0;
8685        let total_items = 3;
8686
8687        let check = |range, expected_item_range, expected_level_range| {
8688            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8689                range,
8690                rep.as_ref(),
8691                def.as_ref(),
8692                max_rep,
8693                max_visible_def,
8694                total_items,
8695                PreambleAction::Absent,
8696            );
8697            assert_eq!(item_range, expected_item_range);
8698            assert_eq!(level_range, expected_level_range);
8699        };
8700
8701        check(0..1, 0..0, 0..1);
8702        check(1..2, 0..2, 1..3);
8703        check(2..3, 2..3, 3..4);
8704        check(0..2, 0..2, 0..3);
8705        check(1..3, 0..3, 1..4);
8706        check(0..3, 0..3, 0..4);
8707
8708        // Null at end
8709        // [[A], [B, C], NULL]
8710        let rep = Some(vec![1, 1, 0, 1]);
8711        let def = Some(vec![0, 0, 0, 1]);
8712        let max_visible_def = 0;
8713        let total_items = 3;
8714
8715        let check = |range, expected_item_range, expected_level_range| {
8716            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8717                range,
8718                rep.as_ref(),
8719                def.as_ref(),
8720                max_rep,
8721                max_visible_def,
8722                total_items,
8723                PreambleAction::Absent,
8724            );
8725            assert_eq!(item_range, expected_item_range);
8726            assert_eq!(level_range, expected_level_range);
8727        };
8728
8729        check(0..1, 0..1, 0..1);
8730        check(1..2, 1..3, 1..3);
8731        check(2..3, 3..3, 3..4);
8732        check(0..2, 0..3, 0..3);
8733        check(1..3, 1..3, 1..4);
8734        check(0..3, 0..3, 0..4);
8735
8736        // No nulls, with repetition
8737        // [[A, B], [C, D], [E, F]]
8738        let rep = Some(vec![1, 0, 1, 0, 1, 0]);
8739        let def: Option<&[u16]> = None;
8740        let max_visible_def = 0;
8741        let total_items = 6;
8742
8743        let check = |range, expected_item_range, expected_level_range| {
8744            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8745                range,
8746                rep.as_ref(),
8747                def.as_ref(),
8748                max_rep,
8749                max_visible_def,
8750                total_items,
8751                PreambleAction::Absent,
8752            );
8753            assert_eq!(item_range, expected_item_range);
8754            assert_eq!(level_range, expected_level_range);
8755        };
8756
8757        check(0..1, 0..2, 0..2);
8758        check(1..2, 2..4, 2..4);
8759        check(2..3, 4..6, 4..6);
8760        check(0..2, 0..4, 0..4);
8761        check(1..3, 2..6, 2..6);
8762        check(0..3, 0..6, 0..6);
8763
8764        // No repetition, with nulls (this case is trivial)
8765        // [A, B, NULL, C]
8766        let rep: Option<&[u16]> = None;
8767        let def = Some(vec![0, 0, 1, 0]);
8768        let max_visible_def = 1;
8769        let total_items = 4;
8770
8771        let check = |range, expected_item_range, expected_level_range| {
8772            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8773                range,
8774                rep.as_ref(),
8775                def.as_ref(),
8776                max_rep,
8777                max_visible_def,
8778                total_items,
8779                PreambleAction::Absent,
8780            );
8781            assert_eq!(item_range, expected_item_range);
8782            assert_eq!(level_range, expected_level_range);
8783        };
8784
8785        check(0..1, 0..1, 0..1);
8786        check(1..2, 1..2, 1..2);
8787        check(2..3, 2..3, 2..3);
8788        check(0..2, 0..2, 0..2);
8789        check(1..3, 1..3, 1..3);
8790        check(0..3, 0..3, 0..3);
8791
8792        // Tricky case, this chunk is a continuation and starts with a rep-index = 0
8793        // [[..., A] [B, C], NULL]
8794        //
8795        // What we do will depend on the preamble action
8796        let rep = Some(vec![0, 1, 0, 1]);
8797        let def = Some(vec![0, 0, 0, 1]);
8798        let max_visible_def = 0;
8799        let total_items = 3;
8800
8801        let check = |range, expected_item_range, expected_level_range| {
8802            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8803                range,
8804                rep.as_ref(),
8805                def.as_ref(),
8806                max_rep,
8807                max_visible_def,
8808                total_items,
8809                PreambleAction::Take,
8810            );
8811            assert_eq!(item_range, expected_item_range);
8812            assert_eq!(level_range, expected_level_range);
8813        };
8814
8815        // If we are taking the preamble then the range must start at 0
8816        check(0..1, 0..3, 0..3);
8817        check(0..2, 0..3, 0..4);
8818
8819        let check = |range, expected_item_range, expected_level_range| {
8820            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8821                range,
8822                rep.as_ref(),
8823                def.as_ref(),
8824                max_rep,
8825                max_visible_def,
8826                total_items,
8827                PreambleAction::Skip,
8828            );
8829            assert_eq!(item_range, expected_item_range);
8830            assert_eq!(level_range, expected_level_range);
8831        };
8832
8833        check(0..1, 1..3, 1..3);
8834        check(1..2, 3..3, 3..4);
8835        check(0..2, 1..3, 1..4);
8836
8837        // Another preamble case but now it doesn't end with a new list
8838        // [[..., A], NULL, [D, E]]
8839        //
8840        // What we do will depend on the preamble action
8841        let rep = Some(vec![0, 1, 1, 0]);
8842        let def = Some(vec![0, 1, 0, 0]);
8843        let max_visible_def = 0;
8844        let total_items = 4;
8845
8846        let check = |range, expected_item_range, expected_level_range| {
8847            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8848                range,
8849                rep.as_ref(),
8850                def.as_ref(),
8851                max_rep,
8852                max_visible_def,
8853                total_items,
8854                PreambleAction::Take,
8855            );
8856            assert_eq!(item_range, expected_item_range);
8857            assert_eq!(level_range, expected_level_range);
8858        };
8859
8860        // If we are taking the preamble then the range must start at 0
8861        check(0..1, 0..1, 0..2);
8862        check(0..2, 0..3, 0..4);
8863
8864        let check = |range, expected_item_range, expected_level_range| {
8865            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8866                range,
8867                rep.as_ref(),
8868                def.as_ref(),
8869                max_rep,
8870                max_visible_def,
8871                total_items,
8872                PreambleAction::Skip,
8873            );
8874            assert_eq!(item_range, expected_item_range);
8875            assert_eq!(level_range, expected_level_range);
8876        };
8877
8878        // If we are taking the preamble then the range must start at 0
8879        check(0..1, 1..1, 1..2);
8880        check(1..2, 1..3, 2..4);
8881        check(0..2, 1..3, 1..4);
8882
8883        // Now a preamble case without any definition levels
8884        // [[..., A] [B, C], [D]]
8885        let rep = Some(vec![0, 1, 0, 1]);
8886        let def: Option<Vec<u16>> = None;
8887        let max_visible_def = 0;
8888        let total_items = 4;
8889
8890        let check = |range, expected_item_range, expected_level_range| {
8891            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8892                range,
8893                rep.as_ref(),
8894                def.as_ref(),
8895                max_rep,
8896                max_visible_def,
8897                total_items,
8898                PreambleAction::Take,
8899            );
8900            assert_eq!(item_range, expected_item_range);
8901            assert_eq!(level_range, expected_level_range);
8902        };
8903
8904        // If we are taking the preamble then the range must start at 0
8905        check(0..1, 0..3, 0..3);
8906        check(0..2, 0..4, 0..4);
8907
8908        let check = |range, expected_item_range, expected_level_range| {
8909            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8910                range,
8911                rep.as_ref(),
8912                def.as_ref(),
8913                max_rep,
8914                max_visible_def,
8915                total_items,
8916                PreambleAction::Skip,
8917            );
8918            assert_eq!(item_range, expected_item_range);
8919            assert_eq!(level_range, expected_level_range);
8920        };
8921
8922        check(0..1, 1..3, 1..3);
8923        check(1..2, 3..4, 3..4);
8924        check(0..2, 1..4, 1..4);
8925
8926        // If we have nested lists then non-top level lists may be empty/null
8927        // and we need to make sure we still handle them as invisible items (we
8928        // failed to do this previously)
8929        let rep = Some(vec![2, 1, 2, 0, 1, 2]);
8930        let def = Some(vec![0, 1, 2, 0, 0, 0]);
8931        let max_rep = 2;
8932        let max_visible_def = 0;
8933        let total_items = 4;
8934
8935        let check = |range, expected_item_range, expected_level_range| {
8936            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8937                range,
8938                rep.as_ref(),
8939                def.as_ref(),
8940                max_rep,
8941                max_visible_def,
8942                total_items,
8943                PreambleAction::Absent,
8944            );
8945            assert_eq!(item_range, expected_item_range);
8946            assert_eq!(level_range, expected_level_range);
8947        };
8948
8949        check(0..3, 0..4, 0..6);
8950        check(0..1, 0..1, 0..2);
8951        check(1..2, 1..3, 2..5);
8952        check(2..3, 3..4, 5..6);
8953
8954        // Invisible items in a preamble that we are taking (regressing a previous failure)
8955        let rep = Some(vec![0, 0, 1, 0, 1, 1]);
8956        let def = Some(vec![0, 1, 0, 0, 0, 0]);
8957        let max_rep = 1;
8958        let max_visible_def = 0;
8959        let total_items = 5;
8960
8961        let check = |range, expected_item_range, expected_level_range| {
8962            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8963                range,
8964                rep.as_ref(),
8965                def.as_ref(),
8966                max_rep,
8967                max_visible_def,
8968                total_items,
8969                PreambleAction::Take,
8970            );
8971            assert_eq!(item_range, expected_item_range);
8972            assert_eq!(level_range, expected_level_range);
8973        };
8974
8975        check(0..0, 0..1, 0..2);
8976        check(0..1, 0..3, 0..4);
8977        check(0..2, 0..4, 0..5);
8978
8979        // Skip preamble (with invis items) and skip a few rows (with invis items)
8980        // and then take a few rows but not all the rows
8981        let rep = Some(vec![0, 1, 0, 1, 0, 1, 0, 1]);
8982        let def = Some(vec![1, 0, 1, 1, 0, 0, 0, 0]);
8983        let max_rep = 1;
8984        let max_visible_def = 0;
8985        let total_items = 5;
8986
8987        let check = |range, expected_item_range, expected_level_range| {
8988            let (item_range, level_range) = DecodeMiniBlockTask::map_range(
8989                range,
8990                rep.as_ref(),
8991                def.as_ref(),
8992                max_rep,
8993                max_visible_def,
8994                total_items,
8995                PreambleAction::Skip,
8996            );
8997            assert_eq!(item_range, expected_item_range);
8998            assert_eq!(level_range, expected_level_range);
8999        };
9000
9001        check(2..3, 2..4, 5..7);
9002    }
9003
9004    #[test]
9005    fn test_slice_batch_data_and_rebase_offsets_u32() {
9006        let data = LanceBuffer::copy_slice(b"0123456789abcdefghij");
9007        let offsets = LanceBuffer::reinterpret_vec(vec![6_u32, 8_u32, 8_u32, 12_u32]);
9008
9009        let (sliced_data, normalized_offsets) =
9010            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
9011                .unwrap();
9012
9013        assert_eq!(sliced_data.as_ref(), b"6789ab");
9014        let normalized = normalized_offsets.borrow_to_typed_slice::<u32>();
9015        assert_eq!(normalized.as_ref(), &[0, 2, 2, 6]);
9016    }
9017
9018    #[test]
9019    fn test_slice_batch_data_and_rebase_offsets_u64() {
9020        let data = LanceBuffer::copy_slice(b"abcdefghijklmnopqrstuvwxyz");
9021        let offsets = LanceBuffer::reinterpret_vec(vec![10_u64, 12_u64, 16_u64, 20_u64]);
9022
9023        let (sliced_data, normalized_offsets) =
9024            VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 64)
9025                .unwrap();
9026
9027        assert_eq!(sliced_data.as_ref(), b"klmnopqrst");
9028        let normalized = normalized_offsets.borrow_to_typed_slice::<u64>();
9029        assert_eq!(normalized.as_ref(), &[0, 2, 6, 10]);
9030    }
9031
9032    #[test]
9033    fn test_slice_batch_data_and_rebase_offsets_rejects_invalid_offsets() {
9034        let data = LanceBuffer::copy_slice(b"abcd");
9035        let offsets = LanceBuffer::reinterpret_vec(vec![3_u32, 2_u32]);
9036
9037        let err = VariableFullZipDecoder::slice_batch_data_and_rebase_offsets(&data, &offsets, 32)
9038            .expect_err("offset end before start should error");
9039        assert!(err.to_string().contains("less than base"));
9040    }
9041
9042    #[test]
9043    fn test_schedule_instructions() {
9044        // Convert repetition index to bytes for testing
9045        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
9046        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9047        let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9048
9049        let check = |user_ranges, expected_instructions| {
9050            let instructions = ChunkInstructions::schedule_instructions(&chunk_index, user_ranges);
9051            assert_eq!(instructions, expected_instructions);
9052        };
9053
9054        // The instructions we expect if we're grabbing the whole range
9055        let expected_take_all = vec![
9056            ChunkInstructions {
9057                chunk_idx: 0,
9058                preamble: PreambleAction::Absent,
9059                rows_to_skip: 0,
9060                rows_to_take: 6,
9061                take_trailer: true,
9062            },
9063            ChunkInstructions {
9064                chunk_idx: 1,
9065                preamble: PreambleAction::Take,
9066                rows_to_skip: 0,
9067                rows_to_take: 2,
9068                take_trailer: false,
9069            },
9070            ChunkInstructions {
9071                chunk_idx: 2,
9072                preamble: PreambleAction::Absent,
9073                rows_to_skip: 0,
9074                rows_to_take: 5,
9075                take_trailer: true,
9076            },
9077            ChunkInstructions {
9078                chunk_idx: 3,
9079                preamble: PreambleAction::Take,
9080                rows_to_skip: 0,
9081                rows_to_take: 1,
9082                take_trailer: false,
9083            },
9084        ];
9085
9086        // Take all as 1 range
9087        check(&[0..14], expected_take_all.clone());
9088
9089        // Take all a individual rows
9090        check(
9091            &[
9092                0..1,
9093                1..2,
9094                2..3,
9095                3..4,
9096                4..5,
9097                5..6,
9098                6..7,
9099                7..8,
9100                8..9,
9101                9..10,
9102                10..11,
9103                11..12,
9104                12..13,
9105                13..14,
9106            ],
9107            expected_take_all,
9108        );
9109
9110        // Test some partial takes
9111
9112        // 2 rows in the same chunk but not contiguous
9113        check(
9114            &[0..1, 3..4],
9115            vec![
9116                ChunkInstructions {
9117                    chunk_idx: 0,
9118                    preamble: PreambleAction::Absent,
9119                    rows_to_skip: 0,
9120                    rows_to_take: 1,
9121                    take_trailer: false,
9122                },
9123                ChunkInstructions {
9124                    chunk_idx: 0,
9125                    preamble: PreambleAction::Absent,
9126                    rows_to_skip: 3,
9127                    rows_to_take: 1,
9128                    take_trailer: false,
9129                },
9130            ],
9131        );
9132
9133        // Taking just a trailer/preamble
9134        check(
9135            &[5..6],
9136            vec![
9137                ChunkInstructions {
9138                    chunk_idx: 0,
9139                    preamble: PreambleAction::Absent,
9140                    rows_to_skip: 5,
9141                    rows_to_take: 1,
9142                    take_trailer: true,
9143                },
9144                ChunkInstructions {
9145                    chunk_idx: 1,
9146                    preamble: PreambleAction::Take,
9147                    rows_to_skip: 0,
9148                    rows_to_take: 0,
9149                    take_trailer: false,
9150                },
9151            ],
9152        );
9153
9154        // Skipping an entire chunk
9155        check(
9156            &[7..10],
9157            vec![
9158                ChunkInstructions {
9159                    chunk_idx: 1,
9160                    preamble: PreambleAction::Skip,
9161                    rows_to_skip: 1,
9162                    rows_to_take: 1,
9163                    take_trailer: false,
9164                },
9165                ChunkInstructions {
9166                    chunk_idx: 2,
9167                    preamble: PreambleAction::Absent,
9168                    rows_to_skip: 0,
9169                    rows_to_take: 2,
9170                    take_trailer: false,
9171                },
9172            ],
9173        );
9174    }
9175
9176    #[test]
9177    fn test_drain_instructions() {
9178        fn drain_from_instructions(
9179            instructions: &mut VecDeque<ChunkInstructions>,
9180            mut rows_desired: u64,
9181            need_preamble: &mut bool,
9182            skip_in_chunk: &mut u64,
9183        ) -> Vec<ChunkDrainInstructions> {
9184            // Note: instructions.len() is an upper bound, we typically take much fewer
9185            let mut drain_instructions = Vec::with_capacity(instructions.len());
9186            while rows_desired > 0 || *need_preamble {
9187                let (next_instructions, consumed_chunk) = instructions
9188                    .front()
9189                    .unwrap()
9190                    .drain_from_instruction(&mut rows_desired, need_preamble, skip_in_chunk);
9191                if consumed_chunk {
9192                    instructions.pop_front();
9193                }
9194                drain_instructions.push(next_instructions);
9195            }
9196            drain_instructions
9197        }
9198
9199        // Convert repetition index to bytes for testing
9200        let rep_data: Vec<u64> = vec![5, 2, 3, 0, 4, 7, 2, 0];
9201        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9202        let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9203        let user_ranges = vec![1..7, 10..14];
9204
9205        // First, schedule the ranges
9206        let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
9207
9208        let mut to_drain = VecDeque::from(scheduled.clone());
9209
9210        // Now we drain in batches of 4
9211
9212        let mut need_preamble = false;
9213        let mut skip_in_chunk = 0;
9214
9215        let next_batch =
9216            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
9217
9218        assert!(!need_preamble);
9219        assert_eq!(skip_in_chunk, 4);
9220        assert_eq!(
9221            next_batch,
9222            vec![ChunkDrainInstructions {
9223                chunk_instructions: scheduled[0].clone(),
9224                rows_to_take: 4,
9225                rows_to_skip: 0,
9226                preamble_action: PreambleAction::Absent,
9227            }]
9228        );
9229
9230        let next_batch =
9231            drain_from_instructions(&mut to_drain, 4, &mut need_preamble, &mut skip_in_chunk);
9232
9233        assert!(!need_preamble);
9234        assert_eq!(skip_in_chunk, 2);
9235
9236        assert_eq!(
9237            next_batch,
9238            vec![
9239                ChunkDrainInstructions {
9240                    chunk_instructions: scheduled[0].clone(),
9241                    rows_to_take: 1,
9242                    rows_to_skip: 4,
9243                    preamble_action: PreambleAction::Absent,
9244                },
9245                ChunkDrainInstructions {
9246                    chunk_instructions: scheduled[1].clone(),
9247                    rows_to_take: 1,
9248                    rows_to_skip: 0,
9249                    preamble_action: PreambleAction::Take,
9250                },
9251                ChunkDrainInstructions {
9252                    chunk_instructions: scheduled[2].clone(),
9253                    rows_to_take: 2,
9254                    rows_to_skip: 0,
9255                    preamble_action: PreambleAction::Absent,
9256                }
9257            ]
9258        );
9259
9260        let next_batch =
9261            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
9262
9263        assert!(!need_preamble);
9264        assert_eq!(skip_in_chunk, 0);
9265
9266        assert_eq!(
9267            next_batch,
9268            vec![
9269                ChunkDrainInstructions {
9270                    chunk_instructions: scheduled[2].clone(),
9271                    rows_to_take: 1,
9272                    rows_to_skip: 2,
9273                    preamble_action: PreambleAction::Absent,
9274                },
9275                ChunkDrainInstructions {
9276                    chunk_instructions: scheduled[3].clone(),
9277                    rows_to_take: 1,
9278                    rows_to_skip: 0,
9279                    preamble_action: PreambleAction::Take,
9280                },
9281            ]
9282        );
9283
9284        // Regression case.  Need a chunk with preamble, rows, and trailer (the middle chunk here)
9285        let rep_data: Vec<u64> = vec![5, 2, 3, 3, 20, 0];
9286        let rep_bytes: Vec<u8> = rep_data.iter().flat_map(|v| v.to_le_bytes()).collect();
9287        let chunk_index = MiniBlockChunkIndex::new_nested_for_test(&rep_bytes, 2);
9288        let user_ranges = vec![0..28];
9289
9290        // First, schedule the ranges
9291        let scheduled = ChunkInstructions::schedule_instructions(&chunk_index, &user_ranges);
9292
9293        let mut to_drain = VecDeque::from(scheduled.clone());
9294
9295        // Drain first chunk and some of second chunk
9296
9297        let mut need_preamble = false;
9298        let mut skip_in_chunk = 0;
9299
9300        let next_batch =
9301            drain_from_instructions(&mut to_drain, 7, &mut need_preamble, &mut skip_in_chunk);
9302
9303        assert_eq!(
9304            next_batch,
9305            vec![
9306                ChunkDrainInstructions {
9307                    chunk_instructions: scheduled[0].clone(),
9308                    rows_to_take: 6,
9309                    rows_to_skip: 0,
9310                    preamble_action: PreambleAction::Absent,
9311                },
9312                ChunkDrainInstructions {
9313                    chunk_instructions: scheduled[1].clone(),
9314                    rows_to_take: 1,
9315                    rows_to_skip: 0,
9316                    preamble_action: PreambleAction::Take,
9317                },
9318            ]
9319        );
9320
9321        assert!(!need_preamble);
9322        assert_eq!(skip_in_chunk, 1);
9323
9324        // Now, the tricky part.  We drain the second chunk, including the trailer, and need to make sure
9325        // we get a drain task to take the preamble of the third chunk (and nothing else)
9326        let next_batch =
9327            drain_from_instructions(&mut to_drain, 2, &mut need_preamble, &mut skip_in_chunk);
9328
9329        assert_eq!(
9330            next_batch,
9331            vec![
9332                ChunkDrainInstructions {
9333                    chunk_instructions: scheduled[1].clone(),
9334                    rows_to_take: 2,
9335                    rows_to_skip: 1,
9336                    preamble_action: PreambleAction::Skip,
9337                },
9338                ChunkDrainInstructions {
9339                    chunk_instructions: scheduled[2].clone(),
9340                    rows_to_take: 0,
9341                    rows_to_skip: 0,
9342                    preamble_action: PreambleAction::Take,
9343                },
9344            ]
9345        );
9346
9347        assert!(!need_preamble);
9348        assert_eq!(skip_in_chunk, 0);
9349    }
9350
9351    use super::chunk_index::{PrefixSums, RowMapping};
9352    use super::{MINIBLOCK_ALIGNMENT, Words, build_chunk_index};
9353    use bytes::Bytes;
9354    use lance_core::cache::{Context, DeepSizeOf};
9355    use rstest::rstest;
9356
9357    /// Builds a `Words` metadata buffer (u16 words) from `(log_num_values, num_bytes)`
9358    /// pairs, returning the words and the total data-buffer size.
9359    fn words_from(entries: &[(u32, u32)]) -> (Words, u64) {
9360        let mut raw = Vec::with_capacity(entries.len() * 2);
9361        let mut total = 0u64;
9362        for &(log, num_bytes) in entries {
9363            assert!(num_bytes > 0 && num_bytes % MINIBLOCK_ALIGNMENT as u32 == 0);
9364            let divided = num_bytes / MINIBLOCK_ALIGNMENT as u32 - 1;
9365            let word = (divided << 4) | log;
9366            assert!(word <= u16::MAX as u32, "test word {word} exceeds u16");
9367            raw.extend_from_slice(&(word as u16).to_le_bytes());
9368            total += num_bytes as u64;
9369        }
9370        (Words::from_bytes(Bytes::from(raw), false).unwrap(), total)
9371    }
9372
9373    fn rep_bytes_from(values: &[u64]) -> Vec<u8> {
9374        values.iter().flat_map(|v| v.to_le_bytes()).collect()
9375    }
9376
9377    #[rstest]
9378    // Two full chunks of 8 values (log 3) plus a partial last chunk; byte sizes vary
9379    // independently of value counts.
9380    #[case::uniform_partial_last(&[(3, 16), (3, 24), (0, 8)], 19, "uniform_flat", 8, 3)]
9381    // Single chunk covers the whole page.
9382    #[case::single_chunk(&[(0, 24)], 5, "uniform_flat", 5, 5)]
9383    // Last chunk is also full (exact multiple).
9384    #[case::exact_multiple(&[(3, 16), (3, 16)], 16, "uniform_flat", 8, 8)]
9385    // Non-last chunks differ in size, so this is a non-uniform flat page.
9386    #[case::non_uniform(&[(4, 16), (2, 16), (0, 8)], 21, "flat", 16, 1)]
9387    fn test_flat_detection(
9388        #[case] entries: &[(u32, u32)],
9389        #[case] items_in_page: u64,
9390        #[case] expected_kind: &str,
9391        #[case] expected_first_items: u64,
9392        #[case] expected_last_items: u64,
9393    ) {
9394        let base = 100u64;
9395        let (words, data_buf_size) = words_from(entries);
9396        let index = build_chunk_index(&words, items_in_page, base, data_buf_size, None, 0).unwrap();
9397
9398        assert_eq!(index.row_mapping_debug(), expected_kind);
9399        assert_eq!(index.num_chunks(), entries.len());
9400        assert_eq!(index.items_in_chunk(0), expected_first_items);
9401        assert_eq!(index.items_in_chunk(entries.len() - 1), expected_last_items);
9402
9403        // Byte ranges are absolute, contiguous, and exactly cover the data buffer.
9404        let mut expected_start = base;
9405        for (i, &(_, num_bytes)) in entries.iter().enumerate() {
9406            let range = index.byte_range(i);
9407            assert_eq!(range.start, expected_start);
9408            assert_eq!(range.end - range.start, num_bytes as u64);
9409            expected_start = range.end;
9410        }
9411        assert_eq!(expected_start, base + data_buf_size);
9412
9413        // For flat pages rows == items, so the per-chunk items sum to the page total.
9414        let total_items: u64 = (0..index.num_chunks())
9415            .map(|i| index.items_in_chunk(i))
9416            .sum();
9417        assert_eq!(total_items, items_in_page);
9418    }
9419
9420    #[test]
9421    fn test_nested_detection_and_axes() {
9422        // Repetition index (stride 2): three chunks holding 5, 4, 3 rows, no trailers.
9423        let rep = rep_bytes_from(&[5, 0, 4, 0, 3, 0]);
9424
9425        // Uniform leaf chunking: value counts 4, 4, 2.
9426        let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9427        let index = build_chunk_index(&words, 10, 0, data_buf_size, Some(&rep), 1).unwrap();
9428        assert_eq!(index.row_mapping_debug(), "nested");
9429        assert_eq!(index.num_chunks(), 3);
9430        // Rows come from the repetition index, not the value counts.
9431        assert_eq!(index.first_row(0), 0);
9432        assert_eq!(index.rows_in_chunk(0), 5);
9433        assert_eq!(index.first_row(1), 5);
9434        assert_eq!(index.rows_in_chunk(1), 4);
9435        assert_eq!(index.first_row(2), 9);
9436        assert_eq!(index.rows_in_chunk(2), 3);
9437        // Items come from the value words.
9438        assert_eq!(index.items_in_chunk(0), 4);
9439        assert_eq!(index.items_in_chunk(1), 4);
9440        assert_eq!(index.items_in_chunk(2), 2);
9441
9442        // Non-uniform leaf chunking: value counts 8, 2, 5.
9443        let (words_nu, dbs_nu) = words_from(&[(3, 8), (1, 8), (0, 8)]);
9444        let index_nu = build_chunk_index(&words_nu, 15, 0, dbs_nu, Some(&rep), 1).unwrap();
9445        assert_eq!(index_nu.row_mapping_debug(), "nested");
9446        assert_eq!(index_nu.items_in_chunk(0), 8);
9447        assert_eq!(index_nu.items_in_chunk(1), 2);
9448        assert_eq!(index_nu.items_in_chunk(2), 5);
9449        // The row axis is unchanged by the leaf chunking.
9450        assert_eq!(index_nu.rows_in_chunk(0), 5);
9451    }
9452
9453    #[test]
9454    fn test_uniform_flat_matches_prefix_sum_flat() {
9455        // Distribution: 4 chunks of 4 values, last chunk 3 (15 items total).
9456        let (words, data_buf_size) = words_from(&[(2, 8), (2, 8), (2, 8), (0, 8)]);
9457        let uniform = build_chunk_index(&words, 15, 0, data_buf_size, None, 0).unwrap();
9458        assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
9459
9460        // The same distribution expressed as a non-uniform Flat prefix-sum index.
9461        let byte_starts = PrefixSums::from_deltas([8u64, 8, 8, 8].into_iter(), 4, 32);
9462        let value_starts = PrefixSums::from_deltas([4u64, 4, 4, 3].into_iter(), 4, 15);
9463        let flat = MiniBlockChunkIndex::new(0, byte_starts, RowMapping::Flat { value_starts });
9464        assert_eq!(flat.row_mapping_debug(), "flat");
9465
9466        // Lookup parity: identical byte ranges and item counts.
9467        for i in 0..4 {
9468            assert_eq!(uniform.byte_range(i), flat.byte_range(i));
9469            assert_eq!(uniform.items_in_chunk(i), flat.items_in_chunk(i));
9470        }
9471
9472        // Scheduler parity across scan / single-row / partial / scattered multi-range.
9473        let range_sets: Vec<Vec<std::ops::Range<u64>>> = vec![
9474            vec![0..15],
9475            vec![0..1],
9476            vec![7..8],
9477            vec![14..15],
9478            vec![3..10],
9479            vec![0..2, 5..6, 12..15],
9480        ];
9481        for ranges in &range_sets {
9482            let from_uniform = ChunkInstructions::schedule_instructions(&uniform, ranges);
9483            let from_flat = ChunkInstructions::schedule_instructions(&flat, ranges);
9484            assert_eq!(from_uniform, from_flat, "mismatch for ranges {ranges:?}");
9485        }
9486
9487        // A full scan yields one Absent, no-trailer instruction per chunk.
9488        let full = ChunkInstructions::schedule_instructions(&uniform, &[0..15]);
9489        assert_eq!(full.len(), 4);
9490        for (i, inst) in full.iter().enumerate() {
9491            assert_eq!(inst.chunk_idx, i);
9492            assert_eq!(inst.preamble, PreambleAction::Absent);
9493            assert_eq!(inst.rows_to_skip, 0);
9494            assert!(!inst.take_trailer);
9495        }
9496        assert_eq!(full.iter().map(|i| i.rows_to_take).sum::<u64>(), 15);
9497    }
9498
9499    #[test]
9500    fn test_deep_size_per_variant_below_legacy() {
9501        // The previous representation cached 48 bytes per chunk (24 for ChunkMeta plus
9502        // 24 for a rep-index block); every variant's heap must be well below that.
9503        const LEGACY_PER_CHUNK: usize = 48;
9504        let num_chunks = 3;
9505        let heap = |index: &MiniBlockChunkIndex| index.deep_size_of_children(&mut Context::new());
9506
9507        let (uniform_words, uniform_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9508        let uniform = build_chunk_index(&uniform_words, 10, 0, uniform_dbs, None, 0).unwrap();
9509        assert_eq!(uniform.row_mapping_debug(), "uniform_flat");
9510        assert!(heap(&uniform) < LEGACY_PER_CHUNK * num_chunks);
9511
9512        let (flat_words, flat_dbs) = words_from(&[(3, 8), (1, 8), (0, 8)]);
9513        let flat = build_chunk_index(&flat_words, 11, 0, flat_dbs, None, 0).unwrap();
9514        assert_eq!(flat.row_mapping_debug(), "flat");
9515        assert!(heap(&flat) < LEGACY_PER_CHUNK * num_chunks);
9516        // Flat carries a value-starts array that UniformFlat derives arithmetically.
9517        assert!(heap(&flat) > heap(&uniform));
9518
9519        let rep = rep_bytes_from(&[4, 0, 3, 0, 3, 0]);
9520        let (nested_words, nested_dbs) = words_from(&[(2, 8), (2, 8), (0, 8)]);
9521        let nested = build_chunk_index(&nested_words, 10, 0, nested_dbs, Some(&rep), 1).unwrap();
9522        assert_eq!(nested.row_mapping_debug(), "nested");
9523        assert!(heap(&nested) < LEGACY_PER_CHUNK * num_chunks);
9524    }
9525
9526    #[tokio::test]
9527    async fn test_fullzip_initialize_is_lazy() {
9528        use futures::{FutureExt, future::BoxFuture};
9529        use std::ops::Range;
9530        use std::sync::Mutex;
9531
9532        #[derive(Debug, Clone)]
9533        struct RecordingScheduler {
9534            data: bytes::Bytes,
9535            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
9536        }
9537
9538        impl RecordingScheduler {
9539            fn new(data: bytes::Bytes) -> Self {
9540                Self {
9541                    data,
9542                    requests: Arc::new(Mutex::new(Vec::new())),
9543                }
9544            }
9545
9546            fn requests(&self) -> Vec<Vec<Range<u64>>> {
9547                self.requests.lock().unwrap().clone()
9548            }
9549        }
9550
9551        impl crate::EncodingsIo for RecordingScheduler {
9552            fn submit_request(
9553                &self,
9554                ranges: Vec<Range<u64>>,
9555                _priority: u64,
9556            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
9557                self.requests.lock().unwrap().push(ranges.clone());
9558                let data = ranges
9559                    .into_iter()
9560                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
9561                    .collect::<Vec<_>>();
9562                std::future::ready(Ok(data)).boxed()
9563            }
9564        }
9565
9566        #[derive(Debug)]
9567        struct TestFixedDecompressor;
9568
9569        impl FixedPerValueDecompressor for TestFixedDecompressor {
9570            fn decompress(
9571                &self,
9572                _data: FixedWidthDataBlock,
9573                _num_rows: u64,
9574            ) -> crate::Result<DataBlock> {
9575                unimplemented!("Test decompressor")
9576            }
9577
9578            fn bits_per_value(&self) -> u64 {
9579                32
9580            }
9581        }
9582
9583        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(vec![
9584            0;
9585            16 * 1024
9586        ])));
9587        let mut scheduler = FullZipScheduler {
9588            data_buf_position: 0,
9589            data_buf_size: 4096,
9590            rep_index: Some(FullZipRepIndexDetails {
9591                buf_position: 1000,
9592                bytes_per_value: 4,
9593            }),
9594            priority: 0,
9595            rows_in_page: 100,
9596            bits_per_offset: 32,
9597            details: Arc::new(FullZipDecodeDetails {
9598                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
9599                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
9600                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
9601                max_rep: 0,
9602                max_visible_def: 0,
9603            }),
9604            cached_state: None,
9605            enable_cache: false,
9606        };
9607
9608        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
9609        let initialization = scheduler.init_layout().unwrap();
9610        assert!(
9611            initialization.required_ranges().next().is_none(),
9612            "FullZip should not request any ranges when caching is disabled"
9613        );
9614        let buffers = initialization.with_buffers(vec![]).unwrap();
9615        let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap();
9616
9617        assert!(
9618            cached_data
9619                .as_arc_any()
9620                .downcast_ref::<super::NoCachedPageData>()
9621                .is_some(),
9622            "FullZip initialize should not eagerly load repetition index data"
9623        );
9624        assert!(scheduler.cached_state.is_none());
9625        assert!(
9626            io.requests().is_empty(),
9627            "FullZip initialize should not issue any I/O"
9628        );
9629    }
9630
9631    /// Prove that an N-page cache miss issues one `submit_request` carrying every
9632    /// page's metadata range, rather than one submission per page. The production
9633    /// I/O scheduler can then coalesce adjacent ranges within that batch.
9634    #[tokio::test]
9635    async fn test_initialize_coalesces_missed_page_metadata() {
9636        use std::ops::Range;
9637
9638        use futures::FutureExt;
9639        use futures::future::BoxFuture;
9640
9641        use super::{
9642            CachedPage, CachedPageData, Error, PageCacheView, PageDataCacheKey,
9643            PageInfoAndScheduler, PageInitialization, PageInitializationBuffers, PageLoadTask,
9644            StructuralPrimitiveFieldScheduler,
9645        };
9646        use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler};
9647
9648        // Give each fake page one fixed-size metadata buffer. The ranges are far
9649        // apart because this test observes batching, not the I/O layer's distance-
9650        // based range coalescing.
9651        #[derive(Debug)]
9652        struct FakePageScheduler {
9653            meta_range: Range<u64>,
9654        }
9655        impl StructuralPageScheduler for FakePageScheduler {
9656            fn init_layout(&self) -> crate::Result<PageInitialization> {
9657                Ok(PageInitialization::new([Some(self.meta_range.clone())]))
9658            }
9659            fn init_from_buffers<'a>(
9660                &'a mut self,
9661                buffers: PageInitializationBuffers,
9662                _io: &Arc<dyn EncodingsIo>,
9663            ) -> BoxFuture<'a, crate::Result<Arc<dyn CachedPageData>>> {
9664                let [buffer] = buffers.try_into_array("fake").unwrap();
9665                let buffer = buffer.unwrap();
9666                assert_eq!(
9667                    &buffer[..8],
9668                    &self.meta_range.start.to_le_bytes(),
9669                    "each page gets the buffer for its own range"
9670                );
9671                std::future::ready(Ok(Arc::new(CacheMarker(42)) as Arc<dyn CachedPageData>)).boxed()
9672            }
9673            fn try_load(&mut self, data: &Arc<dyn CachedPageData>) -> crate::Result<()> {
9674                data.clone()
9675                    .as_arc_any()
9676                    .downcast::<CacheMarker>()
9677                    .map(|_| ())
9678                    .map_err(|_| {
9679                        Error::invalid_input_source(
9680                            "Fake page cache data has an unexpected type".into(),
9681                        )
9682                    })
9683            }
9684            fn schedule_ranges(
9685                &self,
9686                _ranges: &[Range<u64>],
9687                _io: &Arc<dyn EncodingsIo>,
9688            ) -> crate::Result<Vec<PageLoadTask>> {
9689                unreachable!("initialize never schedules data reads");
9690            }
9691        }
9692
9693        const NUM_PAGES: u64 = 64;
9694        const META_LEN: u64 = 64;
9695        const STRIDE: u64 = 1024 * 1024;
9696        let page_schedulers = (0..NUM_PAGES)
9697            .map(|i| PageInfoAndScheduler {
9698                page_index: i as usize,
9699                row_range: (i * 100)..((i + 1) * 100),
9700                scheduler: Box::new(FakePageScheduler {
9701                    meta_range: (i * STRIDE)..(i * STRIDE + META_LEN),
9702                }),
9703            })
9704            .collect();
9705        let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9706            page_schedulers,
9707            0,
9708            Arc::from("test"),
9709        );
9710
9711        let io = Arc::new(RecordingScheduler::default());
9712        // A no-op cache always misses, so every page is a miss.
9713        let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
9714        let context = SchedulerContext::new(io.clone(), cache.clone());
9715
9716        // `None` requests every page; against an always-miss cache that's an
9717        // N-page miss.
9718        let filter = FilterExpression::no_filter();
9719        scheduler.initialize(None, &filter, &context).await.unwrap();
9720
9721        {
9722            let requests = io.requests.lock().unwrap();
9723            assert_eq!(
9724                requests.len(),
9725                1,
9726                "an N-page miss must issue exactly one request, not one per page"
9727            );
9728            assert_eq!(
9729                requests[0].len() as u64,
9730                NUM_PAGES,
9731                "the single request must carry every missed page's metadata range"
9732            );
9733        }
9734
9735        io.requests.lock().unwrap().clear();
9736        let cache_misses_before = cache.stats().await.misses;
9737        let requested_ranges = Arc::<[Range<u64>]>::from([150..151, 6_350..6_351]);
9738        scheduler
9739            .initialize(Some(&requested_ranges), &filter, &context)
9740            .await
9741            .unwrap();
9742        assert_eq!(
9743            cache.stats().await.misses - cache_misses_before,
9744            2,
9745            "range-scoped initialization must only look up touched pages"
9746        );
9747
9748        {
9749            let requests = io.requests.lock().unwrap();
9750            assert_eq!(requests.len(), 1);
9751            assert_eq!(
9752                requests[0],
9753                vec![
9754                    STRIDE..(STRIDE + META_LEN),
9755                    63 * STRIDE..(63 * STRIDE + META_LEN),
9756                ],
9757                "range-scoped initialization must omit untouched pages"
9758            );
9759        }
9760
9761        io.requests.lock().unwrap().clear();
9762        let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024));
9763        let cache_key = PageDataCacheKey {
9764            column_index: 9,
9765            page_index: 0,
9766            view: PageCacheView::Native,
9767            target_data_type_tag: Arc::from("test"),
9768        };
9769        cache
9770            .insert_with_key(
9771                &cache_key,
9772                Arc::new(CachedPage {
9773                    data: Arc::new(super::NoCachedPageData),
9774                }),
9775            )
9776            .await;
9777        let mut recovering_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9778            vec![PageInfoAndScheduler {
9779                page_index: 0,
9780                row_range: 0..100,
9781                scheduler: Box::new(FakePageScheduler {
9782                    meta_range: 0..META_LEN,
9783                }),
9784            }],
9785            9,
9786            Arc::from("test"),
9787        );
9788        let recovering_context = SchedulerContext::new(io.clone(), cache.clone());
9789        recovering_scheduler
9790            .initialize(None, &filter, &recovering_context)
9791            .await
9792            .unwrap();
9793        assert_eq!(io.requests.lock().unwrap().as_slice(), &[vec![0..META_LEN]]);
9794
9795        io.requests.lock().unwrap().clear();
9796        let mut warm_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9797            vec![PageInfoAndScheduler {
9798                page_index: 0,
9799                row_range: 0..100,
9800                scheduler: Box::new(FakePageScheduler {
9801                    meta_range: 0..META_LEN,
9802                }),
9803            }],
9804            9,
9805            Arc::from("test"),
9806        );
9807        warm_scheduler
9808            .initialize(None, &filter, &recovering_context)
9809            .await
9810            .unwrap();
9811        assert!(
9812            io.requests.lock().unwrap().is_empty(),
9813            "the rebuilt cache entry must be reusable"
9814        );
9815
9816        let mut out_of_order_scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9817            vec![
9818                PageInfoAndScheduler {
9819                    page_index: 0,
9820                    row_range: 0..100,
9821                    scheduler: Box::new(FakePageScheduler {
9822                        meta_range: STRIDE..(STRIDE + META_LEN),
9823                    }),
9824                },
9825                PageInfoAndScheduler {
9826                    page_index: 1,
9827                    row_range: 100..200,
9828                    scheduler: Box::new(FakePageScheduler {
9829                        meta_range: 0..META_LEN,
9830                    }),
9831                },
9832            ],
9833            0,
9834            Arc::from("test"),
9835        );
9836        out_of_order_scheduler
9837            .initialize(None, &filter, &context)
9838            .await
9839            .unwrap();
9840        assert_eq!(
9841            io.requests.lock().unwrap().as_slice(),
9842            &[vec![0..META_LEN, STRIDE..(STRIDE + META_LEN)]],
9843            "metadata reads must be sorted without changing page buffer assignment"
9844        );
9845    }
9846
9847    #[tokio::test]
9848    async fn page_cache_key_distinguishes_page_view_and_target_type() {
9849        let cache = lance_core::cache::LanceCache::with_capacity(4 * 1024);
9850        let cases = [
9851            (
9852                PageDataCacheKey {
9853                    column_index: 7,
9854                    page_index: 0,
9855                    view: PageCacheView::Native,
9856                    target_data_type_tag: Arc::from("Boolean"),
9857                },
9858                10,
9859            ),
9860            (
9861                PageDataCacheKey {
9862                    column_index: 7,
9863                    page_index: 1,
9864                    view: PageCacheView::Native,
9865                    target_data_type_tag: Arc::from("Boolean"),
9866                },
9867                11,
9868            ),
9869            (
9870                PageDataCacheKey {
9871                    column_index: 7,
9872                    page_index: 0,
9873                    view: PageCacheView::Native,
9874                    target_data_type_tag: Arc::from("UInt8"),
9875                },
9876                12,
9877            ),
9878            (
9879                PageDataCacheKey {
9880                    column_index: 7,
9881                    page_index: 0,
9882                    view: PageCacheView::BlobDescriptor,
9883                    target_data_type_tag: Arc::from("Struct"),
9884                },
9885                20,
9886            ),
9887            (
9888                PageDataCacheKey {
9889                    column_index: 7,
9890                    page_index: 0,
9891                    view: PageCacheView::BlobPayload,
9892                    target_data_type_tag: Arc::from("LargeBinary"),
9893                },
9894                21,
9895            ),
9896        ];
9897
9898        for (key, marker) in &cases {
9899            cache
9900                .insert_with_key(
9901                    key,
9902                    Arc::new(CachedPage {
9903                        data: Arc::new(CacheMarker(*marker)),
9904                    }),
9905                )
9906                .await;
9907        }
9908
9909        assert_eq!(cache.size().await, cases.len());
9910        for (key, expected) in &cases {
9911            let cached = cache.get_with_key(key).await.unwrap();
9912            let marker = cached
9913                .data
9914                .clone()
9915                .as_arc_any()
9916                .downcast::<CacheMarker>()
9917                .unwrap();
9918            assert_eq!(marker.0, *expected);
9919        }
9920    }
9921
9922    #[tokio::test]
9923    async fn no_op_pages_do_not_create_cache_entries() {
9924        let mut scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
9925            vec![PageInfoAndScheduler {
9926                page_index: 0,
9927                row_range: 0..10,
9928                scheduler: Box::new(SimpleAllNullScheduler::default()),
9929            }],
9930            3,
9931            Arc::from("test"),
9932        );
9933        let io: Arc<dyn crate::EncodingsIo> =
9934            Arc::new(SimulatedScheduler::new(bytes::Bytes::new()));
9935        let cache = Arc::new(lance_core::cache::LanceCache::with_capacity(4 * 1024));
9936        let context = SchedulerContext::new(io, cache.clone());
9937
9938        scheduler
9939            .initialize(None, &FilterExpression::no_filter(), &context)
9940            .await
9941            .unwrap();
9942
9943        assert_eq!(cache.size().await, 0);
9944    }
9945
9946    #[test]
9947    fn page_initialization_preserves_optional_buffer_slots() {
9948        let initialization = PageInitialization::new([None, Some(20..22), None, Some(10..11)]);
9949        assert_eq!(
9950            initialization
9951                .required_ranges()
9952                .cloned()
9953                .collect::<Vec<_>>(),
9954            [20..22, 10..11]
9955        );
9956
9957        let buffers = initialization
9958            .with_buffers(vec![
9959                bytes::Bytes::from_static(b"ab"),
9960                bytes::Bytes::from_static(b"c"),
9961            ])
9962            .unwrap();
9963        let [first, second, third, fourth] = buffers.try_into_array("test").unwrap();
9964        assert!(first.is_none());
9965        assert_eq!(second.unwrap(), bytes::Bytes::from_static(b"ab"));
9966        assert!(third.is_none());
9967        assert_eq!(fourth.unwrap(), bytes::Bytes::from_static(b"c"));
9968
9969        let error = PageInitialization::new([Some(0..1)])
9970            .with_buffers(Vec::new())
9971            .unwrap_err();
9972        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
9973        assert!(error.to_string().contains("1 declared ranges"));
9974    }
9975
9976    #[tokio::test]
9977    async fn initialize_builds_only_requested_page_schedulers() {
9978        use crate::compression::DefaultDecompressionStrategy;
9979        use crate::decoder::{
9980            ColumnInfo, FilterExpression, PageEncoding, PageInfo, SchedulerContext,
9981            StructuralFieldScheduler,
9982        };
9983        use crate::format::pb;
9984        use lance_core::datatypes::Field;
9985
9986        const NUM_PAGES: u64 = 64;
9987        const ROWS_PER_PAGE: u64 = 100;
9988        // All-null constant pages need no initialization I/O, so this test
9989        // observes construction alone.
9990        let page_infos = (0..NUM_PAGES)
9991            .map(|page_index| PageInfo {
9992                num_rows: ROWS_PER_PAGE,
9993                priority: page_index * ROWS_PER_PAGE,
9994                encoding: PageEncoding::Structural(pb21::PageLayout {
9995                    layout: Some(pb21::page_layout::Layout::ConstantLayout(
9996                        pb21::ConstantLayout {
9997                            layers: vec![pb21::RepDefLayer::RepdefNullableItem as i32],
9998                            ..Default::default()
9999                        },
10000                    )),
10001                }),
10002                buffer_offsets_and_sizes: Arc::new([]),
10003            })
10004            .collect::<Vec<_>>();
10005        let column_info = Arc::new(ColumnInfo::new(
10006            3,
10007            Arc::from(page_infos),
10008            Vec::new(),
10009            pb::ColumnEncoding::default(),
10010        ));
10011        let field = Field::try_from(&ArrowField::new("x", DataType::Int32, true)).unwrap();
10012        #[allow(deprecated)]
10013        let eager = StructuralPrimitiveFieldScheduler::try_new(
10014            &column_info,
10015            &DefaultDecompressionStrategy {},
10016            false,
10017            &field,
10018        )
10019        .unwrap();
10020        assert_eq!(eager.page_schedulers.len(), NUM_PAGES as usize);
10021        assert!(eager.page_source.is_none());
10022
10023        let mut scheduler = StructuralPrimitiveFieldScheduler::try_new_lazy(
10024            column_info,
10025            Arc::new(DefaultDecompressionStrategy {}),
10026            false,
10027            &field,
10028        )
10029        .unwrap();
10030        assert!(
10031            scheduler.page_schedulers.is_empty(),
10032            "construction must not build page schedulers"
10033        );
10034        assert_eq!(scheduler.page_row_ends.len(), NUM_PAGES as usize);
10035
10036        let io: Arc<dyn EncodingsIo> = Arc::new(SimulatedScheduler::new(bytes::Bytes::new()));
10037        let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
10038        let mut context = SchedulerContext::new(io, cache);
10039        let filter = FilterExpression::no_filter();
10040
10041        let requested_ranges = [150..151, 6_350..6_351];
10042        scheduler
10043            .initialize(Some(&requested_ranges), &filter, &context)
10044            .await
10045            .unwrap();
10046        let built_pages = |scheduler: &StructuralPrimitiveFieldScheduler| {
10047            scheduler
10048                .page_schedulers
10049                .iter()
10050                .map(|page| page.page_index)
10051                .collect::<Vec<_>>()
10052        };
10053        assert_eq!(built_pages(&scheduler), [1, 63]);
10054        assert_eq!(scheduler.page_schedulers[1].row_range, 6_300..6_400);
10055        assert!(scheduler.page_source.is_some());
10056
10057        {
10058            let mut job = scheduler
10059                .schedule_ranges(&requested_ranges, &filter)
10060                .unwrap();
10061            // Both pages fit one read batch, so one call schedules them together.
10062            assert_eq!(job.schedule_next(&mut context).unwrap().len(), 2);
10063            assert!(job.schedule_next(&mut context).unwrap().is_empty());
10064
10065            // Rows outside the initialized ranges must fail cleanly, not panic
10066            // on an uninitialized page.
10067            let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap();
10068            let error = job.schedule_next(&mut context).unwrap_err();
10069            assert!(matches!(error, lance_core::Error::Internal { .. }));
10070            assert!(error.to_string().contains("not initialized"));
10071        }
10072
10073        // Widening the request builds only the new pages and keeps page order.
10074        scheduler
10075            .initialize(Some(&[0..250]), &filter, &context)
10076            .await
10077            .unwrap();
10078        assert_eq!(built_pages(&scheduler), [0, 1, 2, 63]);
10079
10080        // The eager path builds every page and releases the page source.
10081        scheduler.initialize(None, &filter, &context).await.unwrap();
10082        assert_eq!(scheduler.page_schedulers.len(), NUM_PAGES as usize);
10083        assert!(scheduler.page_source.is_none());
10084        let mut job = scheduler.schedule_ranges(&[0..1], &filter).unwrap();
10085        assert_eq!(job.schedule_next(&mut context).unwrap().len(), 1);
10086    }
10087
10088    /// Records every `submit_request` (ranges and priority) and answers each
10089    /// range with a zero buffer that starts with the range's offset, so a test
10090    /// can count requests and check that every reader got its own ranges back.
10091    #[derive(Debug, Default)]
10092    struct RecordingScheduler {
10093        requests: std::sync::Mutex<Vec<Vec<Range<u64>>>>,
10094        priorities: std::sync::Mutex<Vec<u64>>,
10095    }
10096
10097    impl crate::EncodingsIo for RecordingScheduler {
10098        fn submit_request(
10099            &self,
10100            ranges: Vec<Range<u64>>,
10101            priority: u64,
10102        ) -> futures::future::BoxFuture<'static, crate::Result<Vec<Bytes>>> {
10103            use futures::FutureExt;
10104
10105            let buffers = ranges
10106                .iter()
10107                .map(|range| {
10108                    let mut buffer = vec![0u8; (range.end - range.start) as usize];
10109                    let marker = range.start.to_le_bytes();
10110                    let marker_len = marker.len().min(buffer.len());
10111                    buffer[..marker_len].copy_from_slice(&marker[..marker_len]);
10112                    Bytes::from(buffer)
10113                })
10114                .collect();
10115            self.requests.lock().unwrap().push(ranges);
10116            self.priorities.lock().unwrap().push(priority);
10117            std::future::ready(Ok(buffers)).boxed()
10118        }
10119    }
10120
10121    #[tokio::test]
10122    async fn page_read_batch_submits_pending_reads_as_one_request() {
10123        let io = Arc::new(RecordingScheduler::default());
10124        let batch = RequestBatch::new(io.clone());
10125        let first = batch.submit_request(vec![100..108, 300..308], 7);
10126        let second = batch.submit_request(vec![200..208], 3);
10127        assert_eq!(batch.pending_bytes(), 24);
10128        assert!(
10129            io.requests.lock().unwrap().is_empty(),
10130            "nothing is read before the flush"
10131        );
10132
10133        batch.flush();
10134        {
10135            let requests = io.requests.lock().unwrap();
10136            assert_eq!(requests.len(), 1, "the batch is one request");
10137            assert_eq!(
10138                requests[0],
10139                vec![100..108, 200..208, 300..308],
10140                "the batch is sorted by file offset"
10141            );
10142        }
10143        assert_eq!(
10144            io.priorities.lock().unwrap()[0],
10145            3,
10146            "the batch takes the lowest priority"
10147        );
10148        let first = first.await.unwrap();
10149        assert_eq!(first.len(), 2);
10150        assert_eq!(first[0].as_ref(), &100_u64.to_le_bytes());
10151        assert_eq!(first[1].as_ref(), &300_u64.to_le_bytes());
10152        let second = second.await.unwrap();
10153        assert_eq!(second.len(), 1);
10154        assert_eq!(second[0].as_ref(), &200_u64.to_le_bytes());
10155
10156        // A read submitted after the flush is batched until it is awaited.
10157        let late = batch.submit_request(vec![400..408], 9).await.unwrap();
10158        assert_eq!(late[0].as_ref(), &400_u64.to_le_bytes());
10159        assert_eq!(io.requests.lock().unwrap().len(), 2);
10160
10161        // Queued reads that would exceed the budget are submitted first, so
10162        // no batch grows past it.
10163        let io = Arc::new(RecordingScheduler::default());
10164        let batch = RequestBatch::new(io.clone());
10165        let budget = crate::decoder::IO_REQUEST_BATCH_BYTES;
10166        let first = batch.submit_request(vec![0..budget], 1);
10167        assert!(io.requests.lock().unwrap().is_empty());
10168        let second = batch.submit_request(vec![budget..(budget + 8)], 2);
10169        assert_eq!(
10170            io.requests.lock().unwrap().as_slice(),
10171            &[vec![0..budget]],
10172            "the full batch is submitted before the read that would overflow it"
10173        );
10174        batch.flush();
10175        assert_eq!(io.requests.lock().unwrap().len(), 2);
10176        assert_eq!(first.await.unwrap()[0].len(), budget as usize);
10177        assert_eq!(second.await.unwrap()[0].as_ref(), &budget.to_le_bytes());
10178    }
10179
10180    #[tokio::test]
10181    async fn request_batch_flush_keeps_collecting_and_awaiting_flushes() {
10182        let io = Arc::new(RecordingScheduler::default());
10183        let batch = RequestBatch::new(io.clone());
10184        let first = batch.submit_request(vec![100..108], 1);
10185        let second = batch.submit_request(vec![300..308], 2);
10186        batch.flush();
10187        assert_eq!(
10188            io.requests.lock().unwrap().as_slice(),
10189            &[vec![100..108, 300..308]],
10190            "a flush submits the collected reads as one request"
10191        );
10192        // The batch keeps collecting after a flush.
10193        let third = batch.submit_request(vec![500..508], 3);
10194        let fourth = batch.submit_request(vec![700..708], 4);
10195        assert_eq!(batch.pending_bytes(), 16);
10196        assert_eq!(io.requests.lock().unwrap().len(), 1);
10197        // Awaiting a read submits whatever is queued, so a read queued outside
10198        // a scheduling step (indirect I/O) is never stuck behind a missing flush.
10199        assert_eq!(third.await.unwrap()[0].as_ref(), &500_u64.to_le_bytes());
10200        assert_eq!(
10201            io.requests.lock().unwrap().as_slice(),
10202            &[vec![100..108, 300..308], vec![500..508, 700..708]]
10203        );
10204        assert_eq!(fourth.await.unwrap()[0].as_ref(), &700_u64.to_le_bytes());
10205        assert_eq!(first.await.unwrap()[0].as_ref(), &100_u64.to_le_bytes());
10206        assert_eq!(second.await.unwrap()[0].as_ref(), &300_u64.to_le_bytes());
10207        assert_eq!(io.requests.lock().unwrap().len(), 2);
10208
10209        // Futures that submit when first polled and are polled in the same
10210        // pass (struct children initializing, list columns fetching items)
10211        // end up in one request: the flush waits for one yield.
10212        let (fifth, sixth) = futures::join!(
10213            async { batch.submit_request(vec![900..908], 5).await.unwrap() },
10214            async { batch.submit_request(vec![1100..1108], 6).await.unwrap() },
10215        );
10216        assert_eq!(fifth[0].as_ref(), &900_u64.to_le_bytes());
10217        assert_eq!(sixth[0].as_ref(), &1100_u64.to_le_bytes());
10218        assert_eq!(
10219            io.requests.lock().unwrap().last().unwrap(),
10220            &vec![900..908, 1100..1108],
10221            "reads submitted by sibling futures share one request"
10222        );
10223        assert_eq!(io.requests.lock().unwrap().len(), 3);
10224    }
10225
10226    #[tokio::test]
10227    async fn schedule_next_batches_page_reads_up_to_the_budget() {
10228        use super::{
10229            CachedPageData, IO_REQUEST_BATCH_BYTES, PageInfoAndScheduler, PageInitialization,
10230            PageInitializationBuffers, PageLoadTask, StructuralPageScheduler,
10231            StructuralPrimitiveFieldScheduler,
10232        };
10233        use crate::decoder::{FilterExpression, SchedulerContext, StructuralFieldScheduler};
10234
10235        /// A page that reads `data_ranges`, one shard (and one load task) per
10236        /// range. The bytes it gets back are covered by the batch test, so its
10237        /// load futures just fail and the test never has to build a decoder.
10238        #[derive(Debug)]
10239        struct FakeDataPage {
10240            data_ranges: Vec<Range<u64>>,
10241        }
10242
10243        impl StructuralPageScheduler for FakeDataPage {
10244            fn needs_initialization(&self) -> bool {
10245                false
10246            }
10247
10248            fn init_layout(&self) -> crate::Result<PageInitialization> {
10249                unreachable!("fake data pages need no initialization")
10250            }
10251
10252            fn init_from_buffers<'a>(
10253                &'a mut self,
10254                _buffers: PageInitializationBuffers,
10255                _io: &Arc<dyn EncodingsIo>,
10256            ) -> futures::future::BoxFuture<'a, crate::Result<Arc<dyn CachedPageData>>>
10257            {
10258                unreachable!("fake data pages need no initialization")
10259            }
10260
10261            fn try_load(&mut self, _data: &Arc<dyn CachedPageData>) -> crate::Result<()> {
10262                Ok(())
10263            }
10264
10265            fn schedule_ranges(
10266                &self,
10267                ranges: &[Range<u64>],
10268                io: &Arc<dyn EncodingsIo>,
10269            ) -> crate::Result<Vec<PageLoadTask>> {
10270                use futures::FutureExt;
10271
10272                let num_rows = ranges.iter().map(|r| r.end - r.start).sum();
10273                Ok(self
10274                    .data_ranges
10275                    .iter()
10276                    .enumerate()
10277                    .map(|(shard, range)| {
10278                        drop(io.submit_request(vec![range.clone()], range.start));
10279                        PageLoadTask {
10280                            decoder_fut: std::future::ready(Err(lance_core::Error::internal(
10281                                "fake page",
10282                            )))
10283                            .boxed(),
10284                            // The page's rows are reported once, on the first shard.
10285                            num_rows: if shard == 0 { num_rows } else { 0 },
10286                        }
10287                    })
10288                    .collect())
10289            }
10290        }
10291
10292        fn scheduler_with_pages(
10293            page_bytes: u64,
10294            num_pages: u64,
10295        ) -> StructuralPrimitiveFieldScheduler {
10296            // Leave a gap between pages so the test can tell batching from
10297            // the I/O layer's coalescing of adjacent ranges.
10298            let stride = page_bytes * 2;
10299            let pages = (0..num_pages)
10300                .map(|page| PageInfoAndScheduler {
10301                    page_index: page as usize,
10302                    row_range: (page * 10)..((page + 1) * 10),
10303                    scheduler: Box::new(FakeDataPage {
10304                        data_ranges: vec![(page * stride)..(page * stride + page_bytes)],
10305                    }),
10306                })
10307                .collect();
10308            StructuralPrimitiveFieldScheduler::from_page_schedulers(pages, 0, Arc::from("test"))
10309        }
10310
10311        let filter = FilterExpression::no_filter();
10312        let cache = Arc::new(lance_core::cache::LanceCache::no_cache());
10313
10314        // Small pages: every page of the request is read in one batch.
10315        let io = Arc::new(RecordingScheduler::default());
10316        let scheduler = scheduler_with_pages(64, 6);
10317        let mut context = SchedulerContext::new(io.clone(), cache.clone());
10318        let mut job = scheduler.schedule_ranges(&[5..45], &filter).unwrap();
10319        let scan_lines = job.schedule_next(&mut context).unwrap();
10320        context.flush_io();
10321        assert_eq!(scan_lines.len(), 5, "one scan line per touched page");
10322        assert_eq!(
10323            scan_lines
10324                .iter()
10325                .map(|line| line.rows_scheduled)
10326                .sum::<u64>(),
10327            40
10328        );
10329        assert!(job.schedule_next(&mut context).unwrap().is_empty());
10330        {
10331            let requests = io.requests.lock().unwrap();
10332            assert_eq!(requests.len(), 1, "five pages cost one request");
10333            assert_eq!(
10334                requests[0],
10335                vec![0..64, 128..192, 256..320, 384..448, 512..576],
10336                "the untouched last page is not read"
10337            );
10338        }
10339        assert_eq!(
10340            io.priorities.lock().unwrap()[0],
10341            0,
10342            "the batch carries the first page's priority"
10343        );
10344
10345        // Pages as large as the batch budget are read one request at a time.
10346        let io = Arc::new(RecordingScheduler::default());
10347        let scheduler = scheduler_with_pages(IO_REQUEST_BATCH_BYTES, 2);
10348        let mut context = SchedulerContext::new(io.clone(), cache);
10349        let mut job = scheduler.schedule_ranges(&[0..20], &filter).unwrap();
10350        let first = job.schedule_next(&mut context).unwrap();
10351        context.flush_io();
10352        assert_eq!(first.len(), 1);
10353        assert_eq!(io.requests.lock().unwrap().len(), 1);
10354        let second = job.schedule_next(&mut context).unwrap();
10355        context.flush_io();
10356        assert_eq!(second.len(), 1);
10357        assert_eq!(io.requests.lock().unwrap().len(), 2);
10358        assert!(job.schedule_next(&mut context).unwrap().is_empty());
10359
10360        // A page that shards its own reads (like a blob page) keeps every
10361        // shard in a bounded read instead of one request for the whole page.
10362        let io = Arc::new(RecordingScheduler::default());
10363        let shard = IO_REQUEST_BATCH_BYTES;
10364        let scheduler = StructuralPrimitiveFieldScheduler::from_page_schedulers(
10365            vec![PageInfoAndScheduler {
10366                page_index: 0,
10367                row_range: 0..10,
10368                scheduler: Box::new(FakeDataPage {
10369                    data_ranges: vec![0..shard, (2 * shard)..(3 * shard)],
10370                }),
10371            }],
10372            0,
10373            Arc::from("test"),
10374        );
10375        let mut context = SchedulerContext::new(
10376            io.clone(),
10377            Arc::new(lance_core::cache::LanceCache::no_cache()),
10378        );
10379        let mut job = scheduler.schedule_ranges(&[0..10], &filter).unwrap();
10380        let scan_lines = job.schedule_next(&mut context).unwrap();
10381        context.flush_io();
10382        assert_eq!(scan_lines.len(), 2, "one scan line per shard");
10383        assert_eq!(
10384            io.requests.lock().unwrap().as_slice(),
10385            &[vec![0..shard], vec![(2 * shard)..(3 * shard)]],
10386            "each shard stays its own bounded request"
10387        );
10388        assert!(job.schedule_next(&mut context).unwrap().is_empty());
10389    }
10390
10391    #[test]
10392    fn page_range_mapping_validates_and_splits_ranges() {
10393        // Exclusive end rows of three ten-row pages.
10394        let pages = [10_u64, 20, 30];
10395
10396        let mappings = map_ranges_to_pages(&pages, &[8..22]).unwrap();
10397        assert_eq!(mappings.len(), 3);
10398        assert_eq!(mappings[0].page_idx, 0);
10399        assert_eq!(mappings[0].ranges_in_page, [8..10]);
10400        assert_eq!(mappings[1].page_idx, 1);
10401        assert_eq!(mappings[1].ranges_in_page, [0..10]);
10402        assert_eq!(mappings[2].page_idx, 2);
10403        assert_eq!(mappings[2].ranges_in_page, [0..2]);
10404
10405        let sparse_mappings = map_ranges_to_pages(&pages, &[1..2, 3..4, 28..29]).unwrap();
10406        assert_eq!(sparse_mappings.len(), 2);
10407        assert_eq!(sparse_mappings[0].page_idx, 0);
10408        assert_eq!(sparse_mappings[0].ranges_in_page, [1..2, 3..4]);
10409        assert_eq!(sparse_mappings[1].page_idx, 2);
10410        assert_eq!(sparse_mappings[1].ranges_in_page, [8..9]);
10411        assert_eq!(
10412            pages_overlapping_ranges(&pages, &[1..2, 3..4, 28..29]).unwrap(),
10413            [0, 2]
10414        );
10415
10416        for (ranges, expected_message) in [
10417            (vec![Range { start: 2, end: 1 }], "is reversed"),
10418            (vec![15..16, 5..6], "not strictly ordered"),
10419            (vec![5..15, 10..20], "overlap"),
10420            (vec![5..5, 5..6], "not strictly ordered"),
10421        ] {
10422            let error = map_ranges_to_pages(&pages, &ranges).unwrap_err();
10423            assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
10424            assert!(error.to_string().contains(expected_message));
10425        }
10426    }
10427
10428    #[tokio::test]
10429    async fn test_fullzip_read_source_slices_prefetched_page() {
10430        let page_start = 200_u64;
10431        let page_data = LanceBuffer::copy_slice(&[0, 1, 2, 3, 4, 5, 6, 7]);
10432        let source = FullZipReadSource::PrefetchedPage {
10433            base_offset: page_start,
10434            data: page_data,
10435        };
10436        let ranges = vec![
10437            page_start..(page_start + 3),
10438            (page_start + 4)..(page_start + 8),
10439        ];
10440        let mut data = source.fetch(&ranges, 0).await.unwrap();
10441        assert_eq!(data.pop_front().unwrap().as_ref(), &[0, 1, 2]);
10442        assert_eq!(data.pop_front().unwrap().as_ref(), &[4, 5, 6, 7]);
10443    }
10444
10445    #[tokio::test]
10446    async fn test_fullzip_initialize_caches_rep_index_when_enabled() {
10447        use futures::{FutureExt, future::BoxFuture};
10448        use std::ops::Range;
10449        use std::sync::Mutex;
10450
10451        #[derive(Debug, Clone)]
10452        struct RecordingScheduler {
10453            data: bytes::Bytes,
10454            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
10455        }
10456
10457        impl RecordingScheduler {
10458            fn new(data: bytes::Bytes) -> Self {
10459                Self {
10460                    data,
10461                    requests: Arc::new(Mutex::new(Vec::new())),
10462                }
10463            }
10464
10465            fn requests(&self) -> Vec<Vec<Range<u64>>> {
10466                self.requests.lock().unwrap().clone()
10467            }
10468        }
10469
10470        impl crate::EncodingsIo for RecordingScheduler {
10471            fn submit_request(
10472                &self,
10473                ranges: Vec<Range<u64>>,
10474                _priority: u64,
10475            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
10476                self.requests.lock().unwrap().push(ranges.clone());
10477                let data = ranges
10478                    .into_iter()
10479                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
10480                    .collect::<Vec<_>>();
10481                std::future::ready(Ok(data)).boxed()
10482            }
10483        }
10484
10485        #[derive(Debug)]
10486        struct TestFixedDecompressor;
10487
10488        impl FixedPerValueDecompressor for TestFixedDecompressor {
10489            fn decompress(
10490                &self,
10491                _data: FixedWidthDataBlock,
10492                _num_rows: u64,
10493            ) -> crate::Result<DataBlock> {
10494                unimplemented!("Test decompressor")
10495            }
10496
10497            fn bits_per_value(&self) -> u64 {
10498                32
10499            }
10500        }
10501
10502        let rows_in_page = 100_u64;
10503        let bytes_per_value = 4_u64;
10504        let rep_start = 1000_u64;
10505        let rep_size = ((rows_in_page + 1) * bytes_per_value) as usize;
10506        let mut data = vec![0_u8; 16 * 1024];
10507        data[rep_start as usize..rep_start as usize + rep_size].fill(7);
10508        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(data)));
10509
10510        let mut scheduler = FullZipScheduler {
10511            data_buf_position: 0,
10512            data_buf_size: 4096,
10513            rep_index: Some(FullZipRepIndexDetails {
10514                buf_position: rep_start,
10515                bytes_per_value,
10516            }),
10517            priority: 0,
10518            rows_in_page,
10519            bits_per_offset: 32,
10520            details: Arc::new(FullZipDecodeDetails {
10521                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
10522                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
10523                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
10524                max_rep: 0,
10525                max_visible_def: 0,
10526            }),
10527            cached_state: None,
10528            enable_cache: true,
10529        };
10530
10531        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
10532        let initialization = scheduler.init_layout().unwrap();
10533        let ranges = initialization
10534            .required_ranges()
10535            .cloned()
10536            .collect::<Vec<_>>();
10537        assert_eq!(
10538            ranges,
10539            vec![rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)]
10540        );
10541        // The field scheduler fetches the declared ranges (coalesced across all
10542        // pages) and hands the bytes back to the page to finish initialization.
10543        let buffers = io_dyn.submit_request(ranges, 0).await.unwrap();
10544        let buffers = initialization.with_buffers(buffers).unwrap();
10545        let cached_data = scheduler.init_from_buffers(buffers, &io_dyn).await.unwrap();
10546        assert!(
10547            cached_data
10548                .clone()
10549                .as_arc_any()
10550                .downcast_ref::<FullZipCacheableState>()
10551                .is_some()
10552        );
10553        assert!(scheduler.cached_state.is_none());
10554        scheduler.try_load(&cached_data).unwrap();
10555        assert!(scheduler.cached_state.is_some());
10556        assert_eq!(
10557            io.requests(),
10558            vec![vec![
10559                rep_start..(rep_start + (rows_in_page + 1) * bytes_per_value)
10560            ]]
10561        );
10562    }
10563
10564    #[tokio::test]
10565    async fn test_fullzip_full_page_bypasses_rep_index_io() {
10566        use futures::{FutureExt, future::BoxFuture};
10567        use std::ops::Range;
10568        use std::sync::Mutex;
10569
10570        #[derive(Debug, Clone)]
10571        struct RecordingScheduler {
10572            data: bytes::Bytes,
10573            requests: Arc<Mutex<Vec<Vec<Range<u64>>>>>,
10574        }
10575
10576        impl RecordingScheduler {
10577            fn new(data: bytes::Bytes) -> Self {
10578                Self {
10579                    data,
10580                    requests: Arc::new(Mutex::new(Vec::new())),
10581                }
10582            }
10583
10584            fn requests(&self) -> Vec<Vec<Range<u64>>> {
10585                self.requests.lock().unwrap().clone()
10586            }
10587        }
10588
10589        impl crate::EncodingsIo for RecordingScheduler {
10590            fn submit_request(
10591                &self,
10592                ranges: Vec<Range<u64>>,
10593                _priority: u64,
10594            ) -> BoxFuture<'static, crate::Result<Vec<bytes::Bytes>>> {
10595                self.requests.lock().unwrap().push(ranges.clone());
10596                let data = ranges
10597                    .into_iter()
10598                    .map(|range| self.data.slice(range.start as usize..range.end as usize))
10599                    .collect::<Vec<_>>();
10600                std::future::ready(Ok(data)).boxed()
10601            }
10602        }
10603
10604        #[derive(Debug)]
10605        struct TestFixedDecompressor;
10606
10607        impl FixedPerValueDecompressor for TestFixedDecompressor {
10608            fn decompress(
10609                &self,
10610                _data: FixedWidthDataBlock,
10611                _num_rows: u64,
10612            ) -> crate::Result<DataBlock> {
10613                unimplemented!("Test decompressor")
10614            }
10615
10616            fn bits_per_value(&self) -> u64 {
10617                32
10618            }
10619        }
10620
10621        let rows_in_page = 100_u64;
10622        let data_start = 256_u64;
10623        let data_size = 500_u64;
10624        let rep_start = 4096_u64;
10625        let bytes_per_value = 4_u64;
10626
10627        let mut bytes = vec![0_u8; 16 * 1024];
10628        for i in 0..=rows_in_page {
10629            let offset = (i * 5) as u32;
10630            let pos = rep_start as usize + (i * bytes_per_value) as usize;
10631            bytes[pos..pos + 4].copy_from_slice(&offset.to_le_bytes());
10632        }
10633        let io = Arc::new(RecordingScheduler::new(bytes::Bytes::from(bytes)));
10634
10635        let scheduler = FullZipScheduler {
10636            data_buf_position: data_start,
10637            data_buf_size: data_size,
10638            rep_index: Some(FullZipRepIndexDetails {
10639                buf_position: rep_start,
10640                bytes_per_value,
10641            }),
10642            priority: 0,
10643            rows_in_page,
10644            bits_per_offset: 32,
10645            details: Arc::new(FullZipDecodeDetails {
10646                value_decompressor: PerValueDecompressor::Fixed(Arc::new(TestFixedDecompressor)),
10647                def_meaning: Arc::new([crate::repdef::DefinitionInterpretation::NullableItem]),
10648                ctrl_word_parser: crate::repdef::ControlWordParser::new(0, 1),
10649                max_rep: 0,
10650                max_visible_def: 0,
10651            }),
10652            cached_state: None,
10653            enable_cache: false,
10654        };
10655
10656        let io_dyn: Arc<dyn crate::EncodingsIo> = io.clone();
10657        let tasks = scheduler
10658            .schedule_ranges_rep(
10659                &[0..rows_in_page],
10660                &io_dyn,
10661                FullZipRepIndexDetails {
10662                    buf_position: rep_start,
10663                    bytes_per_value,
10664                },
10665            )
10666            .unwrap();
10667
10668        let requests = io.requests();
10669        assert_eq!(requests.len(), 1);
10670        assert_eq!(requests[0], vec![data_start..(data_start + data_size)]);
10671
10672        let _ = tasks.into_iter().next().unwrap().decoder_fut.await.unwrap();
10673        let requests_after_await = io.requests();
10674        assert_eq!(
10675            requests_after_await.len(),
10676            1,
10677            "full page path should not issue rep-index I/O"
10678        );
10679    }
10680
10681    /// This test is used to reproduce fuzz test https://github.com/lancedb/lance/issues/4492
10682    #[tokio::test]
10683    async fn test_fuzz_issue_4492_empty_rep_values() {
10684        use lance_datagen::{RowCount, Seed, array, gen_batch};
10685
10686        let seed = 1823859942947654717u64;
10687        let num_rows = 2741usize;
10688
10689        // Generate the exact same data that caused the failure
10690        let batch_gen = gen_batch().with_seed(Seed::from(seed));
10691        let base_generator = array::rand_type(&DataType::FixedSizeBinary(32));
10692        let list_generator = array::rand_list_any(base_generator, false);
10693
10694        let batch = batch_gen
10695            .anon_col(list_generator)
10696            .into_batch_rows(RowCount::from(num_rows as u64))
10697            .unwrap();
10698
10699        let list_array = batch.column(0).clone();
10700
10701        // Force miniblock encoding
10702        let mut metadata = HashMap::new();
10703        metadata.insert(
10704            STRUCTURAL_ENCODING_META_KEY.to_string(),
10705            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
10706        );
10707
10708        let test_cases = TestCases::default()
10709            .with_structural_encodings()
10710            .with_batch_size(100)
10711            .with_range(0..num_rows.min(500) as u64)
10712            .with_indices(vec![0, num_rows as u64 / 2, (num_rows - 1) as u64]);
10713
10714        check_round_trip_encoding_of_data(vec![list_array], &test_cases, metadata).await
10715    }
10716
10717    async fn test_minichunk_size_helper(
10718        string_data: Vec<Option<String>>,
10719        minichunk_size: u64,
10720        encodings: &[TestEncoding],
10721    ) {
10722        use crate::constants::MINICHUNK_SIZE_META_KEY;
10723        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
10724        use arrow_array::{ArrayRef, StringArray};
10725        use std::sync::Arc;
10726
10727        let string_array: ArrayRef = Arc::new(StringArray::from(string_data));
10728
10729        let mut metadata = HashMap::new();
10730        metadata.insert(
10731            MINICHUNK_SIZE_META_KEY.to_string(),
10732            minichunk_size.to_string(),
10733        );
10734        metadata.insert(
10735            STRUCTURAL_ENCODING_META_KEY.to_string(),
10736            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
10737        );
10738
10739        let test_cases = TestCases::default()
10740            .with_encodings(encodings.iter().copied())
10741            .with_batch_size(1000);
10742
10743        check_round_trip_encoding_of_data(vec![string_array], &test_cases, metadata).await;
10744    }
10745
10746    #[tokio::test]
10747    async fn test_minichunk_size_roundtrip() {
10748        // Test that minichunk size can be configured and works correctly in round-trip encoding
10749        let mut string_data = Vec::new();
10750        for i in 0..100 {
10751            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
10752        }
10753        // configure minichunk size to 64 bytes (smaller than the default 4kb) for Lance 2.1
10754        test_minichunk_size_helper(
10755            string_data,
10756            64,
10757            &[
10758                TestEncoding::StructuralU16,
10759                TestEncoding::StructuralU32,
10760                TestEncoding::StructuralSparse,
10761            ],
10762        )
10763        .await;
10764    }
10765
10766    #[tokio::test]
10767    async fn test_minichunk_size_128kb_v2_2() {
10768        // Test that minichunk size can be configured to 128KB and works correctly with Lance 2.2
10769        let mut string_data = Vec::new();
10770        // create a 500kb string array
10771        for i in 0..10000 {
10772            string_data.push(Some(format!("test_string_{}", i).repeat(50)));
10773        }
10774        test_minichunk_size_helper(
10775            string_data,
10776            128 * 1024,
10777            &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
10778        )
10779        .await;
10780    }
10781
10782    #[tokio::test]
10783    async fn test_binary_large_minichunk_size_over_max_miniblock_values() {
10784        let mut string_data = Vec::new();
10785        // 128kb/chunk / 6 bytes (t_9999) = 21845 items per chunk
10786        for i in 0..10000 {
10787            string_data.push(Some(format!("t_{}", i)));
10788        }
10789        test_minichunk_size_helper(
10790            string_data,
10791            128 * 1024,
10792            &[TestEncoding::StructuralU32, TestEncoding::StructuralSparse],
10793        )
10794        .await;
10795    }
10796
10797    #[tokio::test]
10798    async fn test_large_dictionary_general_compression() {
10799        use arrow_array::{ArrayRef, StringArray};
10800        use std::collections::HashMap;
10801        use std::sync::Arc;
10802
10803        // Create large string dictionary data (>32KiB) with low cardinality
10804        // Use 100 unique strings, each 500 bytes long = 50KB dictionary
10805        let unique_values: Vec<String> = (0..100)
10806            .map(|i| format!("value_{:04}_{}", i, "x".repeat(500)))
10807            .collect();
10808
10809        // Repeat these strings many times to create a large array
10810        let repeated_strings: Vec<_> = unique_values
10811            .iter()
10812            .cycle()
10813            .take(10_000)
10814            .map(|s| Some(s.as_str()))
10815            .collect();
10816
10817        let string_array = Arc::new(StringArray::from(repeated_strings)) as ArrayRef;
10818
10819        // Configure test to use V2_2 and verify encoding
10820        let test_cases = TestCases::default()
10821            .with_u32_structural_encodings()
10822            .with_verify_encoding(Arc::new(|cols: &[crate::encoder::EncodedColumn], _| {
10823                assert_eq!(cols.len(), 1);
10824                let col = &cols[0];
10825
10826                // Navigate to the dictionary encoding in the page layout
10827                if let Some(PageEncoding::Structural(page_layout)) =
10828                    &col.final_pages.first().map(|p| &p.description)
10829                    && let Some(pb21::page_layout::Layout::MiniBlockLayout(mini_block)) =
10830                        &page_layout.layout
10831                    && let Some(dictionary_encoding) = &mini_block.dictionary
10832                {
10833                    match dictionary_encoding.compression.as_ref() {
10834                        Some(Compression::General(general)) => {
10835                            // Verify it's using LZ4 or Zstd
10836                            let compression = general.compression.as_ref().unwrap();
10837                            assert!(
10838                                compression.scheme()
10839                                    == pb21::CompressionScheme::CompressionAlgorithmLz4
10840                                    || compression.scheme()
10841                                        == pb21::CompressionScheme::CompressionAlgorithmZstd,
10842                                "Expected LZ4 or Zstd compression for large dictionary"
10843                            );
10844                        }
10845                        _ => panic!("Expected General compression for large dictionary"),
10846                    }
10847                }
10848            }));
10849
10850        check_round_trip_encoding_of_data(vec![string_array], &test_cases, HashMap::new()).await;
10851    }
10852
10853    fn dictionary_encoding_from_page(
10854        page: &crate::encoder::EncodedPage,
10855    ) -> &crate::format::pb21::CompressiveEncoding {
10856        let PageEncoding::Structural(layout) = &page.description else {
10857            panic!("Expected structural page encoding");
10858        };
10859        let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap()
10860        else {
10861            panic!("Expected mini-block layout");
10862        };
10863        layout
10864            .dictionary
10865            .as_ref()
10866            .unwrap_or_else(|| panic!("Expected dictionary encoding"))
10867    }
10868
10869    async fn encode_variable_dict_page(
10870        metadata: HashMap<String, String>,
10871    ) -> crate::encoder::EncodedPage {
10872        use arrow_array::types::Int32Type;
10873        use arrow_array::{ArrayRef, DictionaryArray, Int32Array, StringArray};
10874
10875        let values = Arc::new(StringArray::from(
10876            (0..128)
10877                .map(|i| format!("value_{i:04}_{}", "x".repeat(256)))
10878                .collect::<Vec<_>>(),
10879        )) as ArrayRef;
10880        let keys = Int32Array::from_iter_values((0..20_000).map(|i| i % 128));
10881        let dict_array =
10882            Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
10883
10884        let field = arrow_schema::Field::new(
10885            "dict_col",
10886            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
10887            false,
10888        )
10889        .with_metadata(metadata);
10890
10891        encode_first_page(field, dict_array, TestEncoding::StructuralU32).await
10892    }
10893
10894    async fn encode_auto_fixed_dict_page(
10895        metadata: HashMap<String, String>,
10896    ) -> crate::encoder::EncodedPage {
10897        use arrow_array::{ArrayRef, Decimal128Array};
10898
10899        // 128-bit fixed-width values with low cardinality to trigger dictionary encoding.
10900        let values = (0..20_000)
10901            .map(|i| match i % 3 {
10902                0 => 10_i128,
10903                1 => 20_i128,
10904                _ => 30_i128,
10905            })
10906            .collect::<Vec<_>>();
10907        let decimal = Decimal128Array::from_iter_values(values)
10908            .with_precision_and_scale(38, 0)
10909            .unwrap();
10910        let decimal = Arc::new(decimal) as ArrayRef;
10911
10912        let mut field_metadata = metadata;
10913        // Strongly encourage dictionary encoding for this synthetic test data.
10914        field_metadata.insert(
10915            "lance-encoding:dict-size-ratio".to_string(),
10916            "0.99".to_string(),
10917        );
10918        let field = arrow_schema::Field::new("fixed_col", DataType::Decimal128(38, 0), false)
10919            .with_metadata(field_metadata);
10920
10921        encode_first_page(field, decimal, TestEncoding::StructuralU32).await
10922    }
10923
10924    #[tokio::test]
10925    async fn test_dict_values_general_compression_default_lz4_for_variable_dict_values() {
10926        let page = encode_variable_dict_page(HashMap::new()).await;
10927        let dictionary_encoding = dictionary_encoding_from_page(&page);
10928        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10929            panic!("Expected General compression for dictionary values");
10930        };
10931        let compression = general.compression.as_ref().unwrap();
10932        assert_eq!(
10933            compression.scheme(),
10934            pb21::CompressionScheme::CompressionAlgorithmLz4
10935        );
10936    }
10937
10938    #[tokio::test]
10939    async fn test_dict_values_general_compression_default_lz4_for_fixed_dict_values() {
10940        let page = encode_auto_fixed_dict_page(HashMap::new()).await;
10941        let dictionary_encoding = dictionary_encoding_from_page(&page);
10942        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10943            panic!("Expected General compression for dictionary values");
10944        };
10945        let compression = general.compression.as_ref().unwrap();
10946        assert_eq!(
10947            compression.scheme(),
10948            pb21::CompressionScheme::CompressionAlgorithmLz4
10949        );
10950    }
10951
10952    #[tokio::test]
10953    async fn test_dict_values_general_compression_zstd() {
10954        let mut metadata = HashMap::new();
10955        metadata.insert(
10956            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
10957            "zstd".to_string(),
10958        );
10959        let page = encode_variable_dict_page(metadata).await;
10960        let dictionary_encoding = dictionary_encoding_from_page(&page);
10961        let Some(Compression::General(general)) = dictionary_encoding.compression.as_ref() else {
10962            panic!("Expected General compression for dictionary values");
10963        };
10964        let compression = general.compression.as_ref().unwrap();
10965        assert_eq!(
10966            compression.scheme(),
10967            pb21::CompressionScheme::CompressionAlgorithmZstd
10968        );
10969    }
10970
10971    #[tokio::test]
10972    async fn test_dict_values_general_compression_none() {
10973        let mut metadata = HashMap::new();
10974        metadata.insert(
10975            DICT_VALUES_COMPRESSION_META_KEY.to_string(),
10976            "none".to_string(),
10977        );
10978        let page = encode_variable_dict_page(metadata).await;
10979        let dictionary_encoding = dictionary_encoding_from_page(&page);
10980        assert!(
10981            !matches!(
10982                dictionary_encoding.compression.as_ref(),
10983                Some(Compression::General(_))
10984            ),
10985            "Expected dictionary values to avoid General compression"
10986        );
10987    }
10988
10989    #[test]
10990    fn test_resolve_dict_values_compression_metadata_defaults_to_lz4() {
10991        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
10992            &HashMap::new(),
10993            None,
10994            None,
10995        );
10996        assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),);
10997        assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY));
10998    }
10999
11000    #[test]
11001    fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() {
11002        let field_metadata = HashMap::from([
11003            (
11004                DICT_VALUES_COMPRESSION_META_KEY.to_string(),
11005                "none".to_string(),
11006            ),
11007            (
11008                DICT_VALUES_COMPRESSION_LEVEL_META_KEY.to_string(),
11009                "7".to_string(),
11010            ),
11011        ]);
11012        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
11013            &field_metadata,
11014            Some("zstd".to_string()),
11015            Some("3".to_string()),
11016        );
11017        assert_eq!(
11018            metadata.get(COMPRESSION_META_KEY),
11019            Some(&"none".to_string()),
11020        );
11021        assert_eq!(
11022            metadata.get(COMPRESSION_LEVEL_META_KEY),
11023            Some(&"7".to_string()),
11024        );
11025    }
11026
11027    #[test]
11028    fn test_resolve_dict_values_compression_metadata_env_fallback() {
11029        let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata(
11030            &HashMap::new(),
11031            Some("zstd".to_string()),
11032            Some("9".to_string()),
11033        );
11034        assert_eq!(
11035            metadata.get(COMPRESSION_META_KEY),
11036            Some(&"zstd".to_string()),
11037        );
11038        assert_eq!(
11039            metadata.get(COMPRESSION_LEVEL_META_KEY),
11040            Some(&"9".to_string()),
11041        );
11042    }
11043
11044    #[tokio::test]
11045    async fn test_dictionary_encode_int64() {
11046        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
11047        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
11048        use arrow_array::{ArrayRef, Int64Array};
11049        use std::collections::HashMap;
11050        use std::sync::Arc;
11051
11052        // Low cardinality with poor RLE opportunity.
11053        let values = (0..1000)
11054            .map(|i| match i % 3 {
11055                0 => 10i64,
11056                1 => 20i64,
11057                _ => 30i64,
11058            })
11059            .collect::<Vec<_>>();
11060        let array = Arc::new(Int64Array::from(values)) as ArrayRef;
11061
11062        let mut metadata = HashMap::new();
11063        metadata.insert(
11064            STRUCTURAL_ENCODING_META_KEY.to_string(),
11065            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
11066        );
11067        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
11068
11069        let test_cases = TestCases::default()
11070            .with_u32_structural_encodings()
11071            .with_batch_size(1000)
11072            .with_range(0..1000)
11073            .with_indices(vec![0, 1, 10, 999])
11074            .with_expected_encoding("dictionary");
11075
11076        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
11077    }
11078
11079    #[tokio::test]
11080    async fn test_dictionary_encode_float64() {
11081        use crate::constants::{DICT_SIZE_RATIO_META_KEY, STRUCTURAL_ENCODING_META_KEY};
11082        use crate::testing::{TestCases, check_round_trip_encoding_of_data};
11083        use arrow_array::{ArrayRef, Float64Array};
11084        use std::collections::HashMap;
11085        use std::sync::Arc;
11086
11087        // Low cardinality with poor RLE opportunity.
11088        let values = (0..1000)
11089            .map(|i| match i % 3 {
11090                0 => 0.1f64,
11091                1 => 0.2f64,
11092                _ => 0.3f64,
11093            })
11094            .collect::<Vec<_>>();
11095        let array = Arc::new(Float64Array::from(values)) as ArrayRef;
11096
11097        let mut metadata = HashMap::new();
11098        metadata.insert(
11099            STRUCTURAL_ENCODING_META_KEY.to_string(),
11100            STRUCTURAL_ENCODING_MINIBLOCK.to_string(),
11101        );
11102        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.99".to_string());
11103
11104        let test_cases = TestCases::default()
11105            .with_u32_structural_encodings()
11106            .with_batch_size(1000)
11107            .with_range(0..1000)
11108            .with_indices(vec![0, 1, 10, 999])
11109            .with_expected_encoding("dictionary");
11110
11111        check_round_trip_encoding_of_data(vec![array], &test_cases, metadata).await;
11112    }
11113
11114    #[test]
11115    fn test_miniblock_dictionary_out_of_line_bitpacking_decode() {
11116        let rows = 10_000;
11117        let unique_values = 2_000;
11118
11119        let dictionary_encoding =
11120            ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(11, None));
11121        let layout = pb21::MiniBlockLayout {
11122            rep_compression: None,
11123            def_compression: None,
11124            value_compression: Some(ProtobufUtils21::flat(64, None)),
11125            dictionary: Some(dictionary_encoding),
11126            num_dictionary_items: unique_values,
11127            layers: vec![pb21::RepDefLayer::RepdefAllValidItem as i32],
11128            num_buffers: 1,
11129            repetition_index_depth: 0,
11130            num_items: rows,
11131            has_large_chunk: false,
11132        };
11133
11134        let buffer_offsets_and_sizes = vec![(0, 0), (0, 0), (0, 0)];
11135        let scheduler = super::MiniBlockScheduler::try_new(
11136            &buffer_offsets_and_sizes,
11137            /*priority=*/ 0,
11138            /*items_in_page=*/ rows,
11139            &layout,
11140            &DefaultDecompressionStrategy::default(),
11141        )
11142        .unwrap();
11143
11144        let dictionary = scheduler.dictionary.unwrap();
11145        assert_eq!(dictionary.num_dictionary_items, unique_values);
11146        assert_eq!(
11147            dictionary.dictionary_data_alignment,
11148            crate::encoder::MIN_PAGE_BUFFER_ALIGNMENT
11149        );
11150    }
11151
11152    // Dictionary encoding decision tests
11153    fn create_test_fixed_data_block(
11154        num_values: u64,
11155        cardinality: u64,
11156        bits_per_value: u64,
11157    ) -> DataBlock {
11158        assert!(cardinality > 0);
11159        assert!(cardinality <= num_values);
11160        let block_info = BlockInfo::default();
11161
11162        assert_eq!(bits_per_value % 8, 0);
11163        let data = match bits_per_value {
11164            32 => {
11165                let values = (0..num_values)
11166                    .map(|i| (i % cardinality) as u32)
11167                    .collect::<Vec<_>>();
11168                crate::buffer::LanceBuffer::reinterpret_vec(values)
11169            }
11170            64 => {
11171                let values = (0..num_values).map(|i| i % cardinality).collect::<Vec<_>>();
11172                crate::buffer::LanceBuffer::reinterpret_vec(values)
11173            }
11174            128 => {
11175                let values = (0..num_values)
11176                    .map(|i| (i % cardinality) as u128)
11177                    .collect::<Vec<_>>();
11178                crate::buffer::LanceBuffer::reinterpret_vec(values)
11179            }
11180            _ => unreachable!(),
11181        };
11182        DataBlock::FixedWidth(FixedWidthDataBlock {
11183            bits_per_value,
11184            data,
11185            num_values,
11186            block_info,
11187        })
11188    }
11189
11190    /// Helper to create VariableWidth (string) test data block with exact cardinality
11191    fn create_test_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
11192        use arrow_array::StringArray;
11193
11194        assert!(cardinality <= num_values && cardinality > 0);
11195
11196        let mut values = Vec::with_capacity(num_values as usize);
11197        for i in 0..num_values {
11198            values.push(format!("value_{:016}", i % cardinality));
11199        }
11200
11201        let array = StringArray::from(values);
11202        DataBlock::from_array(Arc::new(array) as ArrayRef)
11203    }
11204
11205    fn create_sorted_string_array(num_values: u64, cardinality: u64) -> ArrayRef {
11206        use arrow_array::StringArray;
11207
11208        assert!(cardinality <= num_values && cardinality > 0);
11209
11210        let mut values = Vec::with_capacity(num_values as usize);
11211        for i in 0..num_values {
11212            let value_idx = i * cardinality / num_values;
11213            values.push(format!("value_{:016}", value_idx));
11214        }
11215
11216        Arc::new(StringArray::from(values)) as ArrayRef
11217    }
11218
11219    fn create_sorted_variable_width_block(num_values: u64, cardinality: u64) -> DataBlock {
11220        DataBlock::from_array(create_sorted_string_array(num_values, cardinality))
11221    }
11222
11223    #[test]
11224    fn test_should_dictionary_encode() {
11225        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11226        use lance_core::datatypes::Field as LanceField;
11227
11228        // Create data where dict encoding saves space
11229        let block = create_test_variable_width_block(1000, 10);
11230
11231        let mut metadata = HashMap::new();
11232        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11233        let arrow_field =
11234            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11235        let field = LanceField::try_from(&arrow_field).unwrap();
11236
11237        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11238            &block,
11239            &field,
11240            FixedWidthDictionaryEncoding::Exclude64Bit,
11241        );
11242
11243        assert!(
11244            result.is_some(),
11245            "Should use dictionary encode based on size"
11246        );
11247    }
11248
11249    #[test]
11250    fn test_block_sampling_detects_low_cardinality_in_short_sorted_runs() {
11251        let sample_count: usize = 4096;
11252        let num_values: u64 = 200_000;
11253        let cardinality: u64 = 8_000;
11254        let run_length = num_values / cardinality;
11255        let stride = num_values as usize / sample_count;
11256        assert!(
11257            stride > run_length as usize,
11258            "test must construct the stride > run_length case"
11259        );
11260
11261        let block = create_sorted_variable_width_block(num_values, cardinality);
11262        let sample_unique_ratio =
11263            PrimitiveStructuralEncoder::sample_unique_ratio(&block, sample_count).unwrap();
11264
11265        assert!(
11266            sample_unique_ratio.is_some_and(|ratio| ratio < 0.98),
11267            "sorted low-cardinality data must not be classified as near-unique"
11268        );
11269    }
11270
11271    #[test]
11272    fn test_should_dictionary_encode_sorted_low_cardinality() {
11273        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11274        use lance_core::datatypes::Field as LanceField;
11275
11276        let block = create_sorted_variable_width_block(200_000, 8_000);
11277
11278        let mut metadata = HashMap::new();
11279        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11280        let arrow_field =
11281            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11282        let field = LanceField::try_from(&arrow_field).unwrap();
11283
11284        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11285            &block,
11286            &field,
11287            FixedWidthDictionaryEncoding::Include64Bit,
11288        );
11289
11290        assert!(
11291            result.is_some(),
11292            "sorted low-cardinality data should reach dictionary encoding"
11293        );
11294    }
11295
11296    #[test]
11297    fn test_should_not_dictionary_encode_sorted_high_cardinality_short_runs() {
11298        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11299        use lance_core::datatypes::Field as LanceField;
11300
11301        let num_values = 200_002;
11302        let cardinality = 100_001;
11303        let block = create_sorted_variable_width_block(num_values, cardinality);
11304
11305        let mut metadata = HashMap::new();
11306        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11307        let arrow_field =
11308            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11309        let field = LanceField::try_from(&arrow_field).unwrap();
11310
11311        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11312            &block,
11313            &field,
11314            FixedWidthDictionaryEncoding::Include64Bit,
11315        );
11316
11317        assert!(
11318            result.is_none(),
11319            "sorted high-cardinality short runs should not trigger a full dictionary probe"
11320        );
11321    }
11322
11323    #[tokio::test]
11324    async fn test_encode_sorted_low_cardinality_uses_dictionary_layout() {
11325        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11326
11327        let mut metadata = HashMap::new();
11328        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11329        let field = arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11330        let array = create_sorted_string_array(200_000, 8_000);
11331
11332        let page = encode_first_page(field, array, TestEncoding::StructuralU32).await;
11333        let _ = dictionary_encoding_from_page(&page);
11334    }
11335
11336    #[test]
11337    fn test_should_not_dictionary_encode_unsupported_bits() {
11338        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11339        use lance_core::datatypes::Field as LanceField;
11340
11341        let block = create_test_fixed_data_block(1000, 1000, 32);
11342
11343        let mut metadata = HashMap::new();
11344        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.8".to_string());
11345        let arrow_field =
11346            arrow_schema::Field::new("test", DataType::Int32, false).with_metadata(metadata);
11347        let field = LanceField::try_from(&arrow_field).unwrap();
11348
11349        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11350            &block,
11351            &field,
11352            FixedWidthDictionaryEncoding::Exclude64Bit,
11353        );
11354
11355        assert!(
11356            result.is_none(),
11357            "Should not use dictionary encode for unsupported bit width"
11358        );
11359    }
11360
11361    #[test]
11362    fn test_should_not_dictionary_encode_near_unique_sample() {
11363        use crate::constants::DICT_SIZE_RATIO_META_KEY;
11364        use lance_core::datatypes::Field as LanceField;
11365
11366        let num_values = 5000;
11367        let block = create_test_variable_width_block(num_values, num_values);
11368
11369        let mut metadata = HashMap::new();
11370        metadata.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "1.0".to_string());
11371        let arrow_field =
11372            arrow_schema::Field::new("test", DataType::Utf8, false).with_metadata(metadata);
11373        let field = LanceField::try_from(&arrow_field).unwrap();
11374
11375        let result = PrimitiveStructuralEncoder::should_dictionary_encode(
11376            &block,
11377            &field,
11378            FixedWidthDictionaryEncoding::Exclude64Bit,
11379        );
11380
11381        assert!(
11382            result.is_none(),
11383            "Should not probe dictionary encoding for near-unique data"
11384        );
11385    }
11386
11387    #[test]
11388    fn test_v2_1_miniblock_serializes_log_num_values_15() {
11389        let miniblocks = MiniBlockCompressed {
11390            data: vec![LanceBuffer::from(vec![1_u8; 16])],
11391            chunks: vec![
11392                MiniBlockChunk {
11393                    buffer_sizes: vec![8],
11394                    log_num_values: 15,
11395                },
11396                MiniBlockChunk {
11397                    buffer_sizes: vec![8],
11398                    log_num_values: 0,
11399                },
11400            ],
11401            num_values: 32_769,
11402        };
11403
11404        let serialized = PrimitiveStructuralEncoder::serialize_miniblocks(
11405            miniblocks,
11406            None,
11407            None,
11408            MiniblockChunkSize::U16,
11409        )
11410        .unwrap();
11411
11412        let chunk_metadata = serialized.metadata.borrow_to_typed_slice::<u16>();
11413        assert_eq!(chunk_metadata.len(), 2);
11414        assert_eq!(
11415            chunk_metadata[0] & 0x0F,
11416            15,
11417            "V2.1 metadata should use all 4 bits for log_num_values"
11418        );
11419    }
11420
11421    async fn encode_first_page(
11422        field: arrow_schema::Field,
11423        array: ArrayRef,
11424        version: TestEncoding,
11425    ) -> crate::encoder::EncodedPage {
11426        use crate::repdef::RepDefBuilder;
11427        use crate::{
11428            encoder::{
11429                ColumnIndexSequence, EncodingOptions, MIN_PAGE_BUFFER_ALIGNMENT, OutOfLineBuffers,
11430            },
11431            testing::{create_test_field_encoder, test_encoding_strategy},
11432        };
11433
11434        let lance_field = lance_core::datatypes::Field::try_from(&field).unwrap();
11435        let encoding_strategy = test_encoding_strategy(version);
11436        let mut column_index_seq = ColumnIndexSequence::default();
11437        let encoding_options = EncodingOptions {
11438            cache_bytes_per_column: 1,
11439            max_page_bytes: 32 * 1024 * 1024,
11440            keep_original_array: true,
11441            buffer_alignment: MIN_PAGE_BUFFER_ALIGNMENT,
11442        };
11443
11444        let mut encoder = create_test_field_encoder(
11445            encoding_strategy.as_ref(),
11446            &lance_field,
11447            &mut column_index_seq,
11448            &encoding_options,
11449        )
11450        .unwrap();
11451
11452        let mut external_buffers = OutOfLineBuffers::new(0, MIN_PAGE_BUFFER_ALIGNMENT);
11453        let repdef = RepDefBuilder::default();
11454        let num_rows = array.len() as u64;
11455        let mut pages = Vec::new();
11456        for task in encoder
11457            .maybe_encode(array, &mut external_buffers, repdef, 0, num_rows)
11458            .unwrap()
11459        {
11460            pages.push(task.await.unwrap());
11461        }
11462        for task in encoder.flush(&mut external_buffers).unwrap() {
11463            pages.push(task.await.unwrap());
11464        }
11465        pages.into_iter().next().unwrap()
11466    }
11467
11468    #[tokio::test]
11469    async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() {
11470        use crate::format::pb21::page_layout::Layout;
11471
11472        let val = vec![0xABu8; 33];
11473        let arr: ArrayRef = Arc::new(
11474            arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
11475                std::iter::repeat_n(Some(val.as_slice()), 256),
11476                33,
11477            )
11478            .unwrap(),
11479        );
11480        let field = arrow_schema::Field::new("c", DataType::FixedSizeBinary(33), true);
11481        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11482
11483        let PageEncoding::Structural(layout) = &page.description else {
11484            panic!("Expected structural encoding");
11485        };
11486        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11487            panic!("Expected constant layout in slot 2");
11488        };
11489        assert!(layout.inline_value.is_none());
11490        assert_eq!(page.data.len(), 1);
11491
11492        let test_cases = TestCases::default()
11493            .with_encoding(TestEncoding::StructuralU32)
11494            .with_page_sizes(vec![4096]);
11495        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11496    }
11497
11498    #[tokio::test]
11499    async fn test_constant_layout_out_of_line_utf8_v2_2() {
11500        use crate::format::pb21::page_layout::Layout;
11501
11502        let arr: ArrayRef = Arc::new(arrow_array::StringArray::from_iter_values(
11503            std::iter::repeat_n("hello", 512),
11504        ));
11505        let field = arrow_schema::Field::new("c", DataType::Utf8, true);
11506        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11507
11508        let PageEncoding::Structural(layout) = &page.description else {
11509            panic!("Expected structural encoding");
11510        };
11511        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11512            panic!("Expected constant layout in slot 2");
11513        };
11514        assert!(layout.inline_value.is_none());
11515        assert_eq!(page.data.len(), 1);
11516
11517        let test_cases = TestCases::default()
11518            .with_encoding(TestEncoding::StructuralU32)
11519            .with_page_sizes(vec![4096]);
11520        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11521    }
11522
11523    #[tokio::test]
11524    async fn test_constant_layout_nullable_item_v2_2() {
11525        use crate::format::pb21::page_layout::Layout;
11526
11527        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![
11528            Some(7),
11529            None,
11530            Some(7),
11531            None,
11532            Some(7),
11533        ]));
11534        let field = arrow_schema::Field::new("c", DataType::Int32, true);
11535        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11536
11537        let PageEncoding::Structural(layout) = &page.description else {
11538            panic!("Expected structural encoding");
11539        };
11540        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11541            panic!("Expected constant layout in slot 2");
11542        };
11543        assert!(layout.inline_value.is_some());
11544        assert_eq!(page.data.len(), 2);
11545
11546        let test_cases = TestCases::default()
11547            .with_encoding(TestEncoding::StructuralU32)
11548            .with_page_sizes(vec![4096]);
11549        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11550    }
11551
11552    #[tokio::test]
11553    async fn test_constant_layout_list_repdef_v2_2() {
11554        use crate::format::pb21::page_layout::Layout;
11555        use arrow_array::builder::{Int32Builder, ListBuilder};
11556
11557        let mut builder = ListBuilder::new(Int32Builder::new());
11558        builder.values().append_value(7);
11559        builder.values().append_null();
11560        builder.values().append_value(7);
11561        builder.append(true);
11562
11563        builder.append(true);
11564
11565        builder.values().append_value(7);
11566        builder.append(true);
11567
11568        builder.append_null();
11569
11570        let arr: ArrayRef = Arc::new(builder.finish());
11571        let field = arrow_schema::Field::new(
11572            "c",
11573            DataType::List(Arc::new(arrow_schema::Field::new(
11574                "item",
11575                DataType::Int32,
11576                true,
11577            ))),
11578            true,
11579        );
11580        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11581
11582        let PageEncoding::Structural(layout) = &page.description else {
11583            panic!("Expected structural encoding");
11584        };
11585        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11586            panic!("Expected constant layout in slot 2");
11587        };
11588        assert!(layout.inline_value.is_some());
11589        assert_eq!(page.data.len(), 2);
11590
11591        let test_cases = TestCases::default()
11592            .with_encoding(TestEncoding::StructuralU32)
11593            .with_page_sizes(vec![4096]);
11594        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11595    }
11596
11597    #[tokio::test]
11598    async fn test_constant_layout_fixed_size_list_not_used_v2_2() {
11599        use crate::format::pb21::page_layout::Layout;
11600        use arrow_array::builder::{FixedSizeListBuilder, Int32Builder};
11601
11602        let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
11603        for _ in 0..64 {
11604            builder.values().append_value(1);
11605            builder.values().append_null();
11606            builder.values().append_value(3);
11607            builder.append(true);
11608        }
11609        let arr: ArrayRef = Arc::new(builder.finish());
11610        let field = arrow_schema::Field::new(
11611            "c",
11612            DataType::FixedSizeList(
11613                Arc::new(arrow_schema::Field::new("item", DataType::Int32, true)),
11614                3,
11615            ),
11616            true,
11617        );
11618        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11619
11620        if let PageEncoding::Structural(layout) = &page.description {
11621            assert!(
11622                !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
11623                "FixedSizeList should not use constant layout yet"
11624            );
11625        }
11626
11627        let test_cases = TestCases::default()
11628            .with_encoding(TestEncoding::StructuralU32)
11629            .with_page_sizes(vec![4096]);
11630        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11631    }
11632
11633    #[tokio::test]
11634    async fn test_constant_layout_not_written_before_v2_2() {
11635        use crate::format::pb21::page_layout::Layout;
11636
11637        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![7; 1024]));
11638        let field = arrow_schema::Field::new("c", DataType::Int32, true);
11639        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU16).await;
11640
11641        let PageEncoding::Structural(layout) = &page.description else {
11642            return;
11643        };
11644        assert!(
11645            !matches!(layout.layout.as_ref().unwrap(), Layout::ConstantLayout(_)),
11646            "Should not emit constant layout before v2.2"
11647        );
11648
11649        let test_cases = TestCases::default()
11650            .with_encoding(TestEncoding::StructuralU16)
11651            .with_page_sizes(vec![4096]);
11652        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11653    }
11654
11655    #[tokio::test]
11656    async fn test_all_null_constant_layout_still_works_v2_2() {
11657        use crate::format::pb21::page_layout::Layout;
11658
11659        let arr: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![None, None, None]));
11660        let field = arrow_schema::Field::new("c", DataType::Int32, true);
11661        let page = encode_first_page(field, arr.clone(), TestEncoding::StructuralU32).await;
11662
11663        let PageEncoding::Structural(layout) = &page.description else {
11664            panic!("Expected structural encoding");
11665        };
11666        let Layout::ConstantLayout(layout) = layout.layout.as_ref().unwrap() else {
11667            panic!("Expected layout in slot 2");
11668        };
11669        assert!(layout.inline_value.is_none());
11670        assert_eq!(page.data.len(), 0);
11671
11672        let test_cases = TestCases::default()
11673            .with_encoding(TestEncoding::StructuralU32)
11674            .with_page_sizes(vec![4096]);
11675        check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await;
11676    }
11677
11678    fn hand_built_dictionary_with_out_of_range_null_keys() -> ArrayRef {
11679        use arrow_array::{DictionaryArray, Int32Array, types::Int32Type};
11680        use arrow_buffer::NullBuffer;
11681
11682        let keys = Int32Array::new(
11683            vec![0, 7, 7].into(),
11684            Some(NullBuffer::from(vec![true, false, false])),
11685        );
11686        let values = Arc::new(StringArray::from(vec!["a"]));
11687        Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef
11688    }
11689
11690    fn concatenated_dictionary_with_out_of_range_null_keys() -> ArrayRef {
11691        use arrow_array::{builder::StringDictionaryBuilder, new_null_array, types::Int32Type};
11692
11693        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
11694        builder.append_value("a");
11695        for _ in 0..7 {
11696            builder.append_null();
11697        }
11698        let valued = Arc::new(builder.finish()) as ArrayRef;
11699        let all_null = new_null_array(valued.data_type(), 8);
11700        arrow_select::concat::concat(&[valued.as_ref(), all_null.as_ref()]).unwrap()
11701    }
11702
11703    #[rstest::rstest]
11704    #[case::hand_built(hand_built_dictionary_with_out_of_range_null_keys(), 1)]
11705    #[case::concatenated(concatenated_dictionary_with_out_of_range_null_keys(), 8)]
11706    #[tokio::test]
11707    async fn test_dictionary_out_of_range_null_keys_round_trip(
11708        #[case] dictionary: ArrayRef,
11709        #[case] split_at: usize,
11710    ) {
11711        let num_rows = dictionary.len();
11712        let dictionaries = vec![
11713            dictionary.slice(0, split_at),
11714            dictionary.slice(split_at, num_rows - split_at),
11715        ];
11716        let test_cases = TestCases::default()
11717            .with_encoding(TestEncoding::StructuralU32)
11718            .with_page_sizes(vec![1])
11719            .with_range(split_at as u64..num_rows as u64)
11720            .with_indices(vec![split_at as u64]);
11721
11722        check_round_trip_encoding_of_data(dictionaries, &test_cases, HashMap::new()).await;
11723    }
11724
11725    #[test]
11726    fn test_encode_decode_complex_all_null_vals_roundtrip() {
11727        use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy};
11728
11729        let values: Arc<[u16]> = Arc::from((0..2048).map(|i| (i % 5) as u16).collect::<Vec<u16>>());
11730
11731        let compression_strategy = crate::testing::test_compression_strategy(
11732            TestEncoding::StructuralU16,
11733            crate::compression_config::CompressionParams::default(),
11734        );
11735        let decompression_strategy = DefaultDecompressionStrategy::default();
11736
11737        let (compressed_buf, encoding) = PrimitiveStructuralEncoder::encode_complex_all_null_vals(
11738            &values,
11739            compression_strategy.as_ref(),
11740        )
11741        .unwrap();
11742
11743        let decompressor = decompression_strategy
11744            .create_block_decompressor(&encoding)
11745            .unwrap();
11746        let decompressed = decompressor
11747            .decompress(Some(compressed_buf), values.len() as u64)
11748            .unwrap();
11749        let decompressed_fixed_width = decompressed.as_fixed_width().unwrap();
11750        assert_eq!(decompressed_fixed_width.num_values, values.len() as u64);
11751        assert_eq!(decompressed_fixed_width.bits_per_value, 16);
11752        let rep_result = decompressed_fixed_width.data.borrow_to_typed_slice::<u16>();
11753        assert_eq!(rep_result.as_ref(), values.as_ref());
11754    }
11755
11756    #[tokio::test]
11757    async fn test_complex_all_null_compression_gated_by_version() {
11758        use crate::format::pb21::page_layout::Layout;
11759        use arrow_array::ListArray;
11760
11761        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11762            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
11763        );
11764        let arr: ArrayRef = Arc::new(list_array);
11765        let field = arrow_schema::Field::new(
11766            "c",
11767            DataType::List(Arc::new(arrow_schema::Field::new(
11768                "item",
11769                DataType::Int32,
11770                true,
11771            ))),
11772            true,
11773        );
11774
11775        let page_v21 =
11776            encode_first_page(field.clone(), arr.clone(), TestEncoding::StructuralU16).await;
11777        let PageEncoding::Structural(layout_v21) = &page_v21.description else {
11778            panic!("Expected structural encoding");
11779        };
11780        let Layout::ConstantLayout(layout_v21) = layout_v21.layout.as_ref().unwrap() else {
11781            panic!("Expected constant layout");
11782        };
11783        assert!(layout_v21.rep_compression.is_none());
11784        assert!(layout_v21.def_compression.is_none());
11785        assert_eq!(layout_v21.num_rep_values, 0);
11786        assert_eq!(layout_v21.num_def_values, 0);
11787
11788        let page_v22 = encode_first_page(field, arr, TestEncoding::StructuralU32).await;
11789        let PageEncoding::Structural(layout_v22) = &page_v22.description else {
11790            panic!("Expected structural encoding");
11791        };
11792        let Layout::ConstantLayout(layout_v22) = layout_v22.layout.as_ref().unwrap() else {
11793            panic!("Expected constant layout");
11794        };
11795        assert!(layout_v22.def_compression.is_some());
11796        assert!(layout_v22.num_def_values > 0);
11797    }
11798
11799    #[tokio::test]
11800    async fn test_complex_all_null_round_trip() {
11801        use arrow_array::ListArray;
11802
11803        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11804            (0..1000).map(|i| if i % 2 == 0 { None } else { Some(vec![]) }),
11805        );
11806
11807        let test_cases = TestCases::default().with_u32_structural_encodings();
11808        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
11809            .await;
11810    }
11811
11812    #[tokio::test]
11813    async fn test_complex_all_null_constant_def_round_trip() {
11814        use arrow_array::ListArray;
11815
11816        // Every row is a null list => constant def levels => a single RLE run,
11817        // exercising the lazy run-form decode end to end.
11818        let list_array = ListArray::from_iter_primitive::<arrow_array::types::Int32Type, _, _>(
11819            (0..5000).map(|_| None::<Vec<Option<i32>>>),
11820        );
11821
11822        let test_cases = TestCases::default().with_u32_structural_encodings();
11823        check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new())
11824            .await;
11825    }
11826
11827    fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer {
11828        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
11829            data: LanceBuffer::reinterpret_slice(Arc::from(levels)),
11830            bits_per_value: 16,
11831            num_values: levels.len() as u64,
11832            block_info: BlockInfo::new(),
11833        });
11834        BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block)
11835            .unwrap()
11836            .0
11837            .unwrap()
11838    }
11839
11840    fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns {
11841        let frame = encoded_u16_frame(levels, run_length_width);
11842        RleDecompressor::with_run_length_width(16, run_length_width)
11843            .decode_u16_runs(frame, levels.len() as u64)
11844            .unwrap()
11845    }
11846
11847    #[test]
11848    fn miniblock_levels_use_one_count_for_mixed_codecs() {
11849        let actual_num_levels = usize::from(u16::MAX) + 8;
11850        let levels = vec![1_u16; actual_num_levels];
11851        let rep_frame = encoded_u16_frame(&levels, RunLengthWidth::U8);
11852        let rep_decompressor = RleDecompressor::new(16);
11853        let def_frame = LanceBuffer::reinterpret_slice(Arc::from(levels.clone()));
11854        let def_decompressor = ValueDecompressor::from_flat(&pb21::Flat {
11855            bits_per_value: 16,
11856            data: None,
11857        });
11858
11859        let num_levels = DecodeMiniBlockTask::resolve_num_levels(
11860            Some(&rep_decompressor),
11861            Some(&rep_frame),
11862            Some(&def_decompressor),
11863            Some(&def_frame),
11864            7,
11865        )
11866        .unwrap();
11867        assert_eq!(num_levels, actual_num_levels as u64);
11868
11869        let rep =
11870            DecodeMiniBlockTask::decode_levels(&rep_decompressor, rep_frame, num_levels).unwrap();
11871        let def =
11872            DecodeMiniBlockTask::decode_levels(&def_decompressor, def_frame, num_levels).unwrap();
11873        assert_eq!(rep.as_ref(), levels);
11874        assert_eq!(def, rep);
11875    }
11876
11877    #[test]
11878    fn miniblock_levels_reject_non_wrapped_count_mismatch() {
11879        let frame = encoded_u16_frame(&[1_u16; 8], RunLengthWidth::U8);
11880        let decompressor = RleDecompressor::new(16);
11881
11882        let error = DecodeMiniBlockTask::resolve_num_levels(
11883            Some(&decompressor),
11884            Some(&frame),
11885            None,
11886            None,
11887            7,
11888        )
11889        .unwrap_err();
11890        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
11891        assert!(error.to_string().contains("not congruent modulo 65536"));
11892    }
11893
11894    #[test]
11895    fn miniblock_levels_reject_cross_stream_count_disagreement() {
11896        let rep_levels = vec![1_u16; usize::from(u16::MAX) + 8];
11897        let def_levels = vec![1_u16; rep_levels.len() + usize::from(u16::MAX) + 1];
11898        let rep_frame = encoded_u16_frame(&rep_levels, RunLengthWidth::U8);
11899        let def_frame = encoded_u16_frame(&def_levels, RunLengthWidth::U8);
11900        let decompressor = RleDecompressor::new(16);
11901
11902        let error = DecodeMiniBlockTask::resolve_num_levels(
11903            Some(&decompressor),
11904            Some(&rep_frame),
11905            Some(&decompressor),
11906            Some(&def_frame),
11907            7,
11908        )
11909        .unwrap_err();
11910        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
11911        assert!(
11912            error
11913                .to_string()
11914                .contains("structural streams disagree on the level count")
11915        );
11916    }
11917
11918    fn physical_levels(levels: &[u16]) -> LazyLevels {
11919        LazyLevels::Runs(Arc::new(RunStorage::Physical(
11920            encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(),
11921        )))
11922    }
11923
11924    fn coalesced_levels(levels: &[u16]) -> LazyLevels {
11925        let mut values = Vec::new();
11926        let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len());
11927        for (index, &value) in levels.iter().enumerate() {
11928            if values.last() == Some(&value) {
11929                ends.set_last(index + 1).unwrap();
11930            } else {
11931                values.push(value);
11932                ends.push(index + 1).unwrap();
11933            }
11934        }
11935        LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
11936            values: values.into_boxed_slice(),
11937            ends: ends.finish(),
11938        }))
11939    }
11940
11941    #[test]
11942    fn lazy_levels_runs_match_dense() {
11943        // Runs: 3x2, 1x1, 3x3, 0x2  =>  [3,3,1,3,3,3,0,0]
11944        let expanded: Vec<u16> = vec![3, 3, 1, 3, 3, 3, 0, 0];
11945        let coalesced = coalesced_levels(&expanded);
11946        let physical = physical_levels(&expanded);
11947        let dense = LazyLevels::Dense(ScalarBuffer::<u16>::from(expanded.clone()));
11948        let n = expanded.len();
11949
11950        assert_eq!(coalesced.len(), n);
11951        assert_eq!(physical.len(), n);
11952        assert_eq!(dense.len(), n);
11953
11954        // Rows begin at each `max_rep` (3) position; row `num_rows` maps to `len`.
11955        let max_rep = 3u16;
11956        let row_starts: Vec<usize> = (0..n).filter(|&i| expanded[i] == max_rep).collect();
11957        for target in 0..=row_starts.len() as u64 {
11958            let want = row_starts.get(target as usize).copied().unwrap_or(n);
11959            for runs in [&coalesced, &physical] {
11960                let mut cursor = LevelCursor::default();
11961                assert_eq!(
11962                    runs.seek_row_start(&mut cursor, target, max_rep).unwrap(),
11963                    want,
11964                    "seek_row_start({target})"
11965                );
11966            }
11967            let mut c_dense = LevelCursor::default();
11968            assert_eq!(
11969                dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(),
11970                want
11971            );
11972        }
11973
11974        // `count_le_cursor` (fresh cursor per range) and `extend_into` agree with
11975        // the dense reference on every sub-range.
11976        for start in 0..=n {
11977            for end in start..=n {
11978                for max in [0u16, 1, 2, 3] {
11979                    let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64;
11980                    for runs in [&coalesced, &physical] {
11981                        let mut cursor = RunPosition::default();
11982                        assert_eq!(
11983                            runs.count_le_cursor(&mut cursor, start..end, max).0,
11984                            want,
11985                            "count_le_cursor({start}..{end}, {max})"
11986                        );
11987                    }
11988                    let mut d_cur = RunPosition::default();
11989                    assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want);
11990                }
11991                for runs in [&coalesced, &physical] {
11992                    let mut got = Vec::new();
11993                    runs.extend_into(start..end, RunPosition::default(), &mut got);
11994                    assert_eq!(
11995                        got,
11996                        expanded[start..end].to_vec(),
11997                        "extend_into({start}..{end})"
11998                    );
11999                }
12000                let mut got_dense = Vec::new();
12001                dense.extend_into(start..end, RunPosition::default(), &mut got_dense);
12002                assert_eq!(got_dense, expanded[start..end].to_vec());
12003            }
12004        }
12005    }
12006
12007    #[test]
12008    fn physical_run_hints_support_deferred_materialization() {
12009        let expanded: Vec<u16> = vec![3, 3, 1, 1, 2, 2, 0, 0];
12010        let physical = physical_levels(&expanded);
12011        let LazyLevels::Runs(runs) = &physical else {
12012            panic!("expected physical runs");
12013        };
12014        let mut first_hint = RunPosition::default();
12015        runs.seek(&mut first_hint, 2);
12016        let mut second_hint = RunPosition::default();
12017        runs.seek(&mut second_hint, 6);
12018
12019        let mut second = Vec::new();
12020        physical.extend_into(6..8, second_hint, &mut second);
12021        let mut first = Vec::new();
12022        physical.extend_into(2..4, first_hint, &mut first);
12023        assert_eq!(second, expanded[6..8]);
12024        assert_eq!(first, expanded[2..4]);
12025    }
12026
12027    /// Fuzz parity for the run-oriented complex-all-null drain: the cursor walk
12028    /// over `LazyLevels` must yield the exact level slices and visible
12029    /// count that a brute-force reference over the fully expanded levels does, for
12030    /// dense, physical-run, and coalesced-run forms and arbitrarily shaped range requests.
12031    mod complex_all_null_drain_parity {
12032        use std::ops::Range;
12033
12034        use arrow_buffer::ScalarBuffer;
12035        use proptest::prelude::*;
12036
12037        use super::super::{LazyLevels, LevelCursor, RunPosition};
12038        use super::{coalesced_levels, physical_levels};
12039        use crate::Result;
12040
12041        #[derive(Debug, Clone)]
12042        struct DrainInput {
12043            max_rep: u16,
12044            max_visible: u16,
12045            rep: Option<Vec<u16>>,
12046            def: Option<Vec<u16>>,
12047            ranges: Vec<Range<u64>>,
12048        }
12049
12050        fn dense_levels(levels: &[u16]) -> LazyLevels {
12051            LazyLevels::Dense(ScalarBuffer::from(levels.to_vec()))
12052        }
12053
12054        fn rle_levels(levels: &[u16]) -> LazyLevels {
12055            coalesced_levels(levels)
12056        }
12057
12058        fn seek(
12059            rep: Option<&LazyLevels>,
12060            cursor: &mut LevelCursor,
12061            row: u64,
12062            max_rep: u16,
12063        ) -> Result<usize> {
12064            match rep {
12065                Some(rep) => rep.seek_row_start(cursor, row, max_rep),
12066                None => {
12067                    cursor.row = row;
12068                    cursor.level = row as usize;
12069                    Ok(row as usize)
12070                }
12071            }
12072        }
12073
12074        /// Mirror of `ComplexAllNullPageDecoder::drain`, driving the real
12075        /// `seek_row_start` / `count_le_cursor` with monotonic cursors.
12076        fn simulate_drain(
12077            rep: Option<&LazyLevels>,
12078            def: Option<&LazyLevels>,
12079            max_rep: u16,
12080            max_visible: u16,
12081            ranges: &[Range<u64>],
12082        ) -> Result<(Vec<Range<usize>>, u64)> {
12083            let mut rep_cursor = LevelCursor::default();
12084            let mut def_run_cursor = RunPosition::default();
12085            let mut slices: Vec<Range<usize>> = Vec::new();
12086            let mut visible = 0u64;
12087            for range in ranges {
12088                let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?;
12089                let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?;
12090                visible += match def {
12091                    Some(def) => {
12092                        def.count_le_cursor(
12093                            &mut def_run_cursor,
12094                            level_start..level_end,
12095                            max_visible,
12096                        )
12097                        .0
12098                    }
12099                    None => (level_end - level_start) as u64,
12100                };
12101                match slices.last_mut() {
12102                    Some(last) if last.end == level_start => last.end = level_end,
12103                    _ => slices.push(level_start..level_end),
12104                }
12105            }
12106            Ok((slices, visible))
12107        }
12108
12109        /// Independent brute-force reference over fully expanded levels.
12110        fn reference_drain(
12111            rep: Option<&[u16]>,
12112            def: Option<&[u16]>,
12113            max_rep: u16,
12114            max_visible: u16,
12115            ranges: &[Range<u64>],
12116        ) -> (Vec<Range<usize>>, u64) {
12117            let total_levels = rep
12118                .map(|r| r.len())
12119                .or_else(|| def.map(|d| d.len()))
12120                .unwrap_or(0);
12121            // Level index where each row starts (or `total_levels` for the end row).
12122            let row_starts: Vec<usize> = match rep {
12123                Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(),
12124                None => (0..total_levels).collect(),
12125            };
12126            let level_of_row = |row: u64| {
12127                row_starts
12128                    .get(row as usize)
12129                    .copied()
12130                    .unwrap_or(total_levels)
12131            };
12132
12133            let mut slices: Vec<Range<usize>> = Vec::new();
12134            let mut visible = 0u64;
12135            for range in ranges {
12136                let ls = level_of_row(range.start);
12137                let le = level_of_row(range.end);
12138                visible += match def {
12139                    Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64,
12140                    None => (le - ls) as u64,
12141                };
12142                match slices.last_mut() {
12143                    Some(last) if last.end == ls => last.end = le,
12144                    _ => slices.push(ls..le),
12145                }
12146            }
12147            (slices, visible)
12148        }
12149
12150        fn ranges_strategy(num_rows: u64) -> BoxedStrategy<Vec<Range<u64>>> {
12151            if num_rows == 0 {
12152                return Just(Vec::new()).boxed();
12153            }
12154            // (gap, len) pairs; a zero gap yields ranges adjacent in row space,
12155            // which exercises the level-slice coalescing path.
12156            proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8)
12157                .prop_map(move |pairs| {
12158                    let mut ranges = Vec::new();
12159                    let mut pos = 0u64;
12160                    for (gap, len) in pairs {
12161                        pos = pos.saturating_add(gap);
12162                        if pos >= num_rows {
12163                            break;
12164                        }
12165                        let end = (pos + len).min(num_rows);
12166                        ranges.push(pos..end);
12167                        pos = end;
12168                    }
12169                    ranges
12170                })
12171                .boxed()
12172        }
12173
12174        fn drain_input() -> impl Strategy<Value = DrainInput> {
12175            (
12176                1u16..=3,
12177                0u16..=3,
12178                any::<bool>(),
12179                any::<bool>(),
12180                1usize..=48,
12181            )
12182                .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| {
12183                    // Complex-all-null always has definition levels when there is
12184                    // no repetition, so force `def` present in that case.
12185                    let has_def = has_def || !has_rep;
12186                    let rep = if has_rep {
12187                        proptest::collection::vec(0u16..=max_rep, len)
12188                            .prop_map(move |mut v| {
12189                                // Row 0 must start at a max-rep boundary.
12190                                v[0] = max_rep;
12191                                Some(v)
12192                            })
12193                            .boxed()
12194                    } else {
12195                        Just(None).boxed()
12196                    };
12197                    let def = if has_def {
12198                        proptest::collection::vec(0u16..=(max_visible + 2), len)
12199                            .prop_map(Some)
12200                            .boxed()
12201                    } else {
12202                        Just(None).boxed()
12203                    };
12204                    (Just(max_rep), Just(max_visible), rep, def)
12205                })
12206                .prop_flat_map(|(max_rep, max_visible, rep, def)| {
12207                    let num_rows = match &rep {
12208                        Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64,
12209                        None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0),
12210                    };
12211                    ranges_strategy(num_rows).prop_map(move |ranges| DrainInput {
12212                        max_rep,
12213                        max_visible,
12214                        rep: rep.clone(),
12215                        def: def.clone(),
12216                        ranges,
12217                    })
12218                })
12219        }
12220
12221        proptest! {
12222            #![proptest_config(ProptestConfig::with_cases(512))]
12223
12224            #[test]
12225            fn drain_matches_reference(input in drain_input()) {
12226                let DrainInput { max_rep, max_visible, rep, def, ranges } = input;
12227
12228                let reference =
12229                    reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges);
12230
12231                let rep_dense = rep.as_deref().map(dense_levels);
12232                let def_dense = def.as_deref().map(dense_levels);
12233                let got_dense =
12234                    simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges)
12235                        .unwrap();
12236                prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference");
12237
12238                let rep_rle = rep.as_deref().map(rle_levels);
12239                let def_rle = def.as_deref().map(rle_levels);
12240                let got_rle =
12241                    simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges)
12242                        .unwrap();
12243                prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference");
12244
12245                let rep_physical = rep.as_deref().map(physical_levels);
12246                let def_physical = def.as_deref().map(physical_levels);
12247                let got_physical =
12248                    simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges)
12249                        .unwrap();
12250                prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference");
12251            }
12252        }
12253    }
12254
12255    #[test]
12256    fn lazy_levels_runs_are_compact() {
12257        let single_run = |n: usize| {
12258            let mut ends = RunEndsBuilder::with_capacity(n, 1);
12259            ends.push(n).unwrap();
12260            LazyLevels::Runs(Arc::new(RunStorage::Coalesced {
12261                values: vec![1u16].into_boxed_slice(),
12262                ends: ends.finish(),
12263            }))
12264        };
12265        // Run-form footprint is independent of the logical length within an end width...
12266        assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size());
12267        assert!(single_run(10_000_000).deep_size() < 100);
12268        assert_eq!(single_run(10_000_000).len(), 10_000_000);
12269        // ...while Dense pays 2 bytes per value.
12270        assert_eq!(
12271            LazyLevels::Dense(ScalarBuffer::<u16>::from(vec![1u16; 1000])).deep_size(),
12272            2000
12273        );
12274    }
12275
12276    #[test]
12277    fn lazy_levels_selects_smallest_representation() {
12278        let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8);
12279        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
12280
12281        let equal_size: Vec<u16> = std::iter::repeat_n(0, 256)
12282            .chain(std::iter::repeat_n(1, 100))
12283            .chain(std::iter::repeat_n(2, 100))
12284            .collect();
12285        let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8);
12286        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
12287
12288        let moderate_runs: Vec<u16> = (0..250)
12289            .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
12290            .collect();
12291        let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8);
12292        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
12293
12294        let split_constant = vec![7u16; 5000];
12295        let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8);
12296        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced);
12297
12298        let high_density: Vec<u16> = (0..70_000).map(|index| (index % 2) as u16).collect();
12299        let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32);
12300        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense);
12301    }
12302
12303    #[test]
12304    fn physical_runs_detach_from_large_encoded_frame() {
12305        let levels: Vec<u16> = (0..250)
12306            .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4))
12307            .collect();
12308        let frame = encoded_u16_frame(&levels, RunLengthWidth::U8);
12309        let frame_offset = 4096;
12310        let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000];
12311        allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref());
12312        let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len());
12313        let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8)
12314            .decode_u16_runs(frame, levels.len() as u64)
12315            .unwrap();
12316
12317        assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical);
12318        let cached = LazyLevels::from_rle_runs(runs).unwrap();
12319        assert!(
12320            matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_)))
12321        );
12322        assert_eq!(cached.len(), levels.len());
12323        assert!(cached.deep_size() < 4096);
12324    }
12325
12326    #[test]
12327    fn complex_all_null_levels_reject_invalid_values_and_lengths() {
12328        let invalid_levels = vec![0u16, 3];
12329        for levels in [
12330            LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())),
12331            physical_levels(&invalid_levels),
12332            coalesced_levels(&invalid_levels),
12333        ] {
12334            let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err();
12335            assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12336            assert!(error.to_string().contains("Invalid definition level 3"));
12337        }
12338
12339        let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2])));
12340        let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16])));
12341        let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err();
12342        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12343        assert!(
12344            error
12345                .to_string()
12346                .contains("repetition has 2, definition has 1")
12347        );
12348    }
12349
12350    #[test]
12351    fn block_levels_reject_malformed_payload_size() {
12352        let block = DataBlock::FixedWidth(FixedWidthDataBlock {
12353            data: LanceBuffer::from(vec![0]),
12354            bits_per_value: 16,
12355            num_values: 1,
12356            block_info: BlockInfo::new(),
12357        });
12358        let error = dense_levels_from_block(block, 1, "definition").unwrap_err();
12359        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12360        assert!(
12361            error
12362                .to_string()
12363                .contains("expected 2 bytes for 1 values, got 1")
12364        );
12365    }
12366
12367    #[test]
12368    fn complex_all_null_level_codec_validates_rle_metadata() {
12369        let encoding = pb21::CompressiveEncoding {
12370            compression: Some(Compression::Rle(Box::new(pb21::Rle {
12371                values: None,
12372                run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))),
12373            }))),
12374        };
12375
12376        let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default())
12377            .unwrap_err();
12378        assert!(matches!(error, lance_core::Error::InvalidInput { .. }));
12379        assert!(
12380            error
12381                .to_string()
12382                .contains("RLE compression missing values encoding")
12383        );
12384    }
12385
12386    // https://github.com/lance-format/lance/issues/6681
12387    #[tokio::test]
12388    async fn test_sparse_boolean_list_roundtrip() {
12389        use arrow_array::builder::{BooleanBuilder, ListBuilder};
12390
12391        let mut list_builder = ListBuilder::new(BooleanBuilder::new());
12392        for i in 0..1000i32 {
12393            if i % 64 == 0 {
12394                // Alternate true/false so the array is not constant (constant path avoids the bug).
12395                list_builder.values().append_value(i % 128 == 0);
12396                list_builder.append(true);
12397            } else {
12398                list_builder.append(false);
12399            }
12400        }
12401        let list_array = Arc::new(list_builder.finish());
12402
12403        let test_cases = TestCases::default().with_structural_encodings();
12404        check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await;
12405    }
12406
12407    fn truncated_tail_details() -> std::sync::Arc<super::FullZipDecodeDetails> {
12408        use crate::compression::VariablePerValueDecompressor;
12409        use crate::encodings::physical::binary::VariableDecoder;
12410        use crate::repdef::{ControlWordParser, DefinitionInterpretation};
12411        use std::sync::Arc;
12412        Arc::new(super::FullZipDecodeDetails {
12413            value_decompressor: super::PerValueDecompressor::Variable(Arc::new(
12414                VariableDecoder::default(),
12415            )
12416                as Arc<dyn VariablePerValueDecompressor>),
12417            def_meaning: vec![DefinitionInterpretation::NullableItem].into(),
12418            ctrl_word_parser: ControlWordParser::new(0, 0),
12419            max_rep: 0,
12420            max_visible_def: 0,
12421        })
12422    }
12423
12424    fn decode_variable_full_zip(
12425        buf: Vec<u8>,
12426        bits_per_offset: u8,
12427    ) -> lance_core::Result<super::VariableFullZipDecoder> {
12428        use std::collections::VecDeque;
12429        let mut data = VecDeque::new();
12430        data.push_back(crate::buffer::LanceBuffer::from(buf));
12431        super::VariableFullZipDecoder::new(
12432            truncated_tail_details(),
12433            data,
12434            1,
12435            bits_per_offset,
12436            bits_per_offset,
12437        )
12438    }
12439
12440    /// A well-formed length prefix decodes without incident, for both widths.
12441    #[test]
12442    fn variable_full_zip_wellformed_length_prefix() {
12443        assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok());
12444        assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok());
12445    }
12446
12447    /// A page whose item walk ends with a partial length prefix must surface a
12448    /// corrupt-file error rather than read past the end of the buffer.
12449    ///
12450    /// This asserts the error variant and message rather than merely expecting a
12451    /// panic: before the length prefix was bounds checked, the read was
12452    /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here
12453    /// (which a `#[should_panic]` test would have accepted as a pass) while a
12454    /// release build read up to 8 bytes out of a 4 byte allocation.
12455    #[test]
12456    fn variable_full_zip_truncated_length_prefix_is_corrupt_file() {
12457        use lance_core::Error;
12458
12459        for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] {
12460            let err = decode_variable_full_zip(vec![0xAA; buf_len], bits)
12461                .expect_err("a truncated length prefix must not decode");
12462            assert!(
12463                matches!(err, Error::CorruptFile { .. }),
12464                "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}",
12465                bits,
12466                buf_len,
12467                err
12468            );
12469            let msg = err.to_string();
12470            assert!(
12471                msg.contains("truncated length prefix"),
12472                "error should say what is wrong, got: {msg}"
12473            );
12474        }
12475    }
12476
12477    /// Drains `num_rows` from a single-buffer fixed full-zip page built from `buf`.
12478    ///
12479    /// `bits_rep` picks the control word width (1 -> one byte, 9 -> two bytes) and each
12480    /// visible item is a one-word control plus a 4 byte value.
12481    fn drain_fixed_full_zip(buf: Vec<u8>, bits_rep: u8, num_rows: u64) -> lance_core::Result<()> {
12482        use crate::compression::FixedPerValueDecompressor;
12483        use crate::decoder::StructuralPageDecoder;
12484        use crate::repdef::{ControlWordParser, DefinitionInterpretation};
12485        use std::collections::VecDeque;
12486        use std::sync::Arc;
12487
12488        #[derive(Debug)]
12489        struct UnusedFixedDecompressor;
12490
12491        impl FixedPerValueDecompressor for UnusedFixedDecompressor {
12492            fn decompress(&self, _: FixedWidthDataBlock, _: u64) -> crate::Result<DataBlock> {
12493                unreachable!("these cases fail while slicing, before any value is decoded")
12494            }
12495
12496            fn bits_per_value(&self) -> u64 {
12497                32
12498            }
12499        }
12500
12501        let details = Arc::new(super::FullZipDecodeDetails {
12502            value_decompressor: super::PerValueDecompressor::Fixed(Arc::new(
12503                UnusedFixedDecompressor,
12504            )),
12505            def_meaning: vec![DefinitionInterpretation::AllValidList].into(),
12506            ctrl_word_parser: ControlWordParser::new(bits_rep, 0),
12507            max_rep: 1,
12508            max_visible_def: 0,
12509        });
12510        let bytes_per_word = bits_rep.div_ceil(8) as usize;
12511
12512        let mut data = VecDeque::new();
12513        data.push_back(crate::buffer::LanceBuffer::from(buf));
12514        let mut decoder = super::FixedFullZipDecoder {
12515            details,
12516            data,
12517            offset_in_current: 0,
12518            bytes_per_value: 4,
12519            total_bytes_per_value: 4 + bytes_per_word,
12520            num_rows,
12521        };
12522        decoder.drain(num_rows).map(|_| ())
12523    }
12524
12525    /// Two one-item rows, each a control word plus a 4 byte value, slice cleanly.
12526    #[test]
12527    fn fixed_full_zip_wellformed_page_slices() {
12528        assert!(drain_fixed_full_zip(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 2).is_ok());
12529    }
12530
12531    /// The item walk is driven by the page's own control words, so a page that disagrees
12532    /// with its layout can run off the end of the buffer.  Each way of doing so must
12533    /// surface a corrupt-file error.
12534    ///
12535    /// These assert the error variant and message rather than merely expecting a panic:
12536    /// before the walk was bounds checked these aborted the process inside
12537    /// `LanceBuffer::slice_with_length`, which a `#[should_panic]` test would have
12538    /// accepted as a pass.
12539    #[rstest::rstest]
12540    // Second value is one byte short of a full 4 byte value
12541    #[case::truncated_item(vec![1, 0, 0, 0, 0, 1, 0, 0, 0], 1, 2, "truncated item")]
12542    // Two byte control words, with a lone trailing byte after the first item
12543    #[case::truncated_control_word(
12544        vec![1, 0, 0, 0, 0, 0, 1],
12545        9,
12546        2,
12547        "truncated control word"
12548    )]
12549    // A well-formed two row page cannot satisfy a request for three rows
12550    #[case::page_exhausted(vec![1, 0, 0, 0, 0, 1, 0, 0, 0, 0], 1, 3, "page data ran out")]
12551    fn fixed_full_zip_malformed_page_is_corrupt_file(
12552        #[case] buf: Vec<u8>,
12553        #[case] bits_rep: u8,
12554        #[case] num_rows: u64,
12555        #[case] expected_msg: &str,
12556    ) {
12557        use lance_core::Error;
12558
12559        let err = drain_fixed_full_zip(buf, bits_rep, num_rows)
12560            .expect_err("a malformed page must not slice");
12561        assert!(
12562            matches!(err, Error::CorruptFile { .. }),
12563            "expected CorruptFile, got: {err:?}"
12564        );
12565        let msg = err.to_string();
12566        assert!(
12567            msg.contains(expected_msg),
12568            "error should say what is wrong, got: {msg}"
12569        );
12570    }
12571}