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