Skip to main content

lance_file/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{
5    borrow::Cow,
6    collections::{BTreeMap, BTreeSet},
7    io::Cursor,
8    ops::Range,
9    pin::Pin,
10    sync::Arc,
11};
12
13use arrow_array::RecordBatchReader;
14use arrow_schema::Schema as ArrowSchema;
15use byteorder::{ByteOrder, LittleEndian, ReadBytesExt};
16use bytes::{Bytes, BytesMut};
17use futures::{Stream, StreamExt, stream::BoxStream};
18use lance_core::deepsize::{Context, DeepSizeOf};
19use lance_encoding::{
20    EncodingsIo,
21    decoder::{
22        ColumnInfo, DecoderConfig, DecoderPlugins, FilterExpression, PageEncoding, PageInfo,
23        ReadBatchTask, RequestedRows, SchedulerDecoderConfig, schedule_and_decode,
24        schedule_and_decode_blocking,
25    },
26    encoder::EncodedBatch,
27    version::LanceFileVersion,
28};
29use log::debug;
30use object_store::path::Path;
31use prost::{Message, Name};
32
33use lance_core::{
34    Error, Result,
35    cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache},
36    datatypes::{Field, Schema},
37};
38use lance_encoding::format::pb as pbenc;
39use lance_encoding::format::pb21 as pbenc21;
40use lance_io::{
41    ReadBatchParams,
42    scheduler::FileScheduler,
43    stream::{RecordBatchStream, RecordBatchStreamAdapter},
44};
45
46use crate::{
47    datatypes::{Fields, FieldsWithMeta},
48    format::{MAGIC, MAJOR_VERSION, MINOR_VERSION, pb, pbfile},
49    io::LanceEncodingsIo,
50    version::ConcreteFileVersion,
51    writer::PAGE_BUFFER_ALIGNMENT,
52};
53
54/// Default chunk size for reading large pages (8MiB)
55/// Pages larger than this will be split into multiple chunks during read
56pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024;
57
58// For now, we don't use global buffers for anything other than schema.  If we
59// use these later we should make them lazily loaded and then cached once loaded.
60//
61// We store their position / length for debugging purposes
62#[derive(Debug, DeepSizeOf)]
63pub struct BufferDescriptor {
64    pub position: u64,
65    pub size: u64,
66}
67
68/// Statistics summarize some of the file metadata for quick summary info
69#[derive(Debug)]
70pub struct FileStatistics {
71    /// Statistics about each of the columns in the file
72    pub columns: Vec<ColumnStatistics>,
73}
74
75/// Summary information describing a column
76#[derive(Debug)]
77pub struct ColumnStatistics {
78    /// The number of pages in the column
79    pub num_pages: usize,
80    /// The total number of data & metadata bytes in the column
81    ///
82    /// This is the compressed on-disk size
83    pub size_bytes: u64,
84}
85
86// TODO: Caching
87#[derive(Debug)]
88pub struct CachedFileMetadata {
89    /// The schema of the file
90    pub file_schema: Arc<Schema>,
91    /// The column metadatas
92    pub column_metadatas: Vec<pbfile::ColumnMetadata>,
93    pub column_infos: Vec<Arc<ColumnInfo>>,
94    /// The number of rows in the file
95    pub num_rows: u64,
96    pub file_buffers: Vec<BufferDescriptor>,
97    /// The number of bytes contained in the data page section of the file
98    pub num_data_bytes: u64,
99    /// The number of bytes contained in the column metadata (not including buffers
100    /// referenced by the metadata)
101    pub num_column_metadata_bytes: u64,
102    /// The number of bytes contained in global buffers
103    pub num_global_buffer_bytes: u64,
104    /// The number of bytes contained in the CMO and GBO tables
105    pub num_footer_bytes: u64,
106    pub major_version: u16,
107    pub minor_version: u16,
108    /// The actual total file size in bytes, as reported by the object store.
109    pub file_size_bytes: u64,
110    /// User global buffers (index >= 1) whose bytes were already captured by the
111    /// tail read that `read_all_metadata` performs at open, keyed by buffer index.
112    ///
113    /// All global buffers are laid out contiguously starting at the schema, so on
114    /// small/medium files they land inside the captured tail window. Retaining
115    /// those bytes lets `read_global_buffer` serve them with zero additional I/O.
116    /// The bytes are copied out of the tail (rather than sliced) so the much
117    /// larger tail allocation can be dropped — we only hold what we will serve.
118    ///
119    /// The schema (buffer 0) is excluded: it is already decoded at open and is
120    /// not fetched through `read_global_buffer`. Buffers that fall outside the
121    /// window (large files) are absent here and fall back to a dedicated read.
122    pub retained_global_buffers: BTreeMap<u32, Bytes>,
123}
124
125impl CachedFileMetadata {
126    /// Total file size in bytes.
127    pub fn file_size(&self) -> u64 {
128        self.file_size_bytes
129    }
130}
131
132fn column_metadata_deep_size(column_metadatas: &[pbfile::ColumnMetadata]) -> usize {
133    column_metadatas
134        .iter()
135        .map(|cm| cm.encoded_len() * 4)
136        .sum::<usize>()
137        + std::mem::size_of_val(column_metadatas)
138}
139
140impl DeepSizeOf for CachedFileMetadata {
141    fn deep_size_of_children(&self, context: &mut Context) -> usize {
142        let schema_size = self.file_schema.deep_size_of_children(context);
143
144        let buffers_size: usize = self
145            .file_buffers
146            .iter()
147            .map(|fb| fb.deep_size_of_children(context))
148            .sum();
149
150        // column_metadatas is Vec<pbfile::ColumnMetadata> (protobuf generated,
151        // does not implement DeepSizeOf). We use prost::Message::encoded_len()
152        // as a proxy for in-memory size. The decoded representation is typically
153        // several times larger than the wire format due to heap-allocated
154        // repeated/string/bytes fields, so we apply a 4x multiplier.
155        let column_metadatas_size = column_metadata_deep_size(self.column_metadatas.as_slice());
156
157        // column_infos is Vec<Arc<ColumnInfo>>. Each ColumnInfo contains
158        // page_infos (with protobuf PageEncoding), buffer offsets, and a
159        // column-level ColumnEncoding protobuf.
160        let column_infos_size = self.column_infos.deep_size_of_children(context);
161
162        // Global buffer bytes retained for zero-IO reads (copied out of the tail).
163        let retained_buffers_size = self.retained_global_buffers.deep_size_of_children(context);
164
165        schema_size
166            + buffers_size
167            + column_metadatas_size
168            + column_infos_size
169            + retained_buffers_size
170    }
171}
172
173/// Lightweight file metadata used to locate per-column metadata on demand.
174///
175/// This contains the file-level schema, row count, global buffer descriptors,
176/// and column metadata offset table. Unlike [`CachedFileMetadata`], it does not
177/// hold decoded metadata for every column.
178#[derive(Debug, DeepSizeOf)]
179pub struct FileMetadataIndex {
180    file_schema: Arc<Schema>,
181    num_rows: u64,
182    file_buffers: Vec<BufferDescriptor>,
183    column_metadata_offsets: Arc<[(u64, u64)]>,
184    num_columns: u32,
185    version: LanceFileVersion,
186    file_size_bytes: u64,
187    retained_global_buffers: BTreeMap<u32, Bytes>,
188}
189
190impl FileMetadataIndex {
191    /// Returns the total size of the file in bytes.
192    pub fn file_size(&self) -> u64 {
193        self.file_size_bytes
194    }
195
196    /// Returns the number of physical columns in the file.
197    pub fn num_columns(&self) -> u32 {
198        self.num_columns
199    }
200}
201
202#[derive(Debug)]
203struct CachedColumnMetadata {
204    column_metadata: pbfile::ColumnMetadata,
205    column_info: Arc<ColumnInfo>,
206}
207
208impl DeepSizeOf for CachedColumnMetadata {
209    fn deep_size_of_children(&self, context: &mut Context) -> usize {
210        column_metadata_deep_size(std::slice::from_ref(&self.column_metadata))
211            + self.column_info.deep_size_of_children(context)
212    }
213}
214
215#[derive(Debug, Clone)]
216struct ColumnMetadataCacheKey {
217    column_index: u32,
218}
219
220impl CacheKey for ColumnMetadataCacheKey {
221    type ValueType = CachedColumnMetadata;
222
223    fn key(&self) -> Cow<'_, str> {
224        Cow::Owned(format!("column_metadata/{}", self.column_index))
225    }
226
227    fn type_name() -> &'static str {
228        "ColumnMetadata"
229    }
230
231    fn schema() -> CacheKeySchema {
232        CacheKeySchema::new("lance.file.column-metadata-key", 1)
233    }
234
235    fn write_key(&self, builder: &mut KeyBuilder) {
236        builder.write_u32(self.column_index);
237    }
238}
239
240impl CachedFileMetadata {
241    pub fn version(&self) -> LanceFileVersion {
242        match (self.major_version, self.minor_version) {
243            (0, 3) => LanceFileVersion::V2_0,
244            (2, 0) => LanceFileVersion::V2_0,
245            (2, 1) => LanceFileVersion::V2_1,
246            (2, 2) => LanceFileVersion::V2_2,
247            (2, 3) => LanceFileVersion::V2_3,
248            _ => panic!(
249                "Unsupported version: {}.{}",
250                self.major_version, self.minor_version
251            ),
252        }
253    }
254}
255
256/// Selecting columns from a lance file requires specifying both the
257/// index of the column and the data type of the column
258///
259/// Partly, this is because it is not strictly required that columns
260/// be read into the same type.  For example, a string column may be
261/// read as a string, large_string or string_view type.
262///
263/// A read will only succeed if the decoder for a column is capable
264/// of decoding into the requested type.
265///
266/// Note that this should generally be limited to different in-memory
267/// representations of the same semantic type.  An encoding could
268/// theoretically support "casting" (e.g. int to string, etc.) but
269/// there is little advantage in doing so here.
270///
271/// Note: in order to specify a projection the user will need some way
272/// to figure out the column indices.  In the table format we do this
273/// using field IDs and keeping track of the field id->column index mapping.
274///
275/// If users are not using the table format then they will need to figure
276/// out some way to do this themselves.
277#[derive(Debug, Clone)]
278pub struct ReaderProjection {
279    /// The data types (schema) of the selected columns.  The names
280    /// of the schema are arbitrary and ignored.
281    pub schema: Arc<Schema>,
282    /// The indices of the columns to load.
283    ///
284    /// The content of this vector depends on the file version.
285    ///
286    /// In Lance File Version 2.0 we need ids for structural fields as
287    /// well as leaf fields:
288    ///
289    ///   - Primitive: the index of the column in the schema
290    ///   - List: the index of the list column in the schema
291    ///     followed by the column indices of the children
292    ///   - FixedSizeList (of primitive): the index of the column in the schema
293    ///     (this case is not nested)
294    ///   - FixedSizeList (of non-primitive): not yet implemented
295    ///   - Dictionary: same as primitive
296    ///   - Struct: the index of the struct column in the schema
297    ///     followed by the column indices of the children
298    ///
299    ///   In other words, this should be a DFS listing of the desired schema.
300    ///
301    /// In Lance File Version 2.1 we only need ids for leaf fields.  Any structural
302    /// fields are completely transparent.
303    ///
304    /// For example, if the goal is to load:
305    ///
306    ///   x: int32
307    ///   y: `struct<z: int32, w: string>`
308    ///   z: `list<int32>`
309    ///
310    /// and the schema originally used to store the data was:
311    ///
312    ///   a: `struct<x: int32>`
313    ///   b: int64
314    ///   y: `struct<z: int32, c: int64, w: string>`
315    ///   z: `list<int32>`
316    ///
317    /// Then the column_indices should be:
318    ///
319    /// - 2.0: [1, 3, 4, 6, 7, 8]
320    /// - 2.1: [0, 2, 4, 5]
321    pub column_indices: Vec<u32>,
322}
323
324impl ReaderProjection {
325    fn from_field_ids_helper<'a>(
326        file_version: LanceFileVersion,
327        fields: impl Iterator<Item = &'a Field>,
328        field_id_to_column_index: &BTreeMap<u32, u32>,
329        column_indices: &mut Vec<u32>,
330    ) -> Result<()> {
331        for field in fields {
332            let is_structural = file_version >= LanceFileVersion::V2_1;
333            let (contributes, recurse) = field_column_shape(field, is_structural);
334            // In the 2.0 system we needed ids for intermediate fields.  In 2.1+
335            // we only need ids for leaf fields.
336            if contributes
337                && let Some(column_idx) = field_id_to_column_index.get(&(field.id as u32)).copied()
338            {
339                column_indices.push(column_idx);
340            }
341            if recurse {
342                Self::from_field_ids_helper(
343                    file_version,
344                    field.children.iter(),
345                    field_id_to_column_index,
346                    column_indices,
347                )?;
348            }
349        }
350        Ok(())
351    }
352
353    /// Creates a projection using a mapping from field IDs to column indices
354    ///
355    /// You can obtain such a mapping when the file is written using the
356    /// [`crate::writer::FileWriter::field_id_to_column_indices`] method.
357    pub fn from_field_ids(
358        file_version: LanceFileVersion,
359        schema: &Schema,
360        field_id_to_column_index: &BTreeMap<u32, u32>,
361    ) -> Result<Self> {
362        let mut column_indices = Vec::new();
363        Self::from_field_ids_helper(
364            file_version,
365            schema.fields.iter(),
366            field_id_to_column_index,
367            &mut column_indices,
368        )?;
369        let projection = Self {
370            schema: Arc::new(schema.clone()),
371            column_indices,
372        };
373        Ok(projection)
374    }
375
376    /// Creates a projection that reads the entire file
377    ///
378    /// If the schema provided is not the schema of the entire file then
379    /// the projection will be invalid and the read will fail.
380    /// If the field is a `struct datatype` with `packed` set to true in the field metadata,
381    /// the whole struct has one column index.
382    /// To support nested `packed-struct encoding`, this method need to be further adjusted.
383    pub fn from_whole_schema(schema: &Schema, version: LanceFileVersion) -> Self {
384        let schema = Arc::new(schema.clone());
385        let is_structural = version >= LanceFileVersion::V2_1;
386        let mut column_indices = vec![];
387        let mut curr_column_idx = 0;
388        let mut packed_struct_fields_num = 0;
389        for field in schema.fields_pre_order() {
390            if packed_struct_fields_num > 0 {
391                packed_struct_fields_num -= 1;
392                continue;
393            }
394            if field.is_packed_struct() {
395                column_indices.push(curr_column_idx);
396                curr_column_idx += 1;
397                packed_struct_fields_num = field.children.len();
398            } else if field.children.is_empty() || !is_structural {
399                column_indices.push(curr_column_idx);
400                curr_column_idx += 1;
401            }
402        }
403        Self {
404            schema,
405            column_indices,
406        }
407    }
408
409    /// Creates a projection that reads the specified columns provided by name
410    ///
411    /// The syntax for column names is the same as [`lance_core::datatypes::Schema::project`]
412    ///
413    /// If the schema provided is not the schema of the entire file then
414    /// the projection will be invalid and the read will fail.
415    pub fn from_column_names(
416        file_version: LanceFileVersion,
417        schema: &Schema,
418        column_names: &[&str],
419    ) -> Result<Self> {
420        let field_id_to_column_index = schema
421            .fields_pre_order()
422            // In the 2.0 system we needed ids for intermediate fields.  In 2.1+
423            // we only need ids for leaf fields.
424            .filter(|field| {
425                file_version < LanceFileVersion::V2_1 || field.is_leaf() || field.is_packed_struct()
426            })
427            .enumerate()
428            .map(|(idx, field)| (field.id as u32, idx as u32))
429            .collect::<BTreeMap<_, _>>();
430        let projected = schema.project(column_names)?;
431        let mut column_indices = Vec::new();
432        Self::from_field_ids_helper(
433            file_version,
434            projected.fields.iter(),
435            &field_id_to_column_index,
436            &mut column_indices,
437        )?;
438        Ok(Self {
439            schema: Arc::new(projected),
440            column_indices,
441        })
442    }
443}
444
445/// File Reader Options that can control reading behaviors, such as whether to enable caching on repetition indices
446#[derive(Clone, Debug)]
447pub struct FileReaderOptions {
448    pub decoder_config: DecoderConfig,
449    /// Size of chunks when reading large pages. Pages larger than this
450    /// will be read in multiple chunks to control memory usage.
451    /// Default: 8MB (DEFAULT_READ_CHUNK_SIZE)
452    pub read_chunk_size: u64,
453    /// If set, the reader will produce batches whose total size in bytes
454    /// is approximately this value, overriding the row-based `batch_size`.
455    ///
456    /// This can be set at the dataset level (via `ReadParams::file_reader_options`)
457    /// to provide a default for all scans, or at the scanner level (via
458    /// `Scanner::batch_size_bytes`) to override per scan.
459    pub batch_size_bytes: Option<u64>,
460}
461
462impl Default for FileReaderOptions {
463    fn default() -> Self {
464        Self {
465            decoder_config: DecoderConfig::default(),
466            read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
467            batch_size_bytes: None,
468        }
469    }
470}
471
472#[derive(Debug, Clone)]
473struct PreparedProjection {
474    column_infos: Vec<Arc<ColumnInfo>>,
475    decoder_projection: ReaderProjection,
476}
477
478#[derive(Debug, Clone)]
479enum FileMetadataProvider {
480    Full(Arc<CachedFileMetadata>),
481    Indexed(Arc<FileMetadataIndex>),
482}
483
484#[derive(Debug, Clone)]
485struct FileReadCore {
486    scheduler: Arc<dyn EncodingsIo>,
487    base_projection: ReaderProjection,
488    metadata_provider: FileMetadataProvider,
489    decoder_plugins: Arc<DecoderPlugins>,
490    cache: Arc<LanceCache>,
491    options: FileReaderOptions,
492}
493
494/// A projection-scoped reader for Lance files.
495///
496/// This reader fixes a base projection at construction time. All later reads
497/// must stay within that projection, which lets the reader load only the column
498/// metadata needed by the base projection when opening from a [`FileMetadataIndex`].
499/// It intentionally does not expose APIs that require synchronous access to full
500/// file metadata.
501#[derive(Debug, Clone)]
502pub struct ProjectedFileReader {
503    core: FileReadCore,
504}
505
506/// A Lance file reader backed by fully decoded file metadata.
507#[derive(Debug, Clone)]
508pub struct FileReader {
509    core: FileReadCore,
510    metadata: Arc<CachedFileMetadata>,
511}
512#[derive(Debug)]
513struct Footer {
514    #[allow(dead_code)]
515    column_meta_start: u64,
516    // We don't use this today because we always load metadata for every column
517    // and don't yet support "metadata projection"
518    #[allow(dead_code)]
519    column_meta_offsets_start: u64,
520    global_buff_offsets_start: u64,
521    num_global_buffers: u32,
522    num_columns: u32,
523    major_version: u16,
524    minor_version: u16,
525}
526
527const FOOTER_LEN: usize = 40;
528
529// How a field maps onto physical columns, shared by the projection-building and
530// projection-validation walks so they stay in lockstep. In the 2.0 layout every
531// ordinary field (including structs and lists) has its own column; in 2.1 only
532// leaves do. Blob/packed-struct fields are opaque in all versions: they are a
533// single column with no descent, including unloaded blob descriptor schemas.
534// Returns `(contributes, recurse)`: whether the field has its own column and
535// whether to walk into its children. The DFS order is the field's own column (if
536// any) followed by its children, so a field's root (first) column is always the
537// first entry of its sub-slice.
538fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) {
539    if field.is_blob() || field.is_packed_struct() {
540        return (true, false);
541    }
542
543    let contributes = !is_structural || field.children.is_empty();
544    let recurse = !field.children.is_empty();
545    (contributes, recurse)
546}
547
548// Count the V2.1 physical columns required to reconstruct a projected field.
549// This is the same DFS shape consumed by `ColumnInfoIter`: ordinary structural
550// nodes are transparent and leaves contribute columns. Indexed metadata loading
551// can therefore compact any ordinary structural projection into 0..N while
552// preserving this order.
553//
554// Blob and packed-struct fields remain unsupported by indexed projection. Their
555// opaque decode semantics are handled by the existing full-metadata reader.
556fn indexed_projection_column_count(field: &Field) -> Option<usize> {
557    if field.is_blob() || field.is_packed_struct() {
558        return None;
559    }
560
561    let (contributes, recurse) = field_column_shape(field, true);
562    let initial = usize::from(contributes);
563    if !recurse {
564        return Some(initial);
565    }
566
567    field.children.iter().try_fold(initial, |count, child| {
568        count.checked_add(indexed_projection_column_count(child)?)
569    })
570}
571
572// Whether a field's children each cover the same rows as the field itself. Struct
573// children do (one value per parent row), so they must share its length. List,
574// map, and fixed-size-list items have an independent cardinality (item count, not
575// row count) and are validated only against themselves.
576fn children_share_parent_length(field: &Field) -> bool {
577    field.logical_type.is_struct()
578}
579
580// Validate one field's slice of a projection's flat `column_indices`, returning
581// the field's top-level row count (the page-row sum of its root column). Walks the
582// same DFS order as `from_field_ids_helper`, advancing `cursor` past every column
583// the field contributes.
584//
585// `comparable` tracks whether the field's row count shares the read's top-level
586// cardinality. A struct's children must all match that count -- the decoders
587// combine them assuming equal lengths and would otherwise panic or read past a
588// shorter child -- so the equality check runs only while `comparable` holds. Once
589// the walk descends through a list/map/fixed-size-list its items have an
590// independent cardinality (item count, not row count), so `comparable` turns off
591// for that whole subtree and a nested struct's children are no longer compared.
592fn validate_field_length<F: Fn(usize) -> Result<u64>>(
593    field: &Field,
594    is_structural: bool,
595    comparable: bool,
596    column_indices: &[u32],
597    cursor: &mut usize,
598    column_len: &F,
599) -> Result<u64> {
600    let (contributes, recurse) = field_column_shape(field, is_structural);
601    let mut field_rows: Option<u64> = None;
602    if contributes {
603        let column = *column_indices.get(*cursor).ok_or_else(|| {
604            Error::invalid_input(format!(
605                "projection supplied fewer column indices than its fields require \
606                 (ran out at field '{}')",
607                field.name
608            ))
609        })?;
610        *cursor += 1;
611        field_rows = Some(column_len(column as usize)?);
612    }
613    if recurse {
614        // Only enforce equal-length children for a struct whose own count is still
615        // at the top-level cardinality; below a list/map/fixed-size-list the items
616        // have an independent cardinality, so neither this field nor its
617        // descendants are comparable.
618        let enforce_children = comparable && children_share_parent_length(field);
619        for child in &field.children {
620            let child_rows = validate_field_length(
621                child,
622                is_structural,
623                enforce_children,
624                column_indices,
625                cursor,
626                column_len,
627            )?;
628            // A struct that contributes no column of its own (the 2.1 layout)
629            // takes its row count from its first child.
630            let expected = *field_rows.get_or_insert(child_rows);
631            if enforce_children && child_rows != expected {
632                return Err(Error::invalid_input(format!(
633                    "cannot read field '{}': its children have differing lengths \
634                     (child '{}' has {} rows, but the field has {}); a struct's \
635                     children must all have the same length",
636                    field.name, child.name, child_rows, expected
637                )));
638            }
639        }
640    }
641    field_rows.ok_or_else(|| {
642        Error::invalid_input(format!(
643            "projected field '{}' maps to no columns",
644            field.name
645        ))
646    })
647}
648
649// The reader combines a projection's columns into rectangular batches, so they
650// must all have the same length.  Returns that common length, or a descriptive
651// error (naming each column's length) when they differ. Ordinary files always
652// pass; only files written with `FileWriter::write_column` whose columns ended up
653// unequal can fail, and those must be read separately.
654fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result<u64> {
655    let first = field_lengths.first().map_or(0, |&(_, len)| len);
656    if field_lengths.iter().all(|&(_, len)| len == first) {
657        return Ok(first);
658    }
659    let columns = field_lengths
660        .iter()
661        .map(|(name, len)| format!("{name}={len}"))
662        .collect::<Vec<_>>()
663        .join(", ");
664    Err(Error::invalid_input(format!(
665        "cannot read columns of differing lengths together ({columns}); \
666         read each column (or equal-length group) separately"
667    )))
668}
669
670impl FileReader {
671    pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
672        Self {
673            core: self.core.with_scheduler(scheduler),
674            metadata: self.metadata.clone(),
675        }
676    }
677
678    /// Returns a clone of this reader whose I/O is additionally recorded into
679    /// `stats`, on top of the scheduler's global accounting.
680    ///
681    /// All cached metadata is shared with `self`, so no file is re-opened and
682    /// only a few `Arc` clones are performed.  If the underlying I/O service
683    /// does not support per-scope statistics (e.g. an in-memory scheduler), the
684    /// returned reader is an ordinary, uninstrumented clone.
685    pub fn with_io_stats(
686        &self,
687        stats: Arc<dyn lance_core::utils::io_stats::IoStatsRecorder>,
688    ) -> Self {
689        match self.core.scheduler.with_io_stats(stats) {
690            Some(scheduler) => self.with_scheduler(scheduler),
691            None => self.clone(),
692        }
693    }
694
695    pub fn num_rows(&self) -> u64 {
696        self.core.num_rows()
697    }
698
699    /// The number of rows stored in a single physical column.
700    ///
701    /// For ordinary (rectangular) files every column has the same length, equal
702    /// to [`num_rows`](Self::num_rows). Files written with
703    /// [`FileWriter::write_column`](crate::writer::FileWriter::write_column)
704    /// may have columns of differing lengths; this returns the length of one
705    /// such column, derived by summing its pages' row counts. Errors if
706    /// `column_index` is out of bounds.
707    pub fn column_num_rows(&self, column_index: usize) -> Result<u64> {
708        let column = self
709            .metadata
710            .column_metadatas
711            .get(column_index)
712            .ok_or_else(|| {
713                Error::invalid_input(format!(
714                    "column index {} is out of bounds (file has {} columns)",
715                    column_index,
716                    self.metadata.column_metadatas.len()
717                ))
718            })?;
719        Ok(column.pages.iter().map(|page| page.length).sum())
720    }
721
722    pub fn metadata(&self) -> &Arc<CachedFileMetadata> {
723        &self.metadata
724    }
725
726    fn statistics_from_column_metadata(
727        column_metadatas: &[pbfile::ColumnMetadata],
728    ) -> FileStatistics {
729        let column_stats = column_metadatas
730            .iter()
731            .map(|col_metadata| {
732                let num_pages = col_metadata.pages.len();
733                let size_bytes = col_metadata
734                    .pages
735                    .iter()
736                    .map(|page| page.buffer_sizes.iter().sum::<u64>())
737                    .sum::<u64>();
738                ColumnStatistics {
739                    num_pages,
740                    size_bytes,
741                }
742            })
743            .collect();
744
745        FileStatistics {
746            columns: column_stats,
747        }
748    }
749
750    pub fn file_statistics(&self) -> FileStatistics {
751        Self::statistics_from_column_metadata(&self.metadata().column_metadatas)
752    }
753
754    pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
755        self.core.read_global_buffer(index).await
756    }
757
758    async fn read_tail(scheduler: &FileScheduler) -> Result<(Bytes, u64)> {
759        let file_size = scheduler.reader().size().await? as u64;
760        let begin = if file_size < scheduler.reader().block_size() as u64 {
761            0
762        } else {
763            file_size - scheduler.reader().block_size() as u64
764        };
765        let tail_bytes = scheduler.submit_single(begin..file_size, 0).await?;
766        Ok((tail_bytes, file_size))
767    }
768
769    async fn read_range_from_tail_or_scheduler(
770        tail_bytes: &Bytes,
771        tail_offset: u64,
772        scheduler: &FileScheduler,
773        range: Range<u64>,
774    ) -> Result<Bytes> {
775        let tail_end = tail_offset + tail_bytes.len() as u64;
776        if range.start >= tail_offset && range.end <= tail_end {
777            let rel_start = (range.start - tail_offset) as usize;
778            let rel_end = (range.end - tail_offset) as usize;
779            Ok(tail_bytes.slice(rel_start..rel_end))
780        } else {
781            scheduler.submit_single(range, 0).await
782        }
783    }
784
785    fn retained_global_buffers_from_tail(
786        gbo_table: &[BufferDescriptor],
787        tail_bytes: &Bytes,
788        tail_offset: u64,
789    ) -> BTreeMap<u32, Bytes> {
790        let tail_end = tail_offset + tail_bytes.len() as u64;
791        gbo_table
792            .iter()
793            .enumerate()
794            .skip(1)
795            .filter_map(|(index, buffer)| {
796                let start = buffer.position;
797                let end = buffer.position + buffer.size;
798                if start >= tail_offset && end <= tail_end {
799                    let rel_start = (start - tail_offset) as usize;
800                    let rel_end = (end - tail_offset) as usize;
801                    let bytes = Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]);
802                    Some((index as u32, bytes))
803                } else {
804                    None
805                }
806            })
807            .collect()
808    }
809
810    // Checks to make sure the footer is written correctly and returns the
811    // position of the file descriptor (which comes from the footer)
812    fn decode_footer(footer_bytes: &Bytes) -> Result<Footer> {
813        let len = footer_bytes.len();
814        if len < FOOTER_LEN {
815            return Err(Error::invalid_input(format!(
816                "does not have sufficient data, len: {}, bytes: {:?}",
817                len, footer_bytes
818            )));
819        }
820        let mut cursor = Cursor::new(footer_bytes.slice(len - FOOTER_LEN..));
821
822        let column_meta_start = cursor.read_u64::<LittleEndian>()?;
823        let column_meta_offsets_start = cursor.read_u64::<LittleEndian>()?;
824        let global_buff_offsets_start = cursor.read_u64::<LittleEndian>()?;
825        let num_global_buffers = cursor.read_u32::<LittleEndian>()?;
826        let num_columns = cursor.read_u32::<LittleEndian>()?;
827        let major_version = cursor.read_u16::<LittleEndian>()?;
828        let minor_version = cursor.read_u16::<LittleEndian>()?;
829
830        if major_version == MAJOR_VERSION as u16 && minor_version == MINOR_VERSION as u16 {
831            return Err(Error::version_conflict(
832                "Attempt to use the lance v2 reader to read a legacy file".to_string(),
833                major_version,
834                minor_version,
835            ));
836        }
837
838        let magic_bytes = footer_bytes.slice(len - 4..);
839        if magic_bytes.as_ref() != MAGIC {
840            return Err(Error::invalid_input(format!(
841                "file does not appear to be a Lance file (invalid magic: {:?})",
842                MAGIC
843            )));
844        }
845        Ok(Footer {
846            column_meta_start,
847            column_meta_offsets_start,
848            global_buff_offsets_start,
849            num_global_buffers,
850            num_columns,
851            major_version,
852            minor_version,
853        })
854    }
855
856    // TODO: Once we have coalesced I/O we should only read the column metadatas that we need
857    fn read_all_column_metadata(
858        column_metadata_bytes: Bytes,
859        footer: &Footer,
860    ) -> Result<Vec<pbfile::ColumnMetadata>> {
861        let column_metadata_start = footer.column_meta_start;
862        // cmo == column_metadata_offsets
863        let cmo_table_size = 16 * footer.num_columns as usize;
864        if column_metadata_bytes.len() < cmo_table_size {
865            return Err(Error::invalid_input(format!(
866                "column metadata region has {} bytes but CMO table needs {} bytes for {} columns",
867                column_metadata_bytes.len(),
868                cmo_table_size,
869                footer.num_columns
870            )));
871        }
872        let cmo_table = column_metadata_bytes.slice(column_metadata_bytes.len() - cmo_table_size..);
873        let column_metadata_offsets = Self::decode_cmo_table(cmo_table, footer)?;
874
875        column_metadata_offsets
876            .iter()
877            .map(|(position, length)| {
878                let normalized_position = (*position - column_metadata_start) as usize;
879                let normalized_end = normalized_position + (*length as usize);
880                Ok(pbfile::ColumnMetadata::decode(
881                    &column_metadata_bytes[normalized_position..normalized_end],
882                )?)
883            })
884            .collect::<Result<Vec<_>>>()
885    }
886
887    fn decode_cmo_table(cmo_table: Bytes, footer: &Footer) -> Result<Arc<[(u64, u64)]>> {
888        let expected_size = 16 * footer.num_columns as usize;
889        if cmo_table.len() != expected_size {
890            return Err(Error::invalid_input(format!(
891                "column metadata offset table has {} bytes but expected {} bytes for {} columns",
892                cmo_table.len(),
893                expected_size,
894                footer.num_columns
895            )));
896        }
897
898        let mut offsets = Vec::with_capacity(footer.num_columns as usize);
899        for col_idx in 0..footer.num_columns {
900            let offset = (col_idx * 16) as usize;
901            let position = LittleEndian::read_u64(&cmo_table[offset..offset + 8]);
902            let length = LittleEndian::read_u64(&cmo_table[offset + 8..offset + 16]);
903            let end = position.checked_add(length).ok_or_else(|| {
904                Error::invalid_input(format!(
905                    "column metadata range overflows for column index {}, position={}, length={}",
906                    col_idx, position, length
907                ))
908            })?;
909            if position < footer.column_meta_start || end > footer.column_meta_offsets_start {
910                return Err(Error::invalid_input(format!(
911                    "column metadata range for column index {} is outside metadata region: position={}, length={}, metadata_start={}, cmo_start={}",
912                    col_idx,
913                    position,
914                    length,
915                    footer.column_meta_start,
916                    footer.column_meta_offsets_start
917                )));
918            }
919            offsets.push((position, length));
920        }
921
922        Ok(Arc::from(offsets))
923    }
924
925    async fn optimistic_tail_read(
926        data: &Bytes,
927        start_pos: u64,
928        scheduler: &FileScheduler,
929        file_len: u64,
930    ) -> Result<Bytes> {
931        let num_bytes_needed = (file_len - start_pos) as usize;
932        if data.len() >= num_bytes_needed {
933            Ok(data.slice((data.len() - num_bytes_needed)..))
934        } else {
935            let num_bytes_missing = (num_bytes_needed - data.len()) as u64;
936            let start = file_len - num_bytes_needed as u64;
937            let missing_bytes = scheduler
938                .submit_single(start..start + num_bytes_missing, 0)
939                .await?;
940            let mut combined = BytesMut::with_capacity(data.len() + num_bytes_missing as usize);
941            combined.extend(missing_bytes);
942            combined.extend(data);
943            Ok(combined.freeze())
944        }
945    }
946
947    fn do_decode_gbo_table(
948        gbo_bytes: &Bytes,
949        footer: &Footer,
950        version: LanceFileVersion,
951    ) -> Result<Vec<BufferDescriptor>> {
952        let mut global_bufs_cursor = Cursor::new(gbo_bytes);
953
954        let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize);
955        for buffer_index in 0..footer.num_global_buffers {
956            let buf_pos = global_bufs_cursor.read_u64::<LittleEndian>()?;
957            if version >= LanceFileVersion::V2_1 && buf_pos % PAGE_BUFFER_ALIGNMENT as u64 != 0 {
958                return Err(Error::invalid_input_source(
959                    format!(
960                        "Global buffer {} position {} is not aligned to {} bytes",
961                        buffer_index, buf_pos, PAGE_BUFFER_ALIGNMENT
962                    )
963                    .into(),
964                ));
965            }
966            let buf_size = global_bufs_cursor.read_u64::<LittleEndian>()?;
967            global_buffers.push(BufferDescriptor {
968                position: buf_pos,
969                size: buf_size,
970            });
971        }
972
973        Ok(global_buffers)
974    }
975
976    async fn decode_gbo_table(
977        tail_bytes: &Bytes,
978        file_len: u64,
979        scheduler: &FileScheduler,
980        footer: &Footer,
981        version: LanceFileVersion,
982    ) -> Result<Vec<BufferDescriptor>> {
983        // This could, in theory, trigger another IOP but the GBO table should never be large
984        // enough for that to happen
985        let gbo_bytes = Self::optimistic_tail_read(
986            tail_bytes,
987            footer.global_buff_offsets_start,
988            scheduler,
989            file_len,
990        )
991        .await?;
992        Self::do_decode_gbo_table(&gbo_bytes, footer, version)
993    }
994
995    fn decode_schema(schema_bytes: Bytes) -> Result<(u64, lance_core::datatypes::Schema)> {
996        let file_descriptor = pb::FileDescriptor::decode(schema_bytes)?;
997        let pb_schema = file_descriptor.schema.unwrap();
998        let num_rows = file_descriptor.length;
999        let fields_with_meta = FieldsWithMeta {
1000            fields: Fields(pb_schema.fields),
1001            metadata: pb_schema.metadata,
1002        };
1003        let schema = Schema::try_from(fields_with_meta)?;
1004        Ok((num_rows, schema))
1005    }
1006
1007    // TODO: Support late projection.  Currently, if we want to perform a
1008    // projected read of a file, we load all of the column metadata, and then
1009    // only read the column data that is requested.  This is fine for most cases.
1010    //
1011    // However, if there are many columns then loading all of the column metadata
1012    // may be expensive.  We should support a mode where we only load the column
1013    // metadata for the columns that are requested (the file format supports this).
1014    //
1015    // The main challenge is that we either need to ignore the column metadata cache
1016    // or have a more sophisticated cache that can cache per-column metadata.
1017    //
1018    // Also, if the number of columns is fairly small, it's faster to read them as a
1019    // single IOP, but we can fix this through coalescing.
1020    pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result<CachedFileMetadata> {
1021        // 1. read the footer
1022        let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1023        let tail_offset = file_len - tail_bytes.len() as u64;
1024        let footer = Self::decode_footer(&tail_bytes)?;
1025
1026        let file_version: LanceFileVersion =
1027            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
1028                .into();
1029
1030        let gbo_table =
1031            Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1032        if gbo_table.is_empty() {
1033            return Err(Error::internal(
1034                "File did not contain any global buffers, schema expected".to_string(),
1035            ));
1036        }
1037        let schema_start = gbo_table[0].position;
1038        let schema_size = gbo_table[0].size;
1039
1040        let num_footer_bytes = file_len - schema_start;
1041
1042        // By default we read all column metadatas.  We do NOT read the column metadata buffers
1043        // at this point.  We only want to read the column metadata for columns we are actually loading.
1044        let all_metadata_bytes =
1045            Self::optimistic_tail_read(&tail_bytes, schema_start, scheduler, file_len).await?;
1046
1047        let schema_bytes = all_metadata_bytes.slice(0..schema_size as usize);
1048        let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1049
1050        // Next, read the metadata for the columns
1051        // This is both the column metadata and the CMO table
1052        let column_metadata_start = (footer.column_meta_start - schema_start) as usize;
1053        let column_metadata_end = (footer.global_buff_offsets_start - schema_start) as usize;
1054        let column_metadata_bytes =
1055            all_metadata_bytes.slice(column_metadata_start..column_metadata_end);
1056        let column_metadatas = Self::read_all_column_metadata(column_metadata_bytes, &footer)?;
1057
1058        let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::<u64>();
1059        let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes;
1060        let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start;
1061
1062        let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version)?;
1063
1064        // The tail read above already pulled in any global buffer that lives within
1065        // the captured window. Copy those user buffers (index >= 1; the schema at 0
1066        // is decoded above and never fetched via read_global_buffer) out of the tail
1067        // so read_global_buffer can serve them without I/O. We copy rather than slice
1068        // so the much larger tail allocation can be released once decoding is done.
1069        let retained_global_buffers =
1070            Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1071
1072        Ok(CachedFileMetadata {
1073            file_schema: Arc::new(schema),
1074            column_metadatas,
1075            column_infos,
1076            num_rows,
1077            num_data_bytes,
1078            num_column_metadata_bytes,
1079            num_global_buffer_bytes,
1080            num_footer_bytes,
1081            file_buffers: gbo_table,
1082            major_version: footer.major_version,
1083            minor_version: footer.minor_version,
1084            file_size_bytes: file_len,
1085            retained_global_buffers,
1086        })
1087    }
1088
1089    async fn read_metadata_index_with_known_schema(
1090        scheduler: &FileScheduler,
1091        known_schema: Option<(Arc<Schema>, u64)>,
1092    ) -> Result<FileMetadataIndex> {
1093        let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1094        let tail_offset = file_len - tail_bytes.len() as u64;
1095        let footer = Self::decode_footer(&tail_bytes)?;
1096
1097        let file_version: LanceFileVersion =
1098            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
1099                .into();
1100
1101        let gbo_table =
1102            Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1103        if gbo_table.is_empty() {
1104            return Err(Error::internal(
1105                "File did not contain any global buffers, schema expected".to_string(),
1106            ));
1107        }
1108        let (file_schema, num_rows) = match known_schema {
1109            Some((file_schema, num_rows)) => (file_schema, num_rows),
1110            None => {
1111                let schema_buffer = &gbo_table[0];
1112                let schema_bytes = Self::read_range_from_tail_or_scheduler(
1113                    &tail_bytes,
1114                    tail_offset,
1115                    scheduler,
1116                    schema_buffer.position..schema_buffer.position + schema_buffer.size,
1117                )
1118                .await?;
1119                let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1120                (Arc::new(schema), num_rows)
1121            }
1122        };
1123
1124        let cmo_table = Self::read_range_from_tail_or_scheduler(
1125            &tail_bytes,
1126            tail_offset,
1127            scheduler,
1128            footer.column_meta_offsets_start..footer.global_buff_offsets_start,
1129        )
1130        .await?;
1131        let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?;
1132
1133        let retained_global_buffers =
1134            Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1135
1136        Ok(FileMetadataIndex {
1137            file_schema,
1138            num_rows,
1139            file_buffers: gbo_table,
1140            column_metadata_offsets,
1141            num_columns: footer.num_columns,
1142            version: file_version,
1143            file_size_bytes: file_len,
1144            retained_global_buffers,
1145        })
1146    }
1147
1148    /// Reads the lightweight metadata index from a file.
1149    ///
1150    /// This reads the file schema from the schema global buffer. Use
1151    /// [`Self::read_metadata_index_with_schema`] when the caller already has
1152    /// the schema and row count from a higher-level metadata source.
1153    pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result<FileMetadataIndex> {
1154        Self::read_metadata_index_with_known_schema(scheduler, None).await
1155    }
1156
1157    /// Reads the metadata index without fetching the schema global buffer.
1158    ///
1159    /// Use this when the caller already has the file schema and physical row
1160    /// count from an enclosing metadata layer, such as a dataset manifest.
1161    pub async fn read_metadata_index_with_schema(
1162        scheduler: &FileScheduler,
1163        file_schema: Arc<Schema>,
1164        num_rows: u64,
1165    ) -> Result<FileMetadataIndex> {
1166        Self::read_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))).await
1167    }
1168
1169    fn fetch_encoding<M: Default + Name + Sized>(encoding: &pbfile::Encoding) -> Result<M> {
1170        match &encoding.location {
1171            Some(pbfile::encoding::Location::Indirect(_)) => Err(Error::invalid_input_source(
1172                "Indirect file encodings are not supported".into(),
1173            )),
1174            Some(pbfile::encoding::Location::Direct(encoding)) => {
1175                let encoding_buf = Bytes::from(encoding.encoding.clone());
1176                let encoding_any = prost_types::Any::decode(encoding_buf).map_err(|error| {
1177                    Error::invalid_input_source(
1178                        format!("Invalid direct {} encoding envelope: {error}", M::NAME).into(),
1179                    )
1180                })?;
1181                encoding_any.to_msg::<M>().map_err(|error| {
1182                    Error::invalid_input_source(
1183                        format!("Invalid direct {} encoding: {error}", M::NAME).into(),
1184                    )
1185                })
1186            }
1187            Some(pbfile::encoding::Location::None(_)) => Err(Error::invalid_input_source(
1188                format!("Missing {} encoding description", M::NAME).into(),
1189            )),
1190            None => Err(Error::invalid_input_source(
1191                format!("Missing {} encoding location", M::NAME).into(),
1192            )),
1193        }
1194    }
1195
1196    fn meta_to_col_infos(
1197        column_metadatas: &[pbfile::ColumnMetadata],
1198        file_version: LanceFileVersion,
1199    ) -> Result<Vec<Arc<ColumnInfo>>> {
1200        column_metadatas
1201            .iter()
1202            .enumerate()
1203            .map(|(col_idx, col_meta)| {
1204                let col_idx = u32::try_from(col_idx).map_err(|_| {
1205                    Error::invalid_input_source("File has more than u32::MAX columns".into())
1206                })?;
1207                Self::meta_to_col_info(col_idx, col_meta, file_version)
1208            })
1209            .collect()
1210    }
1211
1212    fn meta_to_col_info(
1213        col_idx: u32,
1214        col_meta: &pbfile::ColumnMetadata,
1215        file_version: LanceFileVersion,
1216    ) -> Result<Arc<ColumnInfo>> {
1217        let page_infos = col_meta
1218            .pages
1219            .iter()
1220            .enumerate()
1221            .map(|(page_idx, page)| {
1222                let num_rows = page.length;
1223                let encoding = match file_version {
1224                    LanceFileVersion::V2_0 => {
1225                        PageEncoding::Legacy(Self::fetch_encoding::<pbenc::ArrayEncoding>(
1226                            page.encoding.as_ref().ok_or_else(|| {
1227                                Error::invalid_input_source(
1228                                    format!(
1229                                        "Column {} page {} is missing its encoding",
1230                                        col_idx, page_idx
1231                                    )
1232                                    .into(),
1233                                )
1234                            })?,
1235                        )?)
1236                    }
1237                    _ => {
1238                        PageEncoding::Structural(Self::fetch_encoding::<pbenc21::PageLayout>(
1239                            page.encoding.as_ref().ok_or_else(|| {
1240                                Error::invalid_input_source(
1241                                    format!(
1242                                        "Column {} page {} is missing its encoding",
1243                                        col_idx, page_idx
1244                                    )
1245                                    .into(),
1246                                )
1247                            })?,
1248                        )?)
1249                    }
1250                };
1251                if page.buffer_offsets.len() != page.buffer_sizes.len() {
1252                    return Err(Error::invalid_input_source(
1253                        format!(
1254                            "Column {} page {} has {} buffer offsets but {} buffer sizes",
1255                            col_idx,
1256                            page_idx,
1257                            page.buffer_offsets.len(),
1258                            page.buffer_sizes.len()
1259                        )
1260                        .into(),
1261                    ));
1262                }
1263                let buffer_offsets_and_sizes = Arc::from(
1264                    page.buffer_offsets
1265                        .iter()
1266                        .zip(page.buffer_sizes.iter())
1267                        .map(|(offset, size)| -> Result<_> {
1268                            if file_version >= LanceFileVersion::V2_1
1269                                && offset % PAGE_BUFFER_ALIGNMENT as u64 != 0
1270                            {
1271                                return Err(Error::invalid_input_source(
1272                                    format!(
1273                                        "Column {} page {} buffer offset {} is not aligned to {} bytes",
1274                                        col_idx, page_idx, offset, PAGE_BUFFER_ALIGNMENT
1275                                    )
1276                                    .into(),
1277                                ));
1278                            }
1279                            Ok((*offset, *size))
1280                        })
1281                        .collect::<Result<Vec<_>>>()?,
1282                );
1283                Ok(PageInfo {
1284                    buffer_offsets_and_sizes,
1285                    encoding,
1286                    num_rows,
1287                    priority: page.priority,
1288                })
1289            })
1290            .collect::<Result<Vec<_>>>()?;
1291        if col_meta.buffer_offsets.len() != col_meta.buffer_sizes.len() {
1292            return Err(Error::invalid_input_source(
1293                format!(
1294                    "Column {} has {} buffer offsets but {} buffer sizes",
1295                    col_idx,
1296                    col_meta.buffer_offsets.len(),
1297                    col_meta.buffer_sizes.len()
1298                )
1299                .into(),
1300            ));
1301        }
1302        let buffer_offsets_and_sizes = Arc::from(
1303            col_meta
1304                .buffer_offsets
1305                .iter()
1306                .zip(col_meta.buffer_sizes.iter())
1307                .map(|(offset, size)| (*offset, *size))
1308                .collect::<Vec<_>>(),
1309        );
1310        Ok(Arc::new(ColumnInfo {
1311            index: col_idx,
1312            page_infos: Arc::from(page_infos),
1313            buffer_offsets_and_sizes,
1314            encoding: Self::fetch_encoding(col_meta.encoding.as_ref().ok_or_else(|| {
1315                Error::invalid_input_source(
1316                    format!("Column {} is missing its encoding", col_idx).into(),
1317                )
1318            })?)?,
1319        }))
1320    }
1321
1322    fn validate_projection(
1323        projection: &ReaderProjection,
1324        metadata: &CachedFileMetadata,
1325    ) -> Result<()> {
1326        if projection.schema.fields.is_empty() {
1327            return Err(Error::invalid_input(
1328                "Attempt to read zero columns from the file, at least one column must be specified"
1329                    .to_string(),
1330            ));
1331        }
1332        let mut column_indices_seen = BTreeSet::new();
1333        for column_index in &projection.column_indices {
1334            if !column_indices_seen.insert(*column_index) {
1335                return Err(Error::invalid_input(format!(
1336                    "The projection specified the column index {} more than once",
1337                    column_index
1338                )));
1339            }
1340            if *column_index >= metadata.column_infos.len() as u32 {
1341                return Err(Error::invalid_input(format!(
1342                    "The projection specified the column index {} but there are only {} columns in the file",
1343                    column_index,
1344                    metadata.column_infos.len()
1345                )));
1346            }
1347        }
1348        Ok(())
1349    }
1350
1351    /// Opens a new file reader without any pre-existing knowledge
1352    ///
1353    /// This will read the file schema from the file itself and thus requires a bit more I/O
1354    ///
1355    /// A `base_projection` can also be provided.  If provided, then the projection will apply
1356    /// to all reads from the file that do not specify their own projection.
1357    pub async fn try_open(
1358        scheduler: FileScheduler,
1359        base_projection: Option<ReaderProjection>,
1360        decoder_plugins: Arc<DecoderPlugins>,
1361        cache: &LanceCache,
1362        options: FileReaderOptions,
1363    ) -> Result<Self> {
1364        let file_metadata = Arc::new(Self::read_all_metadata(&scheduler).await?);
1365        let path = scheduler.reader().path().clone();
1366
1367        // Create LanceEncodingsIo with read chunk size from options
1368        let encodings_io =
1369            LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
1370
1371        Self::try_open_with_file_metadata(
1372            Arc::new(encodings_io),
1373            path,
1374            base_projection,
1375            decoder_plugins,
1376            file_metadata,
1377            cache,
1378            options,
1379        )
1380        .await
1381    }
1382
1383    /// Same as `try_open` but with the file metadata already loaded.
1384    ///
1385    /// This method also can accept any kind of `EncodingsIo` implementation allowing
1386    /// for custom strategies to be used for I/O scheduling (e.g. for takes on fast
1387    /// disks it may be better to avoid asynchronous overhead).
1388    /// Opens a data reader backed by fully decoded file metadata.
1389    pub async fn try_open_with_file_metadata(
1390        scheduler: Arc<dyn EncodingsIo>,
1391        path: Path,
1392        base_projection: Option<ReaderProjection>,
1393        decoder_plugins: Arc<DecoderPlugins>,
1394        file_metadata: Arc<CachedFileMetadata>,
1395        cache: &LanceCache,
1396        options: FileReaderOptions,
1397    ) -> Result<Self> {
1398        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
1399        let core = FileReadCore::try_new(
1400            scheduler,
1401            base_projection,
1402            decoder_plugins,
1403            FileMetadataProvider::Full(file_metadata.clone()),
1404            cache,
1405            options,
1406        )?;
1407        Ok(Self {
1408            core,
1409            metadata: file_metadata,
1410        })
1411    }
1412
1413    // The actual decoder needs all the column infos that make up a type.  In other words, if
1414    // the first type in the schema is Struct<i32, i32> then the decoder will need 3 column infos.
1415    //
1416    // This is a file reader concern because the file reader needs to support late projection of columns
1417    // and so it will need to figure this out anyways.
1418    //
1419    // It's a bit of a tricky process though because the number of column infos may depend on the
1420    // encoding.  Considering the above example, if we wrote it with a packed encoding, then there would
1421    // only be a single column in the file (and not 3).
1422    //
1423    // At the moment this method words because our rules are simple and we just repeat them here.  See
1424    // Self::default_projection for a similar problem.  In the future this is something the encodings
1425    // registry will need to figure out.
1426    fn collect_columns_from_projection(
1427        &self,
1428        _projection: &ReaderProjection,
1429    ) -> Result<Vec<Arc<ColumnInfo>>> {
1430        Ok(self.metadata.column_infos.clone())
1431    }
1432
1433    #[allow(clippy::too_many_arguments)]
1434    async fn do_read_range(
1435        column_infos: Vec<Arc<ColumnInfo>>,
1436        io: Arc<dyn EncodingsIo>,
1437        cache: Arc<LanceCache>,
1438        num_rows: u64,
1439        decoder_plugins: Arc<DecoderPlugins>,
1440        range: Range<u64>,
1441        batch_size: u32,
1442        projection: ReaderProjection,
1443        filter: FilterExpression,
1444        decoder_config: DecoderConfig,
1445        batch_size_bytes: Option<u64>,
1446    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1447        debug!(
1448            "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1449            range,
1450            batch_size,
1451            num_rows,
1452            column_infos.len(),
1453            projection.schema.fields.len(),
1454        );
1455
1456        let config = SchedulerDecoderConfig {
1457            batch_size,
1458            cache,
1459            decoder_plugins,
1460            io,
1461            decoder_config,
1462            batch_size_bytes,
1463        };
1464
1465        let requested_rows = RequestedRows::Ranges(vec![range]);
1466
1467        schedule_and_decode(
1468            column_infos,
1469            requested_rows,
1470            filter,
1471            projection.column_indices,
1472            projection.schema,
1473            config,
1474        )
1475        .await
1476    }
1477
1478    #[allow(clippy::too_many_arguments)]
1479    async fn do_take_rows(
1480        column_infos: Vec<Arc<ColumnInfo>>,
1481        io: Arc<dyn EncodingsIo>,
1482        cache: Arc<LanceCache>,
1483        decoder_plugins: Arc<DecoderPlugins>,
1484        indices: Vec<u64>,
1485        batch_size: u32,
1486        projection: ReaderProjection,
1487        filter: FilterExpression,
1488        decoder_config: DecoderConfig,
1489        batch_size_bytes: Option<u64>,
1490    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1491        debug!(
1492            "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1493            indices.len(),
1494            indices[0],
1495            indices[indices.len() - 1],
1496            batch_size,
1497            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1498        );
1499
1500        let config = SchedulerDecoderConfig {
1501            batch_size,
1502            cache,
1503            decoder_plugins,
1504            io,
1505            decoder_config,
1506            batch_size_bytes,
1507        };
1508
1509        let requested_rows = RequestedRows::Indices(indices);
1510
1511        schedule_and_decode(
1512            column_infos,
1513            requested_rows,
1514            filter,
1515            projection.column_indices,
1516            projection.schema,
1517            config,
1518        )
1519        .await
1520    }
1521
1522    #[allow(clippy::too_many_arguments)]
1523    async fn do_read_ranges(
1524        column_infos: Vec<Arc<ColumnInfo>>,
1525        io: Arc<dyn EncodingsIo>,
1526        cache: Arc<LanceCache>,
1527        decoder_plugins: Arc<DecoderPlugins>,
1528        ranges: Vec<Range<u64>>,
1529        batch_size: u32,
1530        projection: ReaderProjection,
1531        filter: FilterExpression,
1532        decoder_config: DecoderConfig,
1533        batch_size_bytes: Option<u64>,
1534    ) -> Result<BoxStream<'static, ReadBatchTask>> {
1535        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1536        debug!(
1537            "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1538            ranges.len(),
1539            num_rows,
1540            ranges[0].start,
1541            ranges[ranges.len() - 1].end,
1542            batch_size,
1543            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1544        );
1545
1546        let config = SchedulerDecoderConfig {
1547            batch_size,
1548            cache,
1549            decoder_plugins,
1550            io,
1551            decoder_config,
1552            batch_size_bytes,
1553        };
1554
1555        let requested_rows = RequestedRows::Ranges(ranges);
1556
1557        schedule_and_decode(
1558            column_infos,
1559            requested_rows,
1560            filter,
1561            projection.column_indices,
1562            projection.schema,
1563            config,
1564        )
1565        .await
1566    }
1567
1568    /// Creates a stream of "read tasks" to read the data from the file
1569    ///
1570    /// The arguments are similar to [`Self::read_stream_projected`] but instead of returning a stream
1571    /// of record batches it returns a stream of "read tasks".
1572    ///
1573    /// The tasks should be consumed with some kind of `buffered` argument if CPU parallelism is desired.
1574    ///
1575    /// Note that "read task" is probably a bit imprecise.  The tasks are actually "decode tasks".  The
1576    /// reading happens asynchronously in the background.  In other words, a single read task may map to
1577    /// multiple I/O operations or a single I/O operation may map to multiple read tasks.
1578    ///
1579    /// # Why is this async?
1580    ///
1581    /// Constructing the read stream requires running the decode scheduler's
1582    /// `initialize` step, which performs the metadata I/O (chunk metadata,
1583    /// dictionaries, repetition index, ...) needed to plan the read.  We
1584    /// drive that I/O on the awaiting task rather than smuggling it into
1585    /// the stream's first poll.  This way callers control where the
1586    /// scheduling I/O runs (typically inside a per-fragment
1587    /// `tokio::spawn`), planning errors surface from the await instead of
1588    /// from the first stream item, and small reads can also complete the
1589    /// synchronous scheduling step before returning (see
1590    /// [`DecoderConfig::inline_scheduling`]).
1591    pub async fn read_tasks(
1592        &self,
1593        params: ReadBatchParams,
1594        batch_size: u32,
1595        projection: Option<ReaderProjection>,
1596        filter: FilterExpression,
1597    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
1598        self.core
1599            .read_tasks(params, batch_size, projection, filter)
1600            .await
1601    }
1602
1603    /// Reads data from the file as a stream of record batches
1604    ///
1605    /// * `params` - Specifies the range (or indices) of data to read
1606    /// * `batch_size` - The maximum size of a single batch.  A batch may be smaller
1607    ///   if it is the last batch or if it is not possible to create a batch of the
1608    ///   requested size.
1609    ///
1610    ///   For example, if the batch size is 1024 and one of the columns is a string
1611    ///   column then there may be some ranges of 1024 rows that contain more than
1612    ///   2^31 bytes of string data (which is the maximum size of a string column
1613    ///   in Arrow).  In this case smaller batches may be emitted.
1614    /// * `batch_readahead` - The number of batches to read ahead.  This controls the
1615    ///   amount of CPU parallelism of the read.  In other words it controls how many
1616    ///   batches will be decoded in parallel.  It has no effect on the I/O parallelism
1617    ///   of the read (how many I/O requests are in flight at once).
1618    ///
1619    ///   This parameter also is also related to backpressure.  If the consumer of the
1620    ///   stream is slow then the reader will build up RAM.
1621    /// * `projection` - A projection to apply to the read.  This controls which columns
1622    ///   are read from the file.  The projection is NOT applied on top of the base
1623    ///   projection.  The projection is applied directly to the file schema.
1624    ///
1625    /// # Why is this async?
1626    ///
1627    /// This delegates to [`Self::read_tasks`], which awaits the decode
1628    /// scheduler's `initialize` step (and, for small reads, the synchronous
1629    /// scheduling that follows) before returning.  See `read_tasks` for
1630    /// details on why this work is performed up front rather than on the
1631    /// stream's first poll.
1632    pub async fn read_stream_projected(
1633        &self,
1634        params: ReadBatchParams,
1635        batch_size: u32,
1636        batch_readahead: u32,
1637        projection: ReaderProjection,
1638        filter: FilterExpression,
1639    ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1640        let arrow_schema = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
1641        let tasks_stream = self
1642            .read_tasks(params, batch_size, Some(projection), filter)
1643            .await?;
1644        let batch_stream = tasks_stream
1645            .map(|task| task.task)
1646            .buffered(batch_readahead as usize)
1647            .boxed();
1648        Ok(Box::pin(RecordBatchStreamAdapter::new(
1649            arrow_schema,
1650            batch_stream,
1651        )))
1652    }
1653
1654    fn take_rows_blocking(
1655        &self,
1656        indices: Vec<u64>,
1657        batch_size: u32,
1658        projection: ReaderProjection,
1659        filter: FilterExpression,
1660    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1661        let column_infos = self.collect_columns_from_projection(&projection)?;
1662        debug!(
1663            "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1664            indices.len(),
1665            indices[0],
1666            indices[indices.len() - 1],
1667            batch_size,
1668            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1669        );
1670
1671        let config = SchedulerDecoderConfig {
1672            batch_size,
1673            cache: self.core.cache.clone(),
1674            decoder_plugins: self.core.decoder_plugins.clone(),
1675            io: self.core.scheduler.clone(),
1676            decoder_config: self.core.options.decoder_config.clone(),
1677            batch_size_bytes: self.core.options.batch_size_bytes,
1678        };
1679
1680        let requested_rows = RequestedRows::Indices(indices);
1681
1682        schedule_and_decode_blocking(
1683            column_infos,
1684            requested_rows,
1685            filter,
1686            projection.column_indices,
1687            projection.schema,
1688            config,
1689        )
1690    }
1691
1692    fn read_ranges_blocking(
1693        &self,
1694        ranges: Vec<Range<u64>>,
1695        batch_size: u32,
1696        projection: ReaderProjection,
1697        filter: FilterExpression,
1698    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1699        let column_infos = self.collect_columns_from_projection(&projection)?;
1700        let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1701        debug!(
1702            "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1703            ranges.len(),
1704            num_rows,
1705            ranges[0].start,
1706            ranges[ranges.len() - 1].end,
1707            batch_size,
1708            column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1709        );
1710
1711        let config = SchedulerDecoderConfig {
1712            batch_size,
1713            cache: self.core.cache.clone(),
1714            decoder_plugins: self.core.decoder_plugins.clone(),
1715            io: self.core.scheduler.clone(),
1716            decoder_config: self.core.options.decoder_config.clone(),
1717            batch_size_bytes: self.core.options.batch_size_bytes,
1718        };
1719
1720        let requested_rows = RequestedRows::Ranges(ranges);
1721
1722        schedule_and_decode_blocking(
1723            column_infos,
1724            requested_rows,
1725            filter,
1726            projection.column_indices,
1727            projection.schema,
1728            config,
1729        )
1730    }
1731
1732    fn read_range_blocking(
1733        &self,
1734        range: Range<u64>,
1735        batch_size: u32,
1736        projection: ReaderProjection,
1737        filter: FilterExpression,
1738    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1739        let column_infos = self.collect_columns_from_projection(&projection)?;
1740        let num_rows = self.core.num_rows();
1741
1742        debug!(
1743            "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1744            range,
1745            batch_size,
1746            num_rows,
1747            column_infos.len(),
1748            projection.schema.fields.len(),
1749        );
1750
1751        let config = SchedulerDecoderConfig {
1752            batch_size,
1753            cache: self.core.cache.clone(),
1754            decoder_plugins: self.core.decoder_plugins.clone(),
1755            io: self.core.scheduler.clone(),
1756            decoder_config: self.core.options.decoder_config.clone(),
1757            batch_size_bytes: self.core.options.batch_size_bytes,
1758        };
1759
1760        let requested_rows = RequestedRows::Ranges(vec![range]);
1761
1762        schedule_and_decode_blocking(
1763            column_infos,
1764            requested_rows,
1765            filter,
1766            projection.column_indices,
1767            projection.schema,
1768            config,
1769        )
1770    }
1771
1772    /// Read data from the file as an iterator of record batches
1773    ///
1774    /// This is a blocking variant of [`Self::read_stream_projected`] that runs entirely in the
1775    /// calling thread.  It will block on I/O if the decode is faster than the I/O.  It is useful
1776    /// for benchmarking and potentially from "take"ing small batches from fast disks.
1777    ///
1778    /// Large scans of in-memory data will still benefit from threading (and should therefore not
1779    /// use this method) because we can parallelize the decode.
1780    ///
1781    /// Note: calling this from within a tokio runtime will panic.  It is acceptable to call this
1782    /// from a spawn_blocking context.
1783    pub fn read_stream_projected_blocking(
1784        &self,
1785        params: ReadBatchParams,
1786        batch_size: u32,
1787        projection: Option<ReaderProjection>,
1788        filter: FilterExpression,
1789    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1790        let projection = projection.unwrap_or_else(|| self.core.base_projection.clone());
1791        Self::validate_projection(&projection, &self.metadata)?;
1792        // Apply the same projection-length validation as the async path.  This
1793        // reader is always backed by full metadata, so we can build the prepared
1794        // projection synchronously (no column-metadata I/O) and reuse the shared
1795        // check.  `read_len` is the projection's common column length, which
1796        // `RangeFull`/`RangeFrom` resolve against rather than `num_rows`.
1797        let prepared = PreparedProjection {
1798            column_infos: self.metadata.column_infos.clone(),
1799            decoder_projection: projection.clone(),
1800        };
1801        let read_len = self.core.prepared_read_length(&prepared)?;
1802        let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
1803            if bound > read_len || (bound == read_len && inclusive) {
1804                Err(Error::invalid_input(format!(
1805                    "cannot read {params:?} from columns with {read_len} rows"
1806                )))
1807            } else {
1808                Ok(())
1809            }
1810        };
1811        match &params {
1812            ReadBatchParams::Indices(indices) => {
1813                for idx in indices {
1814                    match idx {
1815                        None => {
1816                            return Err(Error::invalid_input("Null value in indices array"));
1817                        }
1818                        Some(idx) => {
1819                            verify_bound(&params, idx as u64, true)?;
1820                        }
1821                    }
1822                }
1823                let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
1824                self.take_rows_blocking(indices, batch_size, projection, filter)
1825            }
1826            ReadBatchParams::Range(range) => {
1827                verify_bound(&params, range.end as u64, false)?;
1828                self.read_range_blocking(
1829                    range.start as u64..range.end as u64,
1830                    batch_size,
1831                    projection,
1832                    filter,
1833                )
1834            }
1835            ReadBatchParams::Ranges(ranges) => {
1836                let mut ranges_u64 = Vec::with_capacity(ranges.len());
1837                for range in ranges.as_ref() {
1838                    verify_bound(&params, range.end, false)?;
1839                    ranges_u64.push(range.start..range.end);
1840                }
1841                self.read_ranges_blocking(ranges_u64, batch_size, projection, filter)
1842            }
1843            ReadBatchParams::RangeFrom(range) => {
1844                verify_bound(&params, range.start as u64, true)?;
1845                self.read_range_blocking(
1846                    range.start as u64..read_len,
1847                    batch_size,
1848                    projection,
1849                    filter,
1850                )
1851            }
1852            ReadBatchParams::RangeTo(range) => {
1853                verify_bound(&params, range.end as u64, false)?;
1854                self.read_range_blocking(0..range.end as u64, batch_size, projection, filter)
1855            }
1856            ReadBatchParams::RangeFull => {
1857                self.read_range_blocking(0..read_len, batch_size, projection, filter)
1858            }
1859        }
1860    }
1861
1862    /// Reads data from the file as a stream of record batches
1863    ///
1864    /// This is similar to [`Self::read_stream_projected`] but uses the base projection
1865    /// provided when the file was opened (or reads all columns if the file was
1866    /// opened without a base projection)
1867    ///
1868    /// # Why is this async?
1869    ///
1870    /// This delegates to [`Self::read_stream_projected`], which awaits the
1871    /// decode scheduler's `initialize` step before returning the stream.
1872    /// See [`Self::read_tasks`] for the rationale.
1873    pub async fn read_stream(
1874        &self,
1875        params: ReadBatchParams,
1876        batch_size: u32,
1877        batch_readahead: u32,
1878        filter: FilterExpression,
1879    ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1880        self.read_stream_projected(
1881            params,
1882            batch_size,
1883            batch_readahead,
1884            self.core.base_projection.clone(),
1885            filter,
1886        )
1887        .await
1888    }
1889
1890    pub fn schema(&self) -> &Arc<Schema> {
1891        self.core.schema()
1892    }
1893}
1894
1895impl FileMetadataProvider {
1896    fn version(&self) -> LanceFileVersion {
1897        match self {
1898            Self::Full(metadata) => metadata.version(),
1899            Self::Indexed(metadata_index) => metadata_index.version,
1900        }
1901    }
1902
1903    fn num_rows(&self) -> u64 {
1904        match self {
1905            Self::Full(metadata) => metadata.num_rows,
1906            Self::Indexed(metadata_index) => metadata_index.num_rows,
1907        }
1908    }
1909
1910    fn schema(&self) -> &Arc<Schema> {
1911        match self {
1912            Self::Full(metadata) => &metadata.file_schema,
1913            Self::Indexed(metadata_index) => &metadata_index.file_schema,
1914        }
1915    }
1916
1917    fn file_buffers(&self) -> &Vec<BufferDescriptor> {
1918        match self {
1919            Self::Full(metadata) => &metadata.file_buffers,
1920            Self::Indexed(metadata_index) => &metadata_index.file_buffers,
1921        }
1922    }
1923
1924    fn retained_global_buffers(&self) -> &BTreeMap<u32, Bytes> {
1925        match self {
1926            Self::Full(metadata) => &metadata.retained_global_buffers,
1927            Self::Indexed(metadata_index) => &metadata_index.retained_global_buffers,
1928        }
1929    }
1930
1931    fn file_statistics(&self) -> Option<FileStatistics> {
1932        let metadata = match self {
1933            Self::Full(metadata) => metadata,
1934            Self::Indexed(_) => return None,
1935        };
1936        Some(FileReader::statistics_from_column_metadata(
1937            &metadata.column_metadatas,
1938        ))
1939    }
1940
1941    fn supports_indexed_projection(
1942        projection: &ReaderProjection,
1943        version: LanceFileVersion,
1944    ) -> bool {
1945        if version < LanceFileVersion::V2_1 || projection.schema.fields.is_empty() {
1946            return false;
1947        }
1948
1949        projection
1950            .schema
1951            .fields
1952            .iter()
1953            .try_fold(0usize, |count, field| {
1954                count.checked_add(indexed_projection_column_count(field)?)
1955            })
1956            == Some(projection.column_indices.len())
1957    }
1958
1959    fn validate_indexed_projection(
1960        projection: &ReaderProjection,
1961        metadata_index: &FileMetadataIndex,
1962    ) -> Result<()> {
1963        if projection.schema.fields.is_empty() {
1964            return Err(Error::invalid_input(
1965                "Attempt to read zero columns from the file, at least one column must be specified"
1966                    .to_string(),
1967            ));
1968        }
1969        let mut column_indices_seen = BTreeSet::new();
1970        for column_index in &projection.column_indices {
1971            if !column_indices_seen.insert(*column_index) {
1972                return Err(Error::invalid_input(format!(
1973                    "The projection specified the column index {} more than once",
1974                    column_index
1975                )));
1976            }
1977            if *column_index >= metadata_index.num_columns {
1978                return Err(Error::invalid_input(format!(
1979                    "The projection specified the column index {} but there are only {} columns in the file",
1980                    column_index, metadata_index.num_columns
1981                )));
1982            }
1983        }
1984        if !Self::supports_indexed_projection(projection, metadata_index.version) {
1985            return Err(Error::not_supported(format!(
1986                "lazy column metadata loading requires a V2.1+ ordinary structural projection without blob or packed-struct fields whose physical-column count matches the projection; got file version {:?}, {} schema fields, and {} column indices",
1987                metadata_index.version,
1988                projection.schema.fields.len(),
1989                projection.column_indices.len()
1990            )));
1991        }
1992        Ok(())
1993    }
1994
1995    fn validate_projection(&self, projection: &ReaderProjection) -> Result<()> {
1996        match self {
1997            Self::Full(metadata) => FileReader::validate_projection(projection, metadata),
1998            Self::Indexed(metadata_index) => {
1999                Self::validate_indexed_projection(projection, metadata_index)
2000            }
2001        }
2002    }
2003
2004    fn column_metadata_range(
2005        metadata_index: &FileMetadataIndex,
2006        column_index: u32,
2007    ) -> Result<Range<u64>> {
2008        let (position, length) = metadata_index
2009            .column_metadata_offsets
2010            .get(column_index as usize)
2011            .copied()
2012            .ok_or_else(|| {
2013                Error::invalid_input(format!(
2014                    "The projection specified the column index {} but there are only {} columns in the file",
2015                    column_index, metadata_index.num_columns
2016                ))
2017            })?;
2018        let end = position.checked_add(length).ok_or_else(|| {
2019            Error::invalid_input(format!(
2020                "column metadata range overflows for column index {}, position={}, length={}",
2021                column_index, position, length
2022            ))
2023        })?;
2024        Ok(position..end)
2025    }
2026
2027    async fn load_indexed_column_infos(
2028        metadata_index: &FileMetadataIndex,
2029        io: &Arc<dyn EncodingsIo>,
2030        cache: &Arc<LanceCache>,
2031        column_indices: &[u32],
2032    ) -> Result<Vec<Arc<ColumnInfo>>> {
2033        let mut column_infos = vec![None; column_indices.len()];
2034        let mut missing_columns = Vec::new();
2035
2036        for (result_index, column_index) in column_indices.iter().copied().enumerate() {
2037            let cache_key = ColumnMetadataCacheKey { column_index };
2038            if let Some(cached) = cache.get_with_key(&cache_key).await {
2039                column_infos[result_index] = Some(cached.column_info.clone());
2040            } else {
2041                let range = Self::column_metadata_range(metadata_index, column_index)?;
2042                missing_columns.push((result_index, column_index, range));
2043            }
2044        }
2045
2046        missing_columns.sort_by_key(|(_, _, range)| range.start);
2047        if !missing_columns.is_empty() {
2048            let ranges = missing_columns
2049                .iter()
2050                .map(|(_, _, range)| range.clone())
2051                .collect::<Vec<_>>();
2052            let metadata_bytes = io.submit_request(ranges, 0).await?;
2053            for ((result_index, column_index, _), bytes) in
2054                missing_columns.into_iter().zip(metadata_bytes)
2055            {
2056                let column_metadata = pbfile::ColumnMetadata::decode(bytes)?;
2057                let column_info = FileReader::meta_to_col_info(
2058                    column_index,
2059                    &column_metadata,
2060                    metadata_index.version,
2061                )?;
2062                let cached = Arc::new(CachedColumnMetadata {
2063                    column_metadata,
2064                    column_info: column_info.clone(),
2065                });
2066                let cache_key = ColumnMetadataCacheKey { column_index };
2067                cache.insert_with_key(&cache_key, cached).await;
2068                column_infos[result_index] = Some(column_info);
2069            }
2070        }
2071
2072        column_infos
2073            .into_iter()
2074            .enumerate()
2075            .map(|(idx, column_info)| {
2076                column_info.ok_or_else(|| {
2077                    Error::internal(format!(
2078                        "lazy metadata loader did not load requested projection column at position {}",
2079                        idx
2080                    ))
2081                })
2082            })
2083            .collect()
2084    }
2085
2086    async fn prepare_projection(
2087        &self,
2088        projection: &ReaderProjection,
2089        io: &Arc<dyn EncodingsIo>,
2090        cache: &Arc<LanceCache>,
2091    ) -> Result<PreparedProjection> {
2092        self.validate_projection(projection)?;
2093        match self {
2094            Self::Full(metadata) => Ok(PreparedProjection {
2095                column_infos: metadata.column_infos.clone(),
2096                decoder_projection: projection.clone(),
2097            }),
2098            Self::Indexed(metadata_index) => {
2099                let column_infos = Self::load_indexed_column_infos(
2100                    metadata_index,
2101                    io,
2102                    cache,
2103                    &projection.column_indices,
2104                )
2105                .await?;
2106                let decoder_projection = ReaderProjection {
2107                    schema: projection.schema.clone(),
2108                    column_indices: (0..projection.column_indices.len())
2109                        .map(|idx| idx as u32)
2110                        .collect(),
2111                };
2112                Ok(PreparedProjection {
2113                    column_infos,
2114                    decoder_projection,
2115                })
2116            }
2117        }
2118    }
2119}
2120
2121impl FileReadCore {
2122    fn try_new(
2123        scheduler: Arc<dyn EncodingsIo>,
2124        base_projection: Option<ReaderProjection>,
2125        decoder_plugins: Arc<DecoderPlugins>,
2126        metadata_provider: FileMetadataProvider,
2127        cache: Arc<LanceCache>,
2128        options: FileReaderOptions,
2129    ) -> Result<Self> {
2130        if let Some(base_projection) = base_projection.as_ref() {
2131            metadata_provider.validate_projection(base_projection)?;
2132        }
2133        let base_projection = base_projection.unwrap_or(ReaderProjection::from_whole_schema(
2134            metadata_provider.schema().as_ref(),
2135            metadata_provider.version(),
2136        ));
2137        Ok(Self {
2138            scheduler,
2139            base_projection,
2140            metadata_provider,
2141            decoder_plugins,
2142            cache,
2143            options,
2144        })
2145    }
2146
2147    fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2148        Self {
2149            scheduler,
2150            base_projection: self.base_projection.clone(),
2151            metadata_provider: self.metadata_provider.clone(),
2152            decoder_plugins: self.decoder_plugins.clone(),
2153            cache: self.cache.clone(),
2154            options: self.options.clone(),
2155        }
2156    }
2157
2158    fn version(&self) -> LanceFileVersion {
2159        self.metadata_provider.version()
2160    }
2161
2162    fn num_rows(&self) -> u64 {
2163        self.metadata_provider.num_rows()
2164    }
2165
2166    fn schema(&self) -> &Arc<Schema> {
2167        self.metadata_provider.schema()
2168    }
2169
2170    async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2171        let file_buffers = self.metadata_provider.file_buffers();
2172        let buffer_desc = file_buffers.get(index as usize).ok_or_else(|| {
2173            Error::invalid_input(format!(
2174                "request for global buffer at index {} but there were only {} global buffers in the file",
2175                index,
2176                file_buffers.len()
2177            ))
2178        })?;
2179
2180        if let Some(bytes) = self.metadata_provider.retained_global_buffers().get(&index) {
2181            return Ok(bytes.clone());
2182        }
2183
2184        let bytes = self
2185            .scheduler
2186            .submit_request(
2187                vec![buffer_desc.position..buffer_desc.position + buffer_desc.size],
2188                0,
2189            )
2190            .await?;
2191        bytes.into_iter().next().ok_or_else(|| {
2192            Error::internal(format!(
2193                "global buffer read for index {} returned no bytes",
2194                index
2195            ))
2196        })
2197    }
2198
2199    // The common length to read across a prepared projection, after validating
2200    // its columns can be combined into rectangular batches. Each top-level field
2201    // is checked for internal consistency (see `validate_field_length`); the
2202    // top-level fields must then share a length, since one read combines them.
2203    // Ordinary files always pass (every column has `num_rows` rows); files
2204    // written with `FileWriter::write_column` whose columns ended up unequal are
2205    // rejected here and must be read separately.
2206    //
2207    // `column_infos` and `decoder_projection.column_indices` line up for both
2208    // metadata providers: the full provider keeps absolute indices into the whole
2209    // file, while the indexed (lazy) provider loads only the projected columns and
2210    // renumbers them 0..N -- in either case `column_infos[column_index]` is the
2211    // requested column.
2212    fn prepared_read_length(&self, prepared: &PreparedProjection) -> Result<u64> {
2213        let is_structural = self.version() >= LanceFileVersion::V2_1;
2214        let column_infos = &prepared.column_infos;
2215        let column_len = |column: usize| -> Result<u64> {
2216            let info = column_infos.get(column).ok_or_else(|| {
2217                Error::invalid_input(format!(
2218                    "projection references column index {} but only {} columns are available",
2219                    column,
2220                    column_infos.len()
2221                ))
2222            })?;
2223            info.page_infos.iter().try_fold(0_u64, |rows, page| {
2224                let page_rows = match &page.encoding {
2225                    PageEncoding::Structural(layout) => match &layout.layout {
2226                        Some(pbenc21::page_layout::Layout::SparseLayout(sparse)) => sparse
2227                            .structural_layers
2228                            .first()
2229                            .and_then(|layer| layer.layer.as_ref())
2230                            .map_or(page.num_rows, |layer| match layer {
2231                                pbenc21::sparse_structural_layer::Layer::Validity(layer) => {
2232                                    layer.num_slots
2233                                }
2234                                pbenc21::sparse_structural_layer::Layer::List(layer) => {
2235                                    layer.num_slots
2236                                }
2237                                pbenc21::sparse_structural_layer::Layer::FixedSizeList(layer) => {
2238                                    layer.num_slots
2239                                }
2240                            }),
2241                        _ => page.num_rows,
2242                    },
2243                    _ => page.num_rows,
2244                };
2245                rows.checked_add(page_rows).ok_or_else(|| {
2246                    Error::invalid_input_source("Column row count overflows u64".into())
2247                })
2248            })
2249        };
2250        let column_indices = &prepared.decoder_projection.column_indices;
2251        let fields = &prepared.decoder_projection.schema.fields;
2252        let mut cursor = 0usize;
2253        let mut field_lengths = Vec::with_capacity(fields.len());
2254        for field in fields {
2255            let rows = validate_field_length(
2256                field,
2257                is_structural,
2258                true,
2259                column_indices,
2260                &mut cursor,
2261                &column_len,
2262            )?;
2263            field_lengths.push((field.name.as_str(), rows));
2264        }
2265        if cursor != column_indices.len() {
2266            return Err(Error::invalid_input(format!(
2267                "projection supplied {} column indices but its fields require {}",
2268                column_indices.len(),
2269                cursor
2270            )));
2271        }
2272        verify_uniform_lengths(&field_lengths)
2273    }
2274
2275    async fn read_range(
2276        &self,
2277        range: Range<u64>,
2278        batch_size: u32,
2279        prepared: PreparedProjection,
2280        filter: FilterExpression,
2281    ) -> Result<BoxStream<'static, ReadBatchTask>> {
2282        FileReader::do_read_range(
2283            prepared.column_infos,
2284            self.scheduler.clone(),
2285            self.cache.clone(),
2286            self.num_rows(),
2287            self.decoder_plugins.clone(),
2288            range,
2289            batch_size,
2290            prepared.decoder_projection,
2291            filter,
2292            self.options.decoder_config.clone(),
2293            self.options.batch_size_bytes,
2294        )
2295        .await
2296    }
2297
2298    async fn take_rows(
2299        &self,
2300        indices: Vec<u64>,
2301        batch_size: u32,
2302        prepared: PreparedProjection,
2303    ) -> Result<BoxStream<'static, ReadBatchTask>> {
2304        FileReader::do_take_rows(
2305            prepared.column_infos,
2306            self.scheduler.clone(),
2307            self.cache.clone(),
2308            self.decoder_plugins.clone(),
2309            indices,
2310            batch_size,
2311            prepared.decoder_projection,
2312            FilterExpression::no_filter(),
2313            self.options.decoder_config.clone(),
2314            self.options.batch_size_bytes,
2315        )
2316        .await
2317    }
2318
2319    async fn read_ranges(
2320        &self,
2321        ranges: Vec<Range<u64>>,
2322        batch_size: u32,
2323        prepared: PreparedProjection,
2324        filter: FilterExpression,
2325    ) -> Result<BoxStream<'static, ReadBatchTask>> {
2326        FileReader::do_read_ranges(
2327            prepared.column_infos,
2328            self.scheduler.clone(),
2329            self.cache.clone(),
2330            self.decoder_plugins.clone(),
2331            ranges,
2332            batch_size,
2333            prepared.decoder_projection,
2334            filter,
2335            self.options.decoder_config.clone(),
2336            self.options.batch_size_bytes,
2337        )
2338        .await
2339    }
2340
2341    async fn read_tasks(
2342        &self,
2343        params: ReadBatchParams,
2344        batch_size: u32,
2345        projection: Option<ReaderProjection>,
2346        filter: FilterExpression,
2347    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2348        let projection = projection.unwrap_or_else(|| self.base_projection.clone());
2349        let prepared = self
2350            .metadata_provider
2351            .prepare_projection(&projection, &self.scheduler, &self.cache)
2352            .await?;
2353        // All projected columns must share a length: the reader combines them
2354        // into rectangular batches.  Ordinary files satisfy this (every column
2355        // has `num_rows` rows); files written with `FileWriter::write_column`
2356        // may not, and such columns must be read separately.  `read_len` is that
2357        // common length, which `RangeFull`/`RangeFrom` resolve against (rather
2358        // than `num_rows`, the file's longest column).
2359        let read_len = self.prepared_read_length(&prepared)?;
2360        let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
2361            if bound > read_len || (bound == read_len && inclusive) {
2362                Err(Error::invalid_input(format!(
2363                    "cannot read {params:?} from columns with {read_len} rows"
2364                )))
2365            } else {
2366                Ok(())
2367            }
2368        };
2369        match &params {
2370            ReadBatchParams::Indices(indices) => {
2371                for idx in indices {
2372                    match idx {
2373                        None => {
2374                            return Err(Error::invalid_input("Null value in indices array"));
2375                        }
2376                        Some(idx) => {
2377                            verify_bound(&params, idx as u64, true)?;
2378                        }
2379                    }
2380                }
2381                let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
2382                self.take_rows(indices, batch_size, prepared).await
2383            }
2384            ReadBatchParams::Range(range) => {
2385                verify_bound(&params, range.end as u64, false)?;
2386                self.read_range(
2387                    range.start as u64..range.end as u64,
2388                    batch_size,
2389                    prepared,
2390                    filter,
2391                )
2392                .await
2393            }
2394            ReadBatchParams::Ranges(ranges) => {
2395                let mut ranges_u64 = Vec::with_capacity(ranges.len());
2396                for range in ranges.as_ref() {
2397                    verify_bound(&params, range.end, false)?;
2398                    ranges_u64.push(range.start..range.end);
2399                }
2400                self.read_ranges(ranges_u64, batch_size, prepared, filter)
2401                    .await
2402            }
2403            ReadBatchParams::RangeFrom(range) => {
2404                verify_bound(&params, range.start as u64, true)?;
2405                self.read_range(range.start as u64..read_len, batch_size, prepared, filter)
2406                    .await
2407            }
2408            ReadBatchParams::RangeTo(range) => {
2409                verify_bound(&params, range.end as u64, false)?;
2410                self.read_range(0..range.end as u64, batch_size, prepared, filter)
2411                    .await
2412            }
2413            ReadBatchParams::RangeFull => {
2414                self.read_range(0..read_len, batch_size, prepared, filter)
2415                    .await
2416            }
2417        }
2418    }
2419}
2420
2421impl ProjectedFileReader {
2422    /// Opens a data reader backed by indexed column metadata.
2423    ///
2424    /// `base_projection` must be a supported indexed projection. Reads that do
2425    /// not pass an explicit projection use this base projection.
2426    pub async fn try_open(
2427        scheduler: FileScheduler,
2428        base_projection: Option<ReaderProjection>,
2429        decoder_plugins: Arc<DecoderPlugins>,
2430        cache: &LanceCache,
2431        options: FileReaderOptions,
2432    ) -> Result<Self> {
2433        let base_projection = Self::require_indexed_base_projection(base_projection)?;
2434        let metadata_index = Arc::new(FileReader::read_metadata_index(&scheduler).await?);
2435        let path = scheduler.reader().path().clone();
2436        let encodings_io =
2437            LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
2438        Self::try_open_with_metadata_index(
2439            Arc::new(encodings_io),
2440            path,
2441            Some(base_projection),
2442            decoder_plugins,
2443            metadata_index,
2444            cache,
2445            options,
2446        )
2447        .await
2448    }
2449
2450    /// Opens a data reader from a previously loaded metadata index.
2451    ///
2452    /// `base_projection` must be a supported indexed projection. Use
2453    /// [`Self::try_open_with_file_metadata`] when the default read should cover
2454    /// the whole file schema.
2455    pub async fn try_open_with_metadata_index(
2456        scheduler: Arc<dyn EncodingsIo>,
2457        path: Path,
2458        base_projection: Option<ReaderProjection>,
2459        decoder_plugins: Arc<DecoderPlugins>,
2460        metadata_index: Arc<FileMetadataIndex>,
2461        cache: &LanceCache,
2462        options: FileReaderOptions,
2463    ) -> Result<Self> {
2464        let base_projection = Self::require_indexed_base_projection(base_projection)?;
2465        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2466        let core = FileReadCore::try_new(
2467            scheduler,
2468            Some(base_projection),
2469            decoder_plugins,
2470            FileMetadataProvider::Indexed(metadata_index),
2471            cache,
2472            options,
2473        )?;
2474        Ok(Self { core })
2475    }
2476
2477    fn require_indexed_base_projection(
2478        base_projection: Option<ReaderProjection>,
2479    ) -> Result<ReaderProjection> {
2480        base_projection.ok_or_else(|| {
2481            Error::invalid_input("ProjectedFileReader requires an explicit base projection")
2482        })
2483    }
2484
2485    pub async fn try_open_with_file_metadata(
2486        scheduler: Arc<dyn EncodingsIo>,
2487        path: Path,
2488        base_projection: Option<ReaderProjection>,
2489        decoder_plugins: Arc<DecoderPlugins>,
2490        file_metadata: Arc<CachedFileMetadata>,
2491        cache: &LanceCache,
2492        options: FileReaderOptions,
2493    ) -> Result<Self> {
2494        let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2495        let core = FileReadCore::try_new(
2496            scheduler,
2497            base_projection,
2498            decoder_plugins,
2499            FileMetadataProvider::Full(file_metadata),
2500            cache,
2501            options,
2502        )?;
2503        Ok(Self { core })
2504    }
2505
2506    /// Returns whether a projection can be served by indexed column metadata.
2507    pub fn supports_projection(projection: &ReaderProjection, version: LanceFileVersion) -> bool {
2508        FileMetadataProvider::supports_indexed_projection(projection, version)
2509    }
2510
2511    /// Returns a clone of this reader using a different scheduler.
2512    pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2513        Self {
2514            core: self.core.with_scheduler(scheduler),
2515        }
2516    }
2517
2518    /// Returns the Lance file version.
2519    pub fn version(&self) -> LanceFileVersion {
2520        self.core.version()
2521    }
2522
2523    /// Returns the number of rows in the file.
2524    pub fn num_rows(&self) -> u64 {
2525        self.core.num_rows()
2526    }
2527
2528    /// Returns the file schema visible to this reader.
2529    pub fn schema(&self) -> &Arc<Schema> {
2530        self.core.schema()
2531    }
2532
2533    /// Returns file statistics when this reader has full file metadata.
2534    pub fn file_statistics(&self) -> Option<FileStatistics> {
2535        self.core.metadata_provider.file_statistics()
2536    }
2537
2538    #[cfg(test)]
2539    fn metadata_index(&self) -> Option<&Arc<FileMetadataIndex>> {
2540        match &self.core.metadata_provider {
2541            FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index),
2542            FileMetadataProvider::Full(_) => None,
2543        }
2544    }
2545
2546    /// Reads a global buffer by index.
2547    pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2548        self.core.read_global_buffer(index).await
2549    }
2550
2551    /// Creates a stream of read tasks for the requested rows and projection.
2552    pub async fn read_tasks(
2553        &self,
2554        params: ReadBatchParams,
2555        batch_size: u32,
2556        projection: Option<ReaderProjection>,
2557        filter: FilterExpression,
2558    ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2559        self.core
2560            .read_tasks(params, batch_size, projection, filter)
2561            .await
2562    }
2563}
2564
2565/// Inspects a page and returns a String describing the page's encoding
2566pub fn describe_encoding(page: &pbfile::column_metadata::Page) -> String {
2567    if let Some(encoding) = &page.encoding {
2568        if let Some(style) = &encoding.location {
2569            match style {
2570                pbfile::encoding::Location::Indirect(indirect) => {
2571                    format!(
2572                        "IndirectEncoding(pos={},size={})",
2573                        indirect.buffer_location, indirect.buffer_length
2574                    )
2575                }
2576                pbfile::encoding::Location::Direct(direct) => {
2577                    let encoding_any =
2578                        prost_types::Any::decode(Bytes::from(direct.encoding.clone()))
2579                            .expect("failed to deserialize encoding as protobuf");
2580                    if encoding_any.type_url == "/lance.encodings.ArrayEncoding" {
2581                        let encoding = encoding_any.to_msg::<pbenc::ArrayEncoding>();
2582                        match encoding {
2583                            Ok(encoding) => {
2584                                format!("{:#?}", encoding)
2585                            }
2586                            Err(err) => {
2587                                format!("Unsupported(decode_err={})", err)
2588                            }
2589                        }
2590                    } else if encoding_any.type_url == "/lance.encodings21.PageLayout" {
2591                        let encoding = encoding_any.to_msg::<pbenc21::PageLayout>();
2592                        match encoding {
2593                            Ok(encoding) => {
2594                                format!("{:#?}", encoding)
2595                            }
2596                            Err(err) => {
2597                                format!("Unsupported(decode_err={})", err)
2598                            }
2599                        }
2600                    } else {
2601                        format!("Unrecognized(type_url={})", encoding_any.type_url)
2602                    }
2603                }
2604                pbfile::encoding::Location::None(_) => "NoEncodingDescription".to_string(),
2605            }
2606        } else {
2607            "MISSING STYLE".to_string()
2608        }
2609    } else {
2610        "MISSING".to_string()
2611    }
2612}
2613
2614pub trait EncodedBatchReaderExt {
2615    fn try_from_mini_lance(
2616        bytes: Bytes,
2617        schema: &Schema,
2618        version: LanceFileVersion,
2619    ) -> Result<Self>
2620    where
2621        Self: Sized;
2622    fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2623    where
2624        Self: Sized;
2625}
2626
2627impl EncodedBatchReaderExt for EncodedBatch {
2628    fn try_from_mini_lance(
2629        bytes: Bytes,
2630        schema: &Schema,
2631        file_version: LanceFileVersion,
2632    ) -> Result<Self>
2633    where
2634        Self: Sized,
2635    {
2636        let projection = ReaderProjection::from_whole_schema(schema, file_version);
2637        let footer = FileReader::decode_footer(&bytes)?;
2638
2639        // Next, read the metadata for the columns
2640        // This is both the column metadata and the CMO table
2641        let column_metadata_start = footer.column_meta_start as usize;
2642        let column_metadata_end = footer.global_buff_offsets_start as usize;
2643        let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2644        let column_metadatas =
2645            FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2646
2647        let file_version: LanceFileVersion =
2648            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
2649                .into();
2650
2651        let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?;
2652
2653        Ok(Self {
2654            data: bytes,
2655            num_rows: page_table
2656                .first()
2657                .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2658                .unwrap_or(0),
2659            page_table,
2660            top_level_columns: projection.column_indices,
2661            schema: Arc::new(schema.clone()),
2662        })
2663    }
2664
2665    fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2666    where
2667        Self: Sized,
2668    {
2669        let footer = FileReader::decode_footer(&bytes)?;
2670        let file_version: LanceFileVersion =
2671            ConcreteFileVersion::from_footer_numbers(footer.major_version, footer.minor_version)?
2672                .into();
2673
2674        let gbo_table = FileReader::do_decode_gbo_table(
2675            &bytes.slice(footer.global_buff_offsets_start as usize..),
2676            &footer,
2677            file_version,
2678        )?;
2679        if gbo_table.is_empty() {
2680            return Err(Error::internal(
2681                "File did not contain any global buffers, schema expected".to_string(),
2682            ));
2683        }
2684        let schema_start = gbo_table[0].position as usize;
2685        let schema_size = gbo_table[0].size as usize;
2686
2687        let schema_bytes = bytes.slice(schema_start..(schema_start + schema_size));
2688        let (_, schema) = FileReader::decode_schema(schema_bytes)?;
2689        let projection = ReaderProjection::from_whole_schema(&schema, file_version);
2690
2691        // Next, read the metadata for the columns
2692        // This is both the column metadata and the CMO table
2693        let column_metadata_start = footer.column_meta_start as usize;
2694        let column_metadata_end = footer.global_buff_offsets_start as usize;
2695        let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2696        let column_metadatas =
2697            FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2698
2699        let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version)?;
2700
2701        Ok(Self {
2702            data: bytes,
2703            num_rows: page_table
2704                .first()
2705                .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2706                .unwrap_or(0),
2707            page_table,
2708            top_level_columns: projection.column_indices,
2709            schema: Arc::new(schema),
2710        })
2711    }
2712}
2713
2714#[cfg(test)]
2715mod tests {
2716    use std::{
2717        collections::{BTreeMap, HashMap},
2718        pin::Pin,
2719        sync::Arc,
2720    };
2721
2722    use arrow_array::{
2723        Int32Array, ListArray, RecordBatch, RecordBatchIterator, UInt32Array,
2724        types::{Float64Type, Int32Type},
2725    };
2726    use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
2727    use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
2728    use bytes::Bytes;
2729    use futures::{StreamExt, prelude::stream::TryStreamExt};
2730    use lance_arrow::{BLOB_META_KEY, RecordBatchExt};
2731    use lance_core::{ArrowResult, datatypes::Schema};
2732    use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch};
2733    use lance_encoding::{
2734        constants::{STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_SPARSE},
2735        decoder::{
2736            DecodeBatchScheduler, DecoderPlugins, EncodedBatchLayout, FilterExpression,
2737            PageEncoding, ReadBatchTask, decode_batch,
2738        },
2739        encoder::{EncodedBatch, EncodingOptions, encode_batch},
2740        format::pb21,
2741        version::LanceFileVersion,
2742    };
2743    use lance_io::{stream::RecordBatchStream, utils::CachedFileSize};
2744    use log::debug;
2745    use rstest::rstest;
2746    use tokio::sync::mpsc;
2747
2748    use crate::reader::{
2749        EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader,
2750        ReaderProjection, validate_field_length, verify_uniform_lengths,
2751    };
2752    use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file};
2753    use crate::writer::{EncodedBatchWriteExt, FileWriter, FileWriterOptions};
2754    use lance_encoding::decoder::DecoderConfig;
2755
2756    fn footer_version(bytes: &[u8]) -> (u16, u16) {
2757        let version_start = bytes.len() - 8;
2758        (
2759            u16::from_le_bytes([bytes[version_start], bytes[version_start + 1]]),
2760            u16::from_le_bytes([bytes[version_start + 2], bytes[version_start + 3]]),
2761        )
2762    }
2763
2764    #[tokio::test]
2765    async fn sparse_file_writer_reader_scan_range_and_take_roundtrip() {
2766        let fs = FsFixture::default();
2767        let sparse_metadata = HashMap::from([(
2768            STRUCTURAL_ENCODING_META_KEY.to_string(),
2769            STRUCTURAL_ENCODING_SPARSE.to_string(),
2770        )]);
2771        let value_field =
2772            Field::new("values", DataType::Int32, true).with_metadata(sparse_metadata.clone());
2773        let item_field = Arc::new(Field::new("item", DataType::Int32, true));
2774        let list_field = Field::new("items", DataType::List(item_field.clone()), true)
2775            .with_metadata(sparse_metadata);
2776        let arrow_schema = Arc::new(ArrowSchema::new(vec![value_field, list_field]));
2777        let list = ListArray::try_new(
2778            item_field,
2779            OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 2, 2, 3, 3, 5])),
2780            Arc::new(Int32Array::from(vec![
2781                Some(1),
2782                None,
2783                Some(3),
2784                Some(4),
2785                Some(5),
2786            ])),
2787            Some(NullBuffer::from(vec![true, false, true, true, true, true])),
2788        )
2789        .unwrap();
2790        let batch = RecordBatch::try_new(
2791            arrow_schema.clone(),
2792            vec![
2793                Arc::new(Int32Array::from(vec![
2794                    Some(10),
2795                    None,
2796                    Some(30),
2797                    Some(40),
2798                    None,
2799                    Some(60),
2800                ])),
2801                Arc::new(list),
2802            ],
2803        )
2804        .unwrap();
2805        let input = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema);
2806        write_lance_file(
2807            input,
2808            &fs,
2809            FileWriterOptions {
2810                format_version: Some(LanceFileVersion::V2_3),
2811                ..Default::default()
2812            },
2813        )
2814        .await;
2815
2816        let file_scheduler = fs
2817            .scheduler
2818            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2819            .await
2820            .unwrap();
2821        let file_reader = FileReader::try_open(
2822            file_scheduler,
2823            None,
2824            Arc::<DecoderPlugins>::default(),
2825            &test_cache(),
2826            FileReaderOptions::default(),
2827        )
2828        .await
2829        .unwrap();
2830        assert_eq!(file_reader.metadata.column_infos.len(), 2);
2831        assert!(
2832            file_reader
2833                .metadata
2834                .column_infos
2835                .iter()
2836                .flat_map(|column| column.page_infos.iter())
2837                .all(|page| {
2838                    matches!(
2839                        &page.encoding,
2840                        PageEncoding::Structural(layout)
2841                            if matches!(
2842                                layout.layout,
2843                                Some(pb21::page_layout::Layout::SparseLayout(_))
2844                            )
2845                    )
2846                })
2847        );
2848
2849        let scan = file_reader
2850            .read_stream(
2851                lance_io::ReadBatchParams::RangeFull,
2852                1024,
2853                1,
2854                FilterExpression::no_filter(),
2855            )
2856            .await
2857            .unwrap()
2858            .try_collect::<Vec<_>>()
2859            .await
2860            .unwrap();
2861        assert_eq!(scan, vec![batch.clone()]);
2862
2863        let range = file_reader
2864            .read_stream(
2865                lance_io::ReadBatchParams::Range(1..5),
2866                1024,
2867                1,
2868                FilterExpression::no_filter(),
2869            )
2870            .await
2871            .unwrap()
2872            .try_collect::<Vec<_>>()
2873            .await
2874            .unwrap();
2875        assert_eq!(range, vec![batch.slice(1, 4)]);
2876
2877        let indices = UInt32Array::from(vec![0, 3, 5]);
2878        let take = file_reader
2879            .read_stream(
2880                lance_io::ReadBatchParams::Indices(indices.clone()),
2881                1024,
2882                1,
2883                FilterExpression::no_filter(),
2884            )
2885            .await
2886            .unwrap()
2887            .try_collect::<Vec<_>>()
2888            .await
2889            .unwrap();
2890        assert_eq!(take, vec![batch.take(&indices).unwrap()]);
2891    }
2892
2893    async fn create_some_file(fs: &FsFixture, version: LanceFileVersion) -> WrittenFile {
2894        let location_type = DataType::Struct(Fields::from(vec![
2895            Field::new("x", DataType::Float64, true),
2896            Field::new("y", DataType::Float64, true),
2897        ]));
2898        let categories_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2899
2900        let mut reader = gen_batch()
2901            .col("score", array::rand::<Float64Type>())
2902            .col("location", array::rand_type(&location_type))
2903            .col("categories", array::rand_type(&categories_type))
2904            .col("binary", array::rand_type(&DataType::Binary));
2905        if version <= LanceFileVersion::V2_0 {
2906            reader = reader.col("large_bin", array::rand_type(&DataType::LargeBinary));
2907        }
2908        let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2909
2910        write_lance_file(
2911            reader,
2912            fs,
2913            FileWriterOptions {
2914                format_version: Some(version),
2915                ..Default::default()
2916            },
2917        )
2918        .await
2919    }
2920
2921    async fn create_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2922        let mut reader = gen_batch();
2923        for column_idx in 0..num_columns {
2924            reader = reader.col(format!("c{column_idx}"), array::step::<Int32Type>());
2925        }
2926        let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2927
2928        write_lance_file(
2929            reader,
2930            fs,
2931            FileWriterOptions {
2932                format_version: Some(LanceFileVersion::V2_1),
2933                ..Default::default()
2934            },
2935        )
2936        .await
2937    }
2938
2939    async fn create_wide_fixed_size_list_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2940        let data_type =
2941            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
2942        let mut reader = gen_batch();
2943        for column_idx in 0..num_columns {
2944            reader = reader.col(
2945                format!("c{column_idx}"),
2946                array::rand_type(&data_type).with_random_nulls(0.1),
2947            );
2948        }
2949        let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2950
2951        write_lance_file(
2952            reader,
2953            fs,
2954            FileWriterOptions {
2955                format_version: Some(LanceFileVersion::V2_1),
2956                ..Default::default()
2957            },
2958        )
2959        .await
2960    }
2961
2962    async fn create_wide_structural_file(fs: &FsFixture, num_groups: usize) -> WrittenFile {
2963        let struct_type = DataType::Struct(Fields::from(vec![
2964            Field::new("x", DataType::Int32, true),
2965            Field::new("y", DataType::Int32, true),
2966        ]));
2967        let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
2968        let mut reader = gen_batch();
2969        for group_idx in 0..num_groups {
2970            reader = reader
2971                .col(
2972                    format!("s{group_idx}"),
2973                    array::rand_type(&struct_type).with_random_nulls(0.5),
2974                )
2975                .col(
2976                    format!("l{group_idx}"),
2977                    array::rand_type(&list_type).with_random_nulls(0.5),
2978                );
2979        }
2980        let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2981
2982        write_lance_file(
2983            reader,
2984            fs,
2985            FileWriterOptions {
2986                format_version: Some(LanceFileVersion::V2_1),
2987                ..Default::default()
2988            },
2989        )
2990        .await
2991    }
2992
2993    type Transformer = Box<dyn Fn(&RecordBatch) -> RecordBatch>;
2994
2995    async fn verify_expected(
2996        expected: &[RecordBatch],
2997        mut actual: Pin<Box<dyn RecordBatchStream>>,
2998        read_size: u32,
2999        transform: Option<Transformer>,
3000    ) {
3001        let mut remaining = expected.iter().map(|batch| batch.num_rows()).sum::<usize>() as u32;
3002        let mut expected_iter = expected.iter().map(|batch| {
3003            if let Some(transform) = &transform {
3004                transform(batch)
3005            } else {
3006                batch.clone()
3007            }
3008        });
3009        let mut next_expected = expected_iter.next().unwrap().clone();
3010        while let Some(actual) = actual.next().await {
3011            let mut actual = actual.unwrap();
3012            let mut rows_to_verify = actual.num_rows() as u32;
3013            let expected_length = remaining.min(read_size);
3014            assert_eq!(expected_length, rows_to_verify);
3015
3016            while rows_to_verify > 0 {
3017                let next_slice_len = (next_expected.num_rows() as u32).min(rows_to_verify);
3018                assert_eq!(
3019                    next_expected.slice(0, next_slice_len as usize),
3020                    actual.slice(0, next_slice_len as usize)
3021                );
3022                remaining -= next_slice_len;
3023                rows_to_verify -= next_slice_len;
3024                if remaining > 0 {
3025                    if next_slice_len == next_expected.num_rows() as u32 {
3026                        next_expected = expected_iter.next().unwrap().clone();
3027                    } else {
3028                        next_expected = next_expected.slice(
3029                            next_slice_len as usize,
3030                            next_expected.num_rows() - next_slice_len as usize,
3031                        );
3032                    }
3033                }
3034                if rows_to_verify > 0 {
3035                    actual = actual.slice(
3036                        next_slice_len as usize,
3037                        actual.num_rows() - next_slice_len as usize,
3038                    );
3039                }
3040            }
3041        }
3042        assert_eq!(remaining, 0);
3043    }
3044
3045    async fn collect_read_tasks(
3046        tasks: Pin<Box<dyn futures::Stream<Item = ReadBatchTask> + Send>>,
3047        readahead: usize,
3048    ) -> Vec<RecordBatch> {
3049        tasks
3050            .map(|task| task.task)
3051            .buffered(readahead)
3052            .try_collect::<Vec<_>>()
3053            .await
3054            .unwrap()
3055    }
3056
3057    /// Writes `batch` to a fresh file, overwrites `patch` bytes at `patch_offset`
3058    /// into the single occurrence of `pattern`, and reads the file back with the
3059    /// default reader configuration.
3060    async fn read_file_with_mutated_bytes(
3061        version: LanceFileVersion,
3062        batch: RecordBatch,
3063        pattern: &[u8],
3064        patch_offset: usize,
3065        patch: &[u8],
3066    ) -> lance_core::Result<Vec<RecordBatch>> {
3067        let fs = FsFixture::default();
3068        let schema = batch.schema();
3069        write_lance_file(
3070            RecordBatchIterator::new(vec![Ok(batch)], schema),
3071            &fs,
3072            FileWriterOptions {
3073                format_version: Some(version),
3074                ..Default::default()
3075            },
3076        )
3077        .await;
3078
3079        let mut bytes = fs
3080            .object_store
3081            .read_one_all(&fs.tmp_path)
3082            .await
3083            .unwrap()
3084            .to_vec();
3085        let matches = bytes
3086            .windows(pattern.len())
3087            .enumerate()
3088            .filter_map(|(position, window)| (window == pattern).then_some(position))
3089            .collect::<Vec<_>>();
3090        assert_eq!(
3091            matches.len(),
3092            1,
3093            "expected the byte pattern to appear exactly once in the file"
3094        );
3095        let patch_start = matches[0] + patch_offset;
3096        bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch);
3097        fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap();
3098
3099        let file_scheduler = fs
3100            .scheduler
3101            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3102            .await
3103            .unwrap();
3104        let file_reader = FileReader::try_open(
3105            file_scheduler,
3106            None,
3107            Arc::<DecoderPlugins>::default(),
3108            &test_cache(),
3109            FileReaderOptions::default(),
3110        )
3111        .await
3112        .unwrap();
3113        file_reader
3114            .read_stream(
3115                lance_io::ReadBatchParams::RangeFull,
3116                1024,
3117                16,
3118                FilterExpression::no_filter(),
3119            )
3120            .await?
3121            .try_collect::<Vec<_>>()
3122            .await
3123    }
3124
3125    /// A corrupt file whose variable-width offsets point outside the value bytes
3126    /// must fail with a typed error under the default reader configuration
3127    /// (`validate_on_decode` disabled) instead of materializing values outside
3128    /// the data buffer.
3129    ///
3130    /// Uses a dictionary-encoded string column because its values page stores
3131    /// the offsets verbatim, so flipping the tail offset in the file reaches the
3132    /// Arrow conversion boundary without being rejected by an intermediate
3133    /// decompressor.
3134    #[rstest]
3135    #[tokio::test]
3136    async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets(
3137        #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3138        version: LanceFileVersion,
3139    ) {
3140        use arrow_array::{Array, DictionaryArray, Int32Array, StringArray};
3141
3142        let values = StringArray::from(vec!["alpha", "beta", "gamma"]);
3143        let indices = Int32Array::from((0..300).map(|i| i % 3).collect::<Vec<i32>>());
3144        let dictionary = DictionaryArray::new(indices, Arc::new(values));
3145        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3146            "category",
3147            dictionary.data_type().clone(),
3148            false,
3149        )]));
3150        let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap();
3151
3152        // The dictionary values page stores the value offsets as plain
3153        // little-endian i32s ending with [5, 9, 14] (2.1 also stores the leading
3154        // zero, 2.2+ omits it).  If a future encoding change stops storing these
3155        // offsets verbatim this lookup fails loudly and the test needs a new
3156        // byte pattern.  The patch rewrites the tail offset so it points far
3157        // beyond the value bytes.
3158        let offsets_tail_pattern = [5_i32, 9, 14]
3159            .iter()
3160            .flat_map(|value| value.to_le_bytes())
3161            .collect::<Vec<u8>>();
3162        let error = read_file_with_mutated_bytes(
3163            version,
3164            batch,
3165            &offsets_tail_pattern,
3166            8,
3167            &100_000_i32.to_le_bytes(),
3168        )
3169        .await
3170        .expect_err("out-of-bounds offsets must fail the read");
3171        assert!(
3172            matches!(error, lance_core::Error::CorruptFile { .. }),
3173            "expected CorruptFile, got: {error}"
3174        );
3175        assert!(
3176            error.to_string().contains("out of bounds"),
3177            "unexpected message: {error}"
3178        );
3179    }
3180
3181    /// Same contract as the test above, but for a plain (non-dictionary) string
3182    /// column: the mini-block chunk stores chunk-relative value offsets that are
3183    /// used to slice the chunk, so a corrupt tail offset must surface as a typed
3184    /// error from the chunk decompressor instead of a panic in the decode task.
3185    #[rstest]
3186    #[tokio::test]
3187    async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets(
3188        #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
3189        version: LanceFileVersion,
3190    ) {
3191        use arrow_array::StringArray;
3192
3193        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
3194            "strings",
3195            DataType::Utf8,
3196            false,
3197        )]));
3198        let batch = RecordBatch::try_new(
3199            arrow_schema,
3200            vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))],
3201        )
3202        .unwrap();
3203
3204        // For ["alpha", "beta", "gamma"] the chunk stores LE i32 offsets
3205        // [16, 21, 25, 30] (chunk-relative: a 16-byte offsets region precedes
3206        // the value bytes).  The patch rewrites the tail offset to point far
3207        // past the chunk.
3208        let chunk_offsets_pattern = [16_i32, 21, 25, 30]
3209            .iter()
3210            .flat_map(|value| value.to_le_bytes())
3211            .collect::<Vec<u8>>();
3212        let error = read_file_with_mutated_bytes(
3213            version,
3214            batch,
3215            &chunk_offsets_pattern,
3216            12,
3217            &100_000_i32.to_le_bytes(),
3218        )
3219        .await
3220        .expect_err("an out-of-bounds chunk offset must fail the read");
3221        assert!(
3222            matches!(error, lance_core::Error::CorruptFile { .. }),
3223            "expected CorruptFile, got: {error}"
3224        );
3225        assert!(
3226            error.to_string().contains("out of bounds"),
3227            "unexpected message: {error}"
3228        );
3229    }
3230
3231    #[tokio::test]
3232    async fn test_round_trip() {
3233        let fs = FsFixture::default();
3234
3235        let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3236
3237        let file_size = fs.object_store.size(&fs.tmp_path).await.unwrap() as usize;
3238        let footer = fs
3239            .object_store
3240            .open(&fs.tmp_path)
3241            .await
3242            .unwrap()
3243            .get_range(file_size - 8..file_size)
3244            .await
3245            .unwrap();
3246        assert_eq!(footer_version(&footer), (0, 3));
3247        assert_eq!(
3248            crate::determine_file_version(&fs.object_store, &fs.tmp_path, Some(file_size))
3249                .await
3250                .unwrap(),
3251            LanceFileVersion::V2_0
3252        );
3253
3254        for read_size in [32, 1024, 1024 * 1024] {
3255            let file_scheduler = fs
3256                .scheduler
3257                .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3258                .await
3259                .unwrap();
3260            let file_reader = FileReader::try_open(
3261                file_scheduler,
3262                None,
3263                Arc::<DecoderPlugins>::default(),
3264                &test_cache(),
3265                FileReaderOptions::default(),
3266            )
3267            .await
3268            .unwrap();
3269
3270            let schema = file_reader.schema();
3271            assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
3272
3273            let batch_stream = file_reader
3274                .read_stream(
3275                    lance_io::ReadBatchParams::RangeFull,
3276                    read_size,
3277                    16,
3278                    FilterExpression::no_filter(),
3279                )
3280                .await
3281                .unwrap();
3282
3283            verify_expected(&data, batch_stream, read_size, None).await;
3284        }
3285    }
3286
3287    #[rstest]
3288    #[test_log::test(tokio::test)]
3289    async fn test_encoded_batch_round_trip(
3290        // TODO: Add V2_1 (currently fails)
3291        #[values(LanceFileVersion::V2_0)] version: LanceFileVersion,
3292    ) {
3293        let data = gen_batch()
3294            .col("x", array::rand::<Int32Type>())
3295            .col("y", array::rand_utf8(ByteCount::from(16), false))
3296            .into_batch_rows(RowCount::from(10000))
3297            .unwrap();
3298
3299        let lance_schema = Arc::new(Schema::try_from(data.schema().as_ref()).unwrap());
3300
3301        let encoding_options = EncodingOptions {
3302            cache_bytes_per_column: 4096,
3303            max_page_bytes: 32 * 1024 * 1024,
3304            keep_original_array: true,
3305            buffer_alignment: 64,
3306        };
3307
3308        let encoding_strategy = crate::versions::v2_0::encoding_strategy();
3309
3310        let encoded_batch = encode_batch(
3311            &data,
3312            lance_schema.clone(),
3313            encoding_strategy.as_ref(),
3314            &encoding_options,
3315        )
3316        .await
3317        .unwrap();
3318
3319        // Test self described
3320        let bytes = encoded_batch.try_to_self_described_lance(version).unwrap();
3321        assert_eq!(footer_version(&bytes), (2, 0));
3322
3323        let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap();
3324
3325        let decoded = decode_batch(
3326            &decoded_batch,
3327            &FilterExpression::no_filter(),
3328            Arc::<DecoderPlugins>::default(),
3329            false,
3330            EncodedBatchLayout::Array,
3331            None,
3332        )
3333        .await
3334        .unwrap();
3335
3336        assert_eq!(data, decoded);
3337
3338        // Test mini
3339        let bytes = encoded_batch.try_to_mini_lance(version).unwrap();
3340        assert_eq!(footer_version(&bytes), (2, 0));
3341        let decoded_batch =
3342            EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref(), LanceFileVersion::V2_0)
3343                .unwrap();
3344        let decoded = decode_batch(
3345            &decoded_batch,
3346            &FilterExpression::no_filter(),
3347            Arc::<DecoderPlugins>::default(),
3348            false,
3349            EncodedBatchLayout::Array,
3350            None,
3351        )
3352        .await
3353        .unwrap();
3354
3355        assert_eq!(data, decoded);
3356    }
3357
3358    #[rstest]
3359    #[test_log::test(tokio::test)]
3360    async fn test_projection(
3361        #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
3362        version: LanceFileVersion,
3363    ) {
3364        let fs = FsFixture::default();
3365
3366        let written_file = create_some_file(&fs, version).await;
3367        let file_scheduler = fs
3368            .scheduler
3369            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3370            .await
3371            .unwrap();
3372
3373        let field_id_mapping = written_file
3374            .field_id_mapping
3375            .iter()
3376            .copied()
3377            .collect::<BTreeMap<_, _>>();
3378
3379        let empty_projection = ReaderProjection {
3380            column_indices: Vec::default(),
3381            schema: Arc::new(Schema::default()),
3382        };
3383
3384        for columns in [
3385            vec!["score"],
3386            vec!["location"],
3387            vec!["categories"],
3388            vec!["score.x"],
3389            vec!["score", "categories"],
3390            vec!["score", "location"],
3391            vec!["location", "categories"],
3392            vec!["score.y", "location", "categories"],
3393        ] {
3394            debug!("Testing round trip with projection {:?}", columns);
3395            for use_field_ids in [true, false] {
3396                // We can specify the projection as part of the read operation via read_stream_projected
3397                let file_reader = FileReader::try_open(
3398                    file_scheduler.clone(),
3399                    None,
3400                    Arc::<DecoderPlugins>::default(),
3401                    &test_cache(),
3402                    FileReaderOptions::default(),
3403                )
3404                .await
3405                .unwrap();
3406
3407                let projected_schema = written_file.schema.project(&columns).unwrap();
3408                let projection = if use_field_ids {
3409                    ReaderProjection::from_field_ids(
3410                        file_reader.metadata.version(),
3411                        &projected_schema,
3412                        &field_id_mapping,
3413                    )
3414                    .unwrap()
3415                } else {
3416                    ReaderProjection::from_column_names(
3417                        file_reader.metadata.version(),
3418                        &written_file.schema,
3419                        &columns,
3420                    )
3421                    .unwrap()
3422                };
3423
3424                let batch_stream = file_reader
3425                    .read_stream_projected(
3426                        lance_io::ReadBatchParams::RangeFull,
3427                        1024,
3428                        16,
3429                        projection.clone(),
3430                        FilterExpression::no_filter(),
3431                    )
3432                    .await
3433                    .unwrap();
3434
3435                let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3436                verify_expected(
3437                    &written_file.data,
3438                    batch_stream,
3439                    1024,
3440                    Some(Box::new(move |batch: &RecordBatch| {
3441                        batch.project_by_schema(&projection_arrow).unwrap()
3442                    })),
3443                )
3444                .await;
3445
3446                // We can also specify the projection as a base projection when we open the file
3447                let file_reader = FileReader::try_open(
3448                    file_scheduler.clone(),
3449                    Some(projection.clone()),
3450                    Arc::<DecoderPlugins>::default(),
3451                    &test_cache(),
3452                    FileReaderOptions::default(),
3453                )
3454                .await
3455                .unwrap();
3456
3457                let batch_stream = file_reader
3458                    .read_stream(
3459                        lance_io::ReadBatchParams::RangeFull,
3460                        1024,
3461                        16,
3462                        FilterExpression::no_filter(),
3463                    )
3464                    .await
3465                    .unwrap();
3466
3467                let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3468                verify_expected(
3469                    &written_file.data,
3470                    batch_stream,
3471                    1024,
3472                    Some(Box::new(move |batch: &RecordBatch| {
3473                        batch.project_by_schema(&projection_arrow).unwrap()
3474                    })),
3475                )
3476                .await;
3477
3478                assert!(
3479                    file_reader
3480                        .read_stream_projected(
3481                            lance_io::ReadBatchParams::RangeFull,
3482                            1024,
3483                            16,
3484                            empty_projection.clone(),
3485                            FilterExpression::no_filter(),
3486                        )
3487                        .await
3488                        .is_err()
3489                );
3490            }
3491        }
3492
3493        assert!(
3494            FileReader::try_open(
3495                file_scheduler.clone(),
3496                Some(empty_projection),
3497                Arc::<DecoderPlugins>::default(),
3498                &test_cache(),
3499                FileReaderOptions::default(),
3500            )
3501            .await
3502            .is_err()
3503        );
3504
3505        let arrow_schema = ArrowSchema::new(vec![
3506            Field::new("x", DataType::Int32, true),
3507            Field::new("y", DataType::Int32, true),
3508        ]);
3509        let schema = Schema::try_from(&arrow_schema).unwrap();
3510
3511        let projection_with_dupes = ReaderProjection {
3512            column_indices: vec![0, 0],
3513            schema: Arc::new(schema),
3514        };
3515
3516        assert!(
3517            FileReader::try_open(
3518                file_scheduler.clone(),
3519                Some(projection_with_dupes),
3520                Arc::<DecoderPlugins>::default(),
3521                &test_cache(),
3522                FileReaderOptions::default(),
3523            )
3524            .await
3525            .is_err()
3526        );
3527    }
3528
3529    #[tokio::test]
3530    async fn test_lazy_reader_direct_projection_matches_eager_reader() {
3531        let fs = FsFixture::default();
3532        let written_file = create_wide_direct_file(&fs, 16).await;
3533
3534        let file_scheduler = fs
3535            .scheduler
3536            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3537            .await
3538            .unwrap();
3539        let projection = ReaderProjection::from_column_names(
3540            LanceFileVersion::V2_1,
3541            &written_file.schema,
3542            &["c10"],
3543        )
3544        .unwrap();
3545
3546        let eager_reader = FileReader::try_open(
3547            file_scheduler.clone(),
3548            None,
3549            Arc::<DecoderPlugins>::default(),
3550            &test_cache(),
3551            FileReaderOptions::default(),
3552        )
3553        .await
3554        .unwrap();
3555        let expected = eager_reader
3556            .read_stream_projected(
3557                lance_io::ReadBatchParams::RangeFull,
3558                127,
3559                16,
3560                projection.clone(),
3561                FilterExpression::no_filter(),
3562            )
3563            .await
3564            .unwrap()
3565            .try_collect::<Vec<_>>()
3566            .await
3567            .unwrap();
3568
3569        let cache = test_cache();
3570        let lazy_reader = ProjectedFileReader::try_open(
3571            file_scheduler,
3572            Some(projection.clone()),
3573            Arc::<DecoderPlugins>::default(),
3574            &cache,
3575            FileReaderOptions::default(),
3576        )
3577        .await
3578        .unwrap();
3579        let tasks = lazy_reader
3580            .read_tasks(
3581                lance_io::ReadBatchParams::RangeFull,
3582                127,
3583                None,
3584                FilterExpression::no_filter(),
3585            )
3586            .await
3587            .unwrap();
3588        let actual = collect_read_tasks(tasks, 16).await;
3589
3590        assert_eq!(expected, actual);
3591    }
3592
3593    #[tokio::test]
3594    async fn test_lazy_reader_loads_only_requested_column_metadata() {
3595        let fs = FsFixture::default();
3596        let written_file = create_wide_direct_file(&fs, 512).await;
3597
3598        let projection = ReaderProjection::from_column_names(
3599            LanceFileVersion::V2_1,
3600            &written_file.schema,
3601            &["c0"],
3602        )
3603        .unwrap();
3604        let file_scheduler = fs
3605            .scheduler
3606            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3607            .await
3608            .unwrap();
3609        let lazy_reader = ProjectedFileReader::try_open(
3610            file_scheduler,
3611            Some(projection.clone()),
3612            Arc::<DecoderPlugins>::default(),
3613            &test_cache(),
3614            FileReaderOptions::default(),
3615        )
3616        .await
3617        .unwrap();
3618        let selected_column = projection.column_indices[0] as usize;
3619        let requested_metadata_bytes = lazy_reader
3620            .metadata_index()
3621            .unwrap()
3622            .column_metadata_offsets[selected_column]
3623            .1;
3624        let total_metadata_bytes = lazy_reader
3625            .metadata_index()
3626            .unwrap()
3627            .column_metadata_offsets
3628            .iter()
3629            .map(|(_, length)| *length)
3630            .sum::<u64>();
3631        assert!(
3632            total_metadata_bytes > 8 * fs.object_store.block_size() as u64,
3633            "test file metadata is too small to prove lazy loading: {total_metadata_bytes} bytes"
3634        );
3635
3636        fs.object_store.io_stats_incremental();
3637        let tasks = lazy_reader
3638            .read_tasks(
3639                lance_io::ReadBatchParams::Range(0..0),
3640                1024,
3641                Some(projection.clone()),
3642                FilterExpression::no_filter(),
3643            )
3644            .await
3645            .unwrap();
3646        let batches = collect_read_tasks(tasks, 1).await;
3647        assert!(batches.is_empty());
3648
3649        let stats = fs.object_store.io_stats_incremental();
3650        assert!(
3651            stats.read_bytes < total_metadata_bytes / 2,
3652            "lazy read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3653            stats.read_bytes,
3654            requested_metadata_bytes,
3655            total_metadata_bytes
3656        );
3657
3658        fs.object_store.io_stats_incremental();
3659        let tasks = lazy_reader
3660            .read_tasks(
3661                lance_io::ReadBatchParams::Range(0..0),
3662                1024,
3663                Some(projection),
3664                FilterExpression::no_filter(),
3665            )
3666            .await
3667            .unwrap();
3668        let batches = collect_read_tasks(tasks, 1).await;
3669        assert!(batches.is_empty());
3670
3671        let stats = fs.object_store.io_stats_incremental();
3672        assert_eq!(
3673            stats.read_iops, 0,
3674            "cached column metadata should avoid repeat metadata I/O"
3675        );
3676        assert_eq!(
3677            stats.read_bytes, 0,
3678            "cached column metadata should avoid repeat metadata reads"
3679        );
3680    }
3681
3682    async fn assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3683        fs: &FsFixture,
3684        projection: ReaderProjection,
3685        shape: &str,
3686    ) -> Vec<RecordBatch> {
3687        let file_scheduler = fs
3688            .scheduler
3689            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3690            .await
3691            .unwrap();
3692        let eager_reader = FileReader::try_open(
3693            file_scheduler.clone(),
3694            None,
3695            Arc::<DecoderPlugins>::default(),
3696            &test_cache(),
3697            FileReaderOptions::default(),
3698        )
3699        .await
3700        .unwrap();
3701        let expected = eager_reader
3702            .read_stream_projected(
3703                lance_io::ReadBatchParams::RangeFull,
3704                127,
3705                16,
3706                projection.clone(),
3707                FilterExpression::no_filter(),
3708            )
3709            .await
3710            .unwrap()
3711            .try_collect::<Vec<_>>()
3712            .await
3713            .unwrap();
3714
3715        let cache = test_cache();
3716        let lazy_reader = ProjectedFileReader::try_open(
3717            file_scheduler,
3718            Some(projection.clone()),
3719            Arc::<DecoderPlugins>::default(),
3720            &cache,
3721            FileReaderOptions::default(),
3722        )
3723        .await
3724        .unwrap();
3725        let metadata_index = lazy_reader.metadata_index().unwrap();
3726        let requested_metadata_bytes = projection
3727            .column_indices
3728            .iter()
3729            .map(|column_index| metadata_index.column_metadata_offsets[*column_index as usize].1)
3730            .sum::<u64>();
3731        let total_metadata_bytes = metadata_index
3732            .column_metadata_offsets
3733            .iter()
3734            .map(|(_, length)| *length)
3735            .sum::<u64>();
3736        assert!(total_metadata_bytes > requested_metadata_bytes * 8);
3737
3738        fs.object_store.io_stats_incremental();
3739        let tasks = lazy_reader
3740            .read_tasks(
3741                lance_io::ReadBatchParams::Range(0..0),
3742                127,
3743                None,
3744                FilterExpression::no_filter(),
3745            )
3746            .await
3747            .unwrap();
3748        assert!(collect_read_tasks(tasks, 1).await.is_empty());
3749        let metadata_stats = fs.object_store.io_stats_incremental();
3750        assert!(
3751            metadata_stats.read_bytes < total_metadata_bytes / 2,
3752            "lazy {shape} read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3753            metadata_stats.read_bytes,
3754            requested_metadata_bytes,
3755            total_metadata_bytes
3756        );
3757
3758        let tasks = lazy_reader
3759            .read_tasks(
3760                lance_io::ReadBatchParams::RangeFull,
3761                127,
3762                None,
3763                FilterExpression::no_filter(),
3764            )
3765            .await
3766            .unwrap();
3767        let actual = collect_read_tasks(tasks, 16).await;
3768        assert_eq!(expected, actual);
3769        actual
3770    }
3771
3772    #[tokio::test]
3773    async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() {
3774        let fs = FsFixture::default();
3775        let written_file = create_wide_fixed_size_list_file(&fs, 512).await;
3776        let projection = ReaderProjection::from_column_names(
3777            LanceFileVersion::V2_1,
3778            &written_file.schema,
3779            &["c17", "c509"],
3780        )
3781        .unwrap();
3782        assert!(ProjectedFileReader::supports_projection(
3783            &projection,
3784            LanceFileVersion::V2_1
3785        ));
3786        assert!(!ProjectedFileReader::supports_projection(
3787            &projection,
3788            LanceFileVersion::V2_0
3789        ));
3790        assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3791            &fs,
3792            projection,
3793            "fixed-size-list",
3794        )
3795        .await;
3796    }
3797
3798    #[tokio::test]
3799    async fn test_lazy_reader_nested_projection_compacts_physical_columns() {
3800        let fs = FsFixture::default();
3801        let written_file = create_wide_structural_file(&fs, 128).await;
3802        let projection = ReaderProjection::from_column_names(
3803            LanceFileVersion::V2_1,
3804            &written_file.schema,
3805            &["s97.y", "l4", "s3"],
3806        )
3807        .unwrap();
3808
3809        assert_eq!(
3810            projection
3811                .schema
3812                .fields
3813                .iter()
3814                .map(|field| field.name.as_str())
3815                .collect::<Vec<_>>(),
3816            vec!["s97", "l4", "s3"]
3817        );
3818        assert_eq!(projection.schema.fields[0].children.len(), 1);
3819        assert_eq!(projection.schema.fields[0].children[0].name, "y");
3820        assert_eq!(projection.schema.fields[2].children.len(), 2);
3821        assert_eq!(projection.column_indices.len(), 4);
3822        assert!(
3823            projection
3824                .column_indices
3825                .windows(2)
3826                .any(|indices| indices[0] > indices[1]),
3827            "the projection must reorder physical columns to exercise compact remapping"
3828        );
3829        assert!(ProjectedFileReader::supports_projection(
3830            &projection,
3831            LanceFileVersion::V2_1
3832        ));
3833        let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3834            &fs, projection, "nested",
3835        )
3836        .await;
3837        assert!(
3838            actual
3839                .iter()
3840                .flat_map(|batch| batch.columns())
3841                .any(|column| column.null_count() > 0),
3842            "the structural projection must exercise nullable arrays"
3843        );
3844    }
3845
3846    #[rstest]
3847    #[case::before_metadata_region(90, 5)]
3848    #[case::after_metadata_region(190, 20)]
3849    fn test_decode_cmo_table_rejects_out_of_range_offsets(
3850        #[case] position: u64,
3851        #[case] length: u64,
3852    ) {
3853        let mut cmo_table = [0; 16];
3854        cmo_table[0..8].copy_from_slice(&position.to_le_bytes());
3855        cmo_table[8..16].copy_from_slice(&length.to_le_bytes());
3856        let footer = super::Footer {
3857            column_meta_start: 100,
3858            column_meta_offsets_start: 200,
3859            global_buff_offsets_start: 200,
3860            num_global_buffers: 0,
3861            num_columns: 1,
3862            major_version: 2,
3863            minor_version: 1,
3864        };
3865
3866        let err = FileReader::decode_cmo_table(Bytes::copy_from_slice(&cmo_table), &footer)
3867            .expect_err("out-of-range CMO entries must be rejected");
3868        assert!(
3869            matches!(err, lance_core::Error::InvalidInput { .. }),
3870            "expected InvalidInput, got {err:?}"
3871        );
3872    }
3873
3874    #[rstest]
3875    #[case::blob(BLOB_META_KEY)]
3876    #[case::packed_struct("lance-encoding:packed")]
3877    #[tokio::test]
3878    async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) {
3879        let fs = FsFixture::default();
3880        let written_file = create_some_file(&fs, LanceFileVersion::V2_1).await;
3881
3882        let ordinary_projection = ReaderProjection::from_column_names(
3883            LanceFileVersion::V2_1,
3884            &written_file.schema,
3885            &["location.x"],
3886        )
3887        .unwrap();
3888        assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1);
3889        assert!(ProjectedFileReader::supports_projection(
3890            &ordinary_projection,
3891            LanceFileVersion::V2_1
3892        ));
3893
3894        let file_scheduler = fs
3895            .scheduler
3896            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3897            .await
3898            .unwrap();
3899        let err = ProjectedFileReader::try_open(
3900            file_scheduler.clone(),
3901            None,
3902            Arc::<DecoderPlugins>::default(),
3903            &test_cache(),
3904            FileReaderOptions::default(),
3905        )
3906        .await
3907        .unwrap_err();
3908        assert!(
3909            matches!(err, lance_core::Error::InvalidInput { .. }),
3910            "expected InvalidInput, got {err:?}"
3911        );
3912
3913        let mut projection = ordinary_projection;
3914        Arc::make_mut(&mut projection.schema).fields[0]
3915            .metadata
3916            .insert(metadata_key.to_string(), "true".to_string());
3917        assert!(!ProjectedFileReader::supports_projection(
3918            &projection,
3919            LanceFileVersion::V2_1
3920        ));
3921
3922        let err = ProjectedFileReader::try_open(
3923            file_scheduler,
3924            Some(projection),
3925            Arc::<DecoderPlugins>::default(),
3926            &test_cache(),
3927            FileReaderOptions::default(),
3928        )
3929        .await
3930        .unwrap_err();
3931        assert!(
3932            matches!(err, lance_core::Error::NotSupported { .. }),
3933            "expected NotSupported for {metadata_key}, got {err:?}"
3934        );
3935    }
3936
3937    // The projection-length validation lives in `FileReadCore`, shared by the
3938    // eager and the lazy (indexed) metadata providers. The indexed provider loads
3939    // only the projected columns and renumbers them 0..N, so this checks that the
3940    // renumbered `column_infos`/`column_indices` still line up for the length
3941    // check: a mismatched-length projection is rejected through
3942    // `ProjectedFileReader`, and a single short column resolves to its own length.
3943    #[tokio::test]
3944    async fn test_lazy_reader_validates_unequal_length_projection() {
3945        use arrow_array::Int32Array;
3946        use lance_io::ReadBatchParams;
3947
3948        let arrow_schema = Arc::new(ArrowSchema::new(vec![
3949            Field::new("a", DataType::Int32, true),
3950            Field::new("c", DataType::Int32, true),
3951        ]));
3952        let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
3953
3954        let fs = FsFixture::default();
3955        let options = FileWriterOptions {
3956            format_version: Some(LanceFileVersion::V2_1),
3957            ..Default::default()
3958        };
3959        let mut writer = FileWriter::try_new(
3960            fs.object_store.create(&fs.tmp_path).await.unwrap(),
3961            lance_schema.clone(),
3962            options,
3963        )
3964        .unwrap();
3965        // "a" has 5 rows, "c" has 1 -- an unequal-length file.
3966        writer
3967            .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])))
3968            .await
3969            .unwrap();
3970        writer
3971            .write_column(1, Arc::new(Int32Array::from(vec![100])))
3972            .await
3973            .unwrap();
3974        writer.finish().await.unwrap();
3975
3976        let file_scheduler = fs
3977            .scheduler
3978            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3979            .await
3980            .unwrap();
3981        let cache = test_cache();
3982        let open_indexed = |names: &[&str]| {
3983            let projection =
3984                ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names)
3985                    .unwrap();
3986            ProjectedFileReader::try_open(
3987                file_scheduler.clone(),
3988                Some(projection),
3989                Arc::<DecoderPlugins>::default(),
3990                &cache,
3991                FileReaderOptions::default(),
3992            )
3993        };
3994
3995        // A mismatched-length projection [a, c] (5 vs 1) is rejected at read time,
3996        // through the indexed provider's renumbered column infos.
3997        let lazy = open_indexed(&["a", "c"]).await.unwrap();
3998        // (`read_tasks` yields a stream, which is not `Debug`, so match rather
3999        // than `unwrap_err`.)
4000        let err = match lazy
4001            .read_tasks(
4002                ReadBatchParams::RangeFull,
4003                1024,
4004                None,
4005                FilterExpression::no_filter(),
4006            )
4007            .await
4008        {
4009            Ok(_) => panic!("expected the mismatched-length projection to be rejected"),
4010            Err(e) => e.to_string(),
4011        };
4012        assert!(
4013            err.contains("a=5") && err.contains("c=1"),
4014            "error should name each column's length, got: {err}"
4015        );
4016
4017        // A single short column resolves to its own length (1), not the file's
4018        // longest column.
4019        let lazy = open_indexed(&["c"]).await.unwrap();
4020        let tasks = lazy
4021            .read_tasks(
4022                ReadBatchParams::RangeFull,
4023                1024,
4024                None,
4025                FilterExpression::no_filter(),
4026            )
4027            .await
4028            .unwrap();
4029        let batches = collect_read_tasks(tasks, 16).await;
4030        let values: Vec<Option<i32>> = batches
4031            .iter()
4032            .flat_map(|b| {
4033                b.column(0)
4034                    .as_any()
4035                    .downcast_ref::<Int32Array>()
4036                    .unwrap()
4037                    .iter()
4038                    .collect::<Vec<_>>()
4039            })
4040            .collect();
4041        assert_eq!(values, vec![Some(100)]);
4042    }
4043
4044    #[test_log::test(tokio::test)]
4045    async fn test_compressing_buffer() {
4046        let fs = FsFixture::default();
4047
4048        let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
4049        let file_scheduler = fs
4050            .scheduler
4051            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4052            .await
4053            .unwrap();
4054
4055        // We can specify the projection as part of the read operation via read_stream_projected
4056        let file_reader = FileReader::try_open(
4057            file_scheduler.clone(),
4058            None,
4059            Arc::<DecoderPlugins>::default(),
4060            &test_cache(),
4061            FileReaderOptions::default(),
4062        )
4063        .await
4064        .unwrap();
4065
4066        let mut projection = written_file.schema.project(&["score"]).unwrap();
4067        for field in projection.fields.iter_mut() {
4068            field
4069                .metadata
4070                .insert("lance:compression".to_string(), "zstd".to_string());
4071        }
4072        let projection = ReaderProjection {
4073            column_indices: projection.fields.iter().map(|f| f.id as u32).collect(),
4074            schema: Arc::new(projection),
4075        };
4076
4077        let batch_stream = file_reader
4078            .read_stream_projected(
4079                lance_io::ReadBatchParams::RangeFull,
4080                1024,
4081                16,
4082                projection.clone(),
4083                FilterExpression::no_filter(),
4084            )
4085            .await
4086            .unwrap();
4087
4088        let projection_arrow = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
4089        verify_expected(
4090            &written_file.data,
4091            batch_stream,
4092            1024,
4093            Some(Box::new(move |batch: &RecordBatch| {
4094                batch.project_by_schema(&projection_arrow).unwrap()
4095            })),
4096        )
4097        .await;
4098    }
4099
4100    #[tokio::test]
4101    async fn test_read_all() {
4102        let fs = FsFixture::default();
4103        let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
4104        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4105
4106        let file_scheduler = fs
4107            .scheduler
4108            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4109            .await
4110            .unwrap();
4111        let file_reader = FileReader::try_open(
4112            file_scheduler.clone(),
4113            None,
4114            Arc::<DecoderPlugins>::default(),
4115            &test_cache(),
4116            FileReaderOptions::default(),
4117        )
4118        .await
4119        .unwrap();
4120
4121        let batches = file_reader
4122            .read_stream(
4123                lance_io::ReadBatchParams::RangeFull,
4124                total_rows as u32,
4125                16,
4126                FilterExpression::no_filter(),
4127            )
4128            .await
4129            .unwrap()
4130            .try_collect::<Vec<_>>()
4131            .await
4132            .unwrap();
4133        assert_eq!(batches.len(), 1);
4134        assert_eq!(batches[0].num_rows(), total_rows);
4135    }
4136
4137    #[rstest]
4138    #[tokio::test]
4139    async fn test_blocking_take(
4140        #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
4141        version: LanceFileVersion,
4142    ) {
4143        let fs = FsFixture::default();
4144        let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await;
4145        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4146
4147        let file_scheduler = fs
4148            .scheduler
4149            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4150            .await
4151            .unwrap();
4152        let file_reader = FileReader::try_open(
4153            file_scheduler.clone(),
4154            Some(ReaderProjection::from_column_names(version, &schema, &["score"]).unwrap()),
4155            Arc::<DecoderPlugins>::default(),
4156            &test_cache(),
4157            FileReaderOptions::default(),
4158        )
4159        .await
4160        .unwrap();
4161
4162        let batches = tokio::task::spawn_blocking(move || {
4163            file_reader
4164                .read_stream_projected_blocking(
4165                    lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![0, 1, 2, 3, 4])),
4166                    total_rows as u32,
4167                    None,
4168                    FilterExpression::no_filter(),
4169                )
4170                .unwrap()
4171                .collect::<ArrowResult<Vec<_>>>()
4172                .unwrap()
4173        })
4174        .await
4175        .unwrap();
4176
4177        assert_eq!(batches.len(), 1);
4178        assert_eq!(batches[0].num_rows(), 5);
4179        assert_eq!(batches[0].num_columns(), 1);
4180    }
4181
4182    #[tokio::test(flavor = "multi_thread")]
4183    async fn test_drop_in_progress() {
4184        let fs = FsFixture::default();
4185        let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
4186        let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
4187
4188        let file_scheduler = fs
4189            .scheduler
4190            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4191            .await
4192            .unwrap();
4193        let file_reader = FileReader::try_open(
4194            file_scheduler.clone(),
4195            None,
4196            Arc::<DecoderPlugins>::default(),
4197            &test_cache(),
4198            FileReaderOptions::default(),
4199        )
4200        .await
4201        .unwrap();
4202
4203        let mut batches = file_reader
4204            .read_stream(
4205                lance_io::ReadBatchParams::RangeFull,
4206                (total_rows / 10) as u32,
4207                16,
4208                FilterExpression::no_filter(),
4209            )
4210            .await
4211            .unwrap();
4212
4213        drop(file_reader);
4214
4215        let batch = batches.next().await.unwrap().unwrap();
4216        assert!(batch.num_rows() > 0);
4217
4218        // Drop in-progress scan
4219        drop(batches);
4220    }
4221
4222    #[tokio::test]
4223    async fn drop_while_scheduling() {
4224        // This is a bit of a white-box test, pokes at the internals.  We want to
4225        // test the case where the read stream is dropped before the scheduling
4226        // thread finishes.  We can't do that in a black-box fashion because the
4227        // scheduling thread runs in the background and there is no easy way to
4228        // pause / gate it.
4229
4230        // It's a regression for a bug where the scheduling thread would panic
4231        // if the stream was dropped before it finished.
4232
4233        let fs = FsFixture::default();
4234        let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
4235        let total_rows = written_file
4236            .data
4237            .iter()
4238            .map(|batch| batch.num_rows())
4239            .sum::<usize>();
4240
4241        let file_scheduler = fs
4242            .scheduler
4243            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4244            .await
4245            .unwrap();
4246        let file_reader = FileReader::try_open(
4247            file_scheduler.clone(),
4248            None,
4249            Arc::<DecoderPlugins>::default(),
4250            &test_cache(),
4251            FileReaderOptions::default(),
4252        )
4253        .await
4254        .unwrap();
4255
4256        let projection =
4257            ReaderProjection::from_whole_schema(&written_file.schema, LanceFileVersion::V2_0);
4258        let column_infos = file_reader
4259            .collect_columns_from_projection(&projection)
4260            .unwrap();
4261        let mut decode_scheduler = DecodeBatchScheduler::try_new(
4262            &projection.schema,
4263            &projection.column_indices,
4264            &column_infos,
4265            &vec![],
4266            total_rows as u64,
4267            Arc::<DecoderPlugins>::default(),
4268            file_reader.core.scheduler.clone(),
4269            test_cache(),
4270            &FilterExpression::no_filter(),
4271            &DecoderConfig::default(),
4272        )
4273        .await
4274        .unwrap();
4275
4276        let range = 0..total_rows as u64;
4277
4278        let (tx, rx) = mpsc::unbounded_channel();
4279
4280        // Simulate the stream / decoder being dropped
4281        drop(rx);
4282
4283        // Scheduling should not panic
4284        decode_scheduler.schedule_range(
4285            range,
4286            &FilterExpression::no_filter(),
4287            tx,
4288            file_reader.core.scheduler.clone(),
4289        )
4290    }
4291
4292    #[tokio::test]
4293    async fn test_read_empty_range() {
4294        let fs = FsFixture::default();
4295        create_some_file(&fs, LanceFileVersion::V2_0).await;
4296
4297        let file_scheduler = fs
4298            .scheduler
4299            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4300            .await
4301            .unwrap();
4302        let file_reader = FileReader::try_open(
4303            file_scheduler.clone(),
4304            None,
4305            Arc::<DecoderPlugins>::default(),
4306            &test_cache(),
4307            FileReaderOptions::default(),
4308        )
4309        .await
4310        .unwrap();
4311
4312        // All ranges empty, no data
4313        let batches = file_reader
4314            .read_stream(
4315                lance_io::ReadBatchParams::Range(0..0),
4316                1024,
4317                16,
4318                FilterExpression::no_filter(),
4319            )
4320            .await
4321            .unwrap()
4322            .try_collect::<Vec<_>>()
4323            .await
4324            .unwrap();
4325
4326        assert_eq!(batches.len(), 0);
4327
4328        // Some ranges empty
4329        let batches = file_reader
4330            .read_stream(
4331                lance_io::ReadBatchParams::Ranges(Arc::new([0..1, 2..2])),
4332                1024,
4333                16,
4334                FilterExpression::no_filter(),
4335            )
4336            .await
4337            .unwrap()
4338            .try_collect::<Vec<_>>()
4339            .await
4340            .unwrap();
4341        assert_eq!(batches.len(), 1);
4342    }
4343
4344    async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) {
4345        let lance_schema =
4346            lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new(
4347                "foo",
4348                DataType::Int32,
4349                true,
4350            )]))
4351            .unwrap();
4352
4353        let mut file_writer = FileWriter::try_new(
4354            fs.object_store.create(&fs.tmp_path).await.unwrap(),
4355            lance_schema,
4356            FileWriterOptions::default(),
4357        )
4358        .unwrap();
4359
4360        let buf_index = file_writer.add_global_buffer(buffer).await.unwrap();
4361        assert_eq!(buf_index, 1);
4362
4363        file_writer.finish().await.unwrap();
4364    }
4365
4366    /// A global buffer that fits inside the tail region captured at open is served
4367    /// from memory with no additional I/O.  A buffer larger than that window cannot
4368    /// fit and falls back to a dedicated read.  Both must round-trip correctly.
4369    #[rstest]
4370    #[case::within_tail_window(true)]
4371    #[case::outside_tail_window(false)]
4372    #[tokio::test]
4373    async fn test_read_global_buffer(#[case] within_window: bool) {
4374        let fs = FsFixture::default();
4375
4376        let block_size = fs.object_store.block_size();
4377        let buffer = if within_window {
4378            Bytes::from_static(b"hello")
4379        } else {
4380            Bytes::from(vec![7u8; 2 * block_size])
4381        };
4382        let expected_read_iops = if within_window { 0 } else { 1 };
4383
4384        write_file_with_global_buffer(&fs, buffer.clone()).await;
4385
4386        let file_scheduler = fs
4387            .scheduler
4388            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4389            .await
4390            .unwrap();
4391        let file_reader = FileReader::try_open(
4392            file_scheduler,
4393            None,
4394            Arc::<DecoderPlugins>::default(),
4395            &test_cache(),
4396            FileReaderOptions::default(),
4397        )
4398        .await
4399        .unwrap();
4400
4401        // The user buffer should be retained only when it fits the tail window, and
4402        // the schema (buffer 0) is never retained.
4403        let retained = &file_reader.metadata().retained_global_buffers;
4404        assert!(!retained.contains_key(&0), "schema must not be retained");
4405        assert_eq!(retained.contains_key(&1), within_window);
4406
4407        // Reset the IO counters so we only measure the read_global_buffer call.
4408        fs.object_store.io_stats_incremental();
4409
4410        let buf = file_reader.read_global_buffer(1).await.unwrap();
4411        assert_eq!(buf, buffer);
4412
4413        let stats = fs.object_store.io_stats_incremental();
4414        assert_eq!(stats.read_iops, expected_read_iops);
4415    }
4416
4417    /// A file whose only global buffer is the schema (i.e. a plain data file, the
4418    /// common case) must retain nothing — there is no user buffer to serve.
4419    #[tokio::test]
4420    async fn test_read_global_buffer_no_user_buffers() {
4421        let fs = FsFixture::default();
4422        create_some_file(&fs, LanceFileVersion::V2_1).await;
4423
4424        let file_scheduler = fs
4425            .scheduler
4426            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4427            .await
4428            .unwrap();
4429        let file_reader = FileReader::try_open(
4430            file_scheduler,
4431            None,
4432            Arc::<DecoderPlugins>::default(),
4433            &test_cache(),
4434            FileReaderOptions::default(),
4435        )
4436        .await
4437        .unwrap();
4438
4439        let metadata = file_reader.metadata();
4440        assert_eq!(metadata.file_buffers.len(), 1, "expected only the schema");
4441        assert!(
4442            metadata.retained_global_buffers.is_empty(),
4443            "a file with no user global buffers must retain nothing"
4444        );
4445    }
4446
4447    #[rstest]
4448    #[tokio::test]
4449    async fn test_deep_size_of_includes_column_metadata(
4450        #[values(
4451            LanceFileVersion::V2_0,
4452            LanceFileVersion::V2_1,
4453            LanceFileVersion::V2_2,
4454            LanceFileVersion::V2_3
4455        )]
4456        version: LanceFileVersion,
4457    ) {
4458        // Regression test: CachedFileMetadata::deep_size_of must account for
4459        // column_metadatas and column_infos, otherwise the moka cache weigher
4460        // dramatically underestimates entry sizes and never evicts, causing
4461        // unbounded memory growth on random-access workloads.
4462        use lance_core::deepsize::DeepSizeOf;
4463
4464        let fs = FsFixture::default();
4465        let _written = create_some_file(&fs, version).await;
4466        let cache = test_cache();
4467        let file_scheduler = fs
4468            .scheduler
4469            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4470            .await
4471            .unwrap();
4472        let file_reader = FileReader::try_open(
4473            file_scheduler,
4474            None,
4475            Arc::<DecoderPlugins>::default(),
4476            &cache,
4477            FileReaderOptions::default(),
4478        )
4479        .await
4480        .unwrap();
4481
4482        let metadata = file_reader.metadata();
4483        let deep_size = metadata.deep_size_of();
4484
4485        // The file has multiple columns (score, location, categories, binary,
4486        // maybe large_bin). The deep_size_of must be substantially more than
4487        // just the schema — it should include column_metadatas + column_infos.
4488        // A naive implementation that ignores these fields reports < 1 KB;
4489        // a correct one should report at least several KB for this test file.
4490        assert!(
4491            deep_size > 1024,
4492            "deep_size_of ({deep_size}) is suspiciously small — \
4493             column_metadatas and column_infos may not be accounted for"
4494        );
4495
4496        // Verify column_metadatas is non-empty (sanity check).
4497        assert!(
4498            !metadata.column_metadatas.is_empty(),
4499            "Expected non-empty column_metadatas"
4500        );
4501
4502        // Verify the size scales with the number of columns: a file with more
4503        // columns should have a larger deep_size_of.
4504        let num_columns = metadata.column_metadatas.len();
4505        assert!(
4506            deep_size > num_columns * 50,
4507            "deep_size_of ({deep_size}) should scale with column count ({num_columns})"
4508        );
4509    }
4510
4511    #[tokio::test]
4512    async fn test_read_global_buffer_out_of_range() {
4513        let fs = FsFixture::default();
4514
4515        write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
4516
4517        let file_scheduler = fs
4518            .scheduler
4519            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4520            .await
4521            .unwrap();
4522        let file_reader = FileReader::try_open(
4523            file_scheduler,
4524            None,
4525            Arc::<DecoderPlugins>::default(),
4526            &test_cache(),
4527            FileReaderOptions::default(),
4528        )
4529        .await
4530        .unwrap();
4531
4532        // The file has two global buffers (schema at 0, "hello" at 1); index 2 is
4533        // out of range and must surface a descriptive error rather than panicking.
4534        let err = file_reader.read_global_buffer(2).await.unwrap_err();
4535        assert!(
4536            matches!(err, lance_core::Error::InvalidInput { .. }),
4537            "expected InvalidInput, got: {err:?}"
4538        );
4539        let msg = err.to_string();
4540        assert!(msg.contains('2'), "error should mention the index: {msg}");
4541    }
4542
4543    // Exercises the projection length-validation walk in isolation, feeding
4544    // synthetic per-column lengths so we can reach cases no current writer can
4545    // actually produce -- in particular a struct whose children diverge in
4546    // length, which the decoders would otherwise panic on or misread.
4547    #[rstest]
4548    fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) {
4549        let run = |dt: DataType, indices: &[u32], lengths: Vec<u64>| -> lance_core::Result<u64> {
4550            let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]);
4551            let schema = Schema::try_from(&arrow).unwrap();
4552            let column_len = |c: usize| Ok(lengths[c]);
4553            let mut cursor = 0usize;
4554            let mut field_lengths = Vec::new();
4555            for field in &schema.fields {
4556                let rows = validate_field_length(
4557                    field,
4558                    is_structural,
4559                    true,
4560                    indices,
4561                    &mut cursor,
4562                    &column_len,
4563                )?;
4564                field_lengths.push((field.name.as_str(), rows));
4565            }
4566            verify_uniform_lengths(&field_lengths)
4567        };
4568
4569        let struct_ty = || {
4570            DataType::Struct(Fields::from(vec![
4571                Field::new("a", DataType::Int32, true),
4572                Field::new("b", DataType::Int32, true),
4573            ]))
4574        };
4575
4576        // In 2.1 a struct contributes no column of its own (just its two leaves);
4577        // in 2.0 it also has its own column first.
4578        let (indices, equal, unequal): (&[u32], Vec<u64>, Vec<u64>) = if is_structural {
4579            (&[0, 1], vec![5, 5], vec![5, 3])
4580        } else {
4581            (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3])
4582        };
4583
4584        assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5);
4585
4586        let err = run(struct_ty(), indices, unequal).unwrap_err();
4587        let msg = err.to_string();
4588        assert!(
4589            msg.contains("differing lengths") && msg.contains('b'),
4590            "expected a child-length error naming 'b', got: {msg}"
4591        );
4592    }
4593
4594    #[test]
4595    fn test_validate_v2_0_unloaded_blob_projection_is_opaque() {
4596        let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]);
4597        let arrow = ArrowSchema::new(vec![
4598            Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata),
4599        ]);
4600        let mut schema = Schema::try_from(&arrow).unwrap();
4601        schema.fields[0].unloaded_mut();
4602        let projection = ReaderProjection {
4603            schema: Arc::new(schema),
4604            column_indices: vec![0],
4605        };
4606        let column_len = |column: usize| {
4607            assert_eq!(column, 0);
4608            Ok(3)
4609        };
4610        let mut cursor = 0usize;
4611
4612        let rows = validate_field_length(
4613            &projection.schema.fields[0],
4614            false,
4615            true,
4616            &projection.column_indices,
4617            &mut cursor,
4618            &column_len,
4619        )
4620        .unwrap();
4621
4622        assert_eq!(rows, 3);
4623        assert_eq!(cursor, 1);
4624    }
4625
4626    #[test]
4627    fn test_validate_length_list_and_empty_struct() {
4628        let validate = |dt: DataType,
4629                        is_structural: bool,
4630                        indices: &[u32],
4631                        lengths: Vec<u64>|
4632         -> lance_core::Result<u64> {
4633            let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]);
4634            let schema = Schema::try_from(&arrow).unwrap();
4635            let column_len = |c: usize| Ok(lengths[c]);
4636            let mut cursor = 0usize;
4637            validate_field_length(
4638                &schema.fields[0],
4639                is_structural,
4640                true,
4641                indices,
4642                &mut cursor,
4643                &column_len,
4644            )
4645        };
4646
4647        // A list's items have a different cardinality than its rows; that gap
4648        // must not be flagged as a mismatch. In 2.0 the list is an offsets column
4649        // (rows) plus an items column (item count); in 2.1 it is a single column
4650        // whose page rows are the top-level row count.
4651        let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
4652        assert_eq!(
4653            validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(),
4654            5
4655        );
4656        assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5);
4657
4658        // A list of structs: the struct's children sit below the list boundary, so
4659        // their (item-count) lengths must NOT be compared against the list's row
4660        // count. In 2.0 each field has a column [list, struct, a, b] = [6, 6, 29,
4661        // 29]; the list resolves to its own row count (6) and the struct's longer
4662        // children are not flagged. (Regression: this previously errored because
4663        // the nested struct's children were checked against the struct's count.)
4664        let list_of_struct = DataType::List(Arc::new(Field::new(
4665            "item",
4666            DataType::Struct(Fields::from(vec![
4667                Field::new("a", DataType::Int32, true),
4668                Field::new("b", DataType::Int32, true),
4669            ])),
4670            true,
4671        )));
4672        assert_eq!(
4673            validate(
4674                list_of_struct.clone(),
4675                false,
4676                &[0, 1, 2, 3],
4677                vec![6, 6, 29, 29]
4678            )
4679            .unwrap(),
4680            6
4681        );
4682        // In 2.1 only the two leaves carry columns; still no false mismatch.
4683        assert_eq!(
4684            validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(),
4685            29
4686        );
4687
4688        // An empty struct contributes a single column and validates to its length.
4689        let empty_struct = DataType::Struct(Fields::empty());
4690        assert_eq!(
4691            validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(),
4692            9
4693        );
4694        assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9);
4695    }
4696}