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