1use 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, 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 writer::PAGE_BUFFER_ALIGNMENT,
51};
52
53pub const DEFAULT_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024;
56
57#[derive(Debug, DeepSizeOf)]
62pub struct BufferDescriptor {
63 pub position: u64,
64 pub size: u64,
65}
66
67#[derive(Debug)]
69pub struct FileStatistics {
70 pub columns: Vec<ColumnStatistics>,
72}
73
74#[derive(Debug)]
76pub struct ColumnStatistics {
77 pub num_pages: usize,
79 pub size_bytes: u64,
83}
84
85#[derive(Debug)]
87pub struct CachedFileMetadata {
88 pub file_schema: Arc<Schema>,
90 pub column_metadatas: Vec<pbfile::ColumnMetadata>,
92 pub column_infos: Vec<Arc<ColumnInfo>>,
93 pub num_rows: u64,
95 pub file_buffers: Vec<BufferDescriptor>,
96 pub num_data_bytes: u64,
98 pub num_column_metadata_bytes: u64,
101 pub num_global_buffer_bytes: u64,
103 pub num_footer_bytes: u64,
105 pub major_version: u16,
106 pub minor_version: u16,
107 pub file_size_bytes: u64,
109 pub retained_global_buffers: BTreeMap<u32, Bytes>,
122}
123
124impl CachedFileMetadata {
125 pub fn file_size(&self) -> u64 {
127 self.file_size_bytes
128 }
129}
130
131fn column_metadata_deep_size(column_metadatas: &[pbfile::ColumnMetadata]) -> usize {
132 column_metadatas
133 .iter()
134 .map(|cm| cm.encoded_len() * 4)
135 .sum::<usize>()
136 + std::mem::size_of_val(column_metadatas)
137}
138
139impl DeepSizeOf for CachedFileMetadata {
140 fn deep_size_of_children(&self, context: &mut Context) -> usize {
141 let schema_size = self.file_schema.deep_size_of_children(context);
142
143 let buffers_size: usize = self
144 .file_buffers
145 .iter()
146 .map(|fb| fb.deep_size_of_children(context))
147 .sum();
148
149 let column_metadatas_size = column_metadata_deep_size(self.column_metadatas.as_slice());
155
156 let column_infos_size = self.column_infos.deep_size_of_children(context);
160
161 let retained_buffers_size = self.retained_global_buffers.deep_size_of_children(context);
163
164 schema_size
165 + buffers_size
166 + column_metadatas_size
167 + column_infos_size
168 + retained_buffers_size
169 }
170}
171
172#[derive(Debug, DeepSizeOf)]
178pub struct FileMetadataIndex {
179 file_schema: Arc<Schema>,
180 num_rows: u64,
181 file_buffers: Vec<BufferDescriptor>,
182 column_metadata_offsets: Arc<[(u64, u64)]>,
183 num_columns: u32,
184 version: LanceFileVersion,
185 file_size_bytes: u64,
186 retained_global_buffers: BTreeMap<u32, Bytes>,
187}
188
189impl FileMetadataIndex {
190 pub fn file_size(&self) -> u64 {
192 self.file_size_bytes
193 }
194
195 pub fn num_columns(&self) -> u32 {
197 self.num_columns
198 }
199}
200
201#[derive(Debug)]
202struct CachedColumnMetadata {
203 column_metadata: pbfile::ColumnMetadata,
204 column_info: Arc<ColumnInfo>,
205}
206
207impl DeepSizeOf for CachedColumnMetadata {
208 fn deep_size_of_children(&self, context: &mut Context) -> usize {
209 column_metadata_deep_size(std::slice::from_ref(&self.column_metadata))
210 + self.column_info.deep_size_of_children(context)
211 }
212}
213
214#[derive(Debug, Clone)]
215struct ColumnMetadataCacheKey {
216 column_index: u32,
217}
218
219impl CacheKey for ColumnMetadataCacheKey {
220 type ValueType = CachedColumnMetadata;
221
222 fn key(&self) -> Cow<'_, str> {
223 Cow::Owned(format!("column_metadata/{}", self.column_index))
224 }
225
226 fn type_name() -> &'static str {
227 "ColumnMetadata"
228 }
229}
230
231impl CachedFileMetadata {
232 pub fn version(&self) -> LanceFileVersion {
233 match (self.major_version, self.minor_version) {
234 (0, 3) => LanceFileVersion::V2_0,
235 (2, 0) => LanceFileVersion::V2_0,
236 (2, 1) => LanceFileVersion::V2_1,
237 (2, 2) => LanceFileVersion::V2_2,
238 (2, 3) => LanceFileVersion::V2_3,
239 _ => panic!(
240 "Unsupported version: {}.{}",
241 self.major_version, self.minor_version
242 ),
243 }
244 }
245}
246
247#[derive(Debug, Clone)]
269pub struct ReaderProjection {
270 pub schema: Arc<Schema>,
273 pub column_indices: Vec<u32>,
313}
314
315impl ReaderProjection {
316 fn from_field_ids_helper<'a>(
317 file_version: LanceFileVersion,
318 fields: impl Iterator<Item = &'a Field>,
319 field_id_to_column_index: &BTreeMap<u32, u32>,
320 column_indices: &mut Vec<u32>,
321 ) -> Result<()> {
322 for field in fields {
323 let is_structural = file_version >= LanceFileVersion::V2_1;
324 let (contributes, recurse) = field_column_shape(field, is_structural);
325 if contributes
328 && let Some(column_idx) = field_id_to_column_index.get(&(field.id as u32)).copied()
329 {
330 column_indices.push(column_idx);
331 }
332 if recurse {
333 Self::from_field_ids_helper(
334 file_version,
335 field.children.iter(),
336 field_id_to_column_index,
337 column_indices,
338 )?;
339 }
340 }
341 Ok(())
342 }
343
344 pub fn from_field_ids(
349 file_version: LanceFileVersion,
350 schema: &Schema,
351 field_id_to_column_index: &BTreeMap<u32, u32>,
352 ) -> Result<Self> {
353 let mut column_indices = Vec::new();
354 Self::from_field_ids_helper(
355 file_version,
356 schema.fields.iter(),
357 field_id_to_column_index,
358 &mut column_indices,
359 )?;
360 let projection = Self {
361 schema: Arc::new(schema.clone()),
362 column_indices,
363 };
364 Ok(projection)
365 }
366
367 pub fn from_whole_schema(schema: &Schema, version: LanceFileVersion) -> Self {
375 let schema = Arc::new(schema.clone());
376 let is_structural = version >= LanceFileVersion::V2_1;
377 let mut column_indices = vec![];
378 let mut curr_column_idx = 0;
379 let mut packed_struct_fields_num = 0;
380 for field in schema.fields_pre_order() {
381 if packed_struct_fields_num > 0 {
382 packed_struct_fields_num -= 1;
383 continue;
384 }
385 if field.is_packed_struct() {
386 column_indices.push(curr_column_idx);
387 curr_column_idx += 1;
388 packed_struct_fields_num = field.children.len();
389 } else if field.children.is_empty() || !is_structural {
390 column_indices.push(curr_column_idx);
391 curr_column_idx += 1;
392 }
393 }
394 Self {
395 schema,
396 column_indices,
397 }
398 }
399
400 pub fn from_column_names(
407 file_version: LanceFileVersion,
408 schema: &Schema,
409 column_names: &[&str],
410 ) -> Result<Self> {
411 let field_id_to_column_index = schema
412 .fields_pre_order()
413 .filter(|field| {
416 file_version < LanceFileVersion::V2_1 || field.is_leaf() || field.is_packed_struct()
417 })
418 .enumerate()
419 .map(|(idx, field)| (field.id as u32, idx as u32))
420 .collect::<BTreeMap<_, _>>();
421 let projected = schema.project(column_names)?;
422 let mut column_indices = Vec::new();
423 Self::from_field_ids_helper(
424 file_version,
425 projected.fields.iter(),
426 &field_id_to_column_index,
427 &mut column_indices,
428 )?;
429 Ok(Self {
430 schema: Arc::new(projected),
431 column_indices,
432 })
433 }
434}
435
436#[derive(Clone, Debug)]
438pub struct FileReaderOptions {
439 pub decoder_config: DecoderConfig,
440 pub read_chunk_size: u64,
444 pub batch_size_bytes: Option<u64>,
451}
452
453impl Default for FileReaderOptions {
454 fn default() -> Self {
455 Self {
456 decoder_config: DecoderConfig::default(),
457 read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
458 batch_size_bytes: None,
459 }
460 }
461}
462
463#[derive(Debug, Clone)]
464struct PreparedProjection {
465 column_infos: Vec<Arc<ColumnInfo>>,
466 decoder_projection: ReaderProjection,
467}
468
469#[derive(Debug, Clone)]
470enum FileMetadataProvider {
471 Full(Arc<CachedFileMetadata>),
472 Indexed(Arc<FileMetadataIndex>),
473}
474
475#[derive(Debug, Clone)]
476struct FileReadCore {
477 scheduler: Arc<dyn EncodingsIo>,
478 base_projection: ReaderProjection,
479 metadata_provider: FileMetadataProvider,
480 decoder_plugins: Arc<DecoderPlugins>,
481 cache: Arc<LanceCache>,
482 options: FileReaderOptions,
483}
484
485#[derive(Debug, Clone)]
493pub struct ProjectedFileReader {
494 core: FileReadCore,
495}
496
497#[derive(Debug, Clone)]
499pub struct FileReader {
500 core: FileReadCore,
501 metadata: Arc<CachedFileMetadata>,
502}
503#[derive(Debug)]
504struct Footer {
505 #[allow(dead_code)]
506 column_meta_start: u64,
507 #[allow(dead_code)]
510 column_meta_offsets_start: u64,
511 global_buff_offsets_start: u64,
512 num_global_buffers: u32,
513 num_columns: u32,
514 major_version: u16,
515 minor_version: u16,
516}
517
518const FOOTER_LEN: usize = 40;
519
520fn field_column_shape(field: &Field, is_structural: bool) -> (bool, bool) {
530 if field.is_blob() || field.is_packed_struct() {
531 return (true, false);
532 }
533
534 let contributes = !is_structural || field.children.is_empty();
535 let recurse = !field.children.is_empty();
536 (contributes, recurse)
537}
538
539fn indexed_projection_column_count(field: &Field) -> Option<usize> {
548 if field.is_blob() || field.is_packed_struct() {
549 return None;
550 }
551
552 let (contributes, recurse) = field_column_shape(field, true);
553 let initial = usize::from(contributes);
554 if !recurse {
555 return Some(initial);
556 }
557
558 field.children.iter().try_fold(initial, |count, child| {
559 count.checked_add(indexed_projection_column_count(child)?)
560 })
561}
562
563fn children_share_parent_length(field: &Field) -> bool {
568 field.logical_type.is_struct()
569}
570
571fn validate_field_length<F: Fn(usize) -> Result<u64>>(
584 field: &Field,
585 is_structural: bool,
586 comparable: bool,
587 column_indices: &[u32],
588 cursor: &mut usize,
589 column_len: &F,
590) -> Result<u64> {
591 let (contributes, recurse) = field_column_shape(field, is_structural);
592 let mut field_rows: Option<u64> = None;
593 if contributes {
594 let column = *column_indices.get(*cursor).ok_or_else(|| {
595 Error::invalid_input(format!(
596 "projection supplied fewer column indices than its fields require \
597 (ran out at field '{}')",
598 field.name
599 ))
600 })?;
601 *cursor += 1;
602 field_rows = Some(column_len(column as usize)?);
603 }
604 if recurse {
605 let enforce_children = comparable && children_share_parent_length(field);
610 for child in &field.children {
611 let child_rows = validate_field_length(
612 child,
613 is_structural,
614 enforce_children,
615 column_indices,
616 cursor,
617 column_len,
618 )?;
619 let expected = *field_rows.get_or_insert(child_rows);
622 if enforce_children && child_rows != expected {
623 return Err(Error::invalid_input(format!(
624 "cannot read field '{}': its children have differing lengths \
625 (child '{}' has {} rows, but the field has {}); a struct's \
626 children must all have the same length",
627 field.name, child.name, child_rows, expected
628 )));
629 }
630 }
631 }
632 field_rows.ok_or_else(|| {
633 Error::invalid_input(format!(
634 "projected field '{}' maps to no columns",
635 field.name
636 ))
637 })
638}
639
640fn verify_uniform_lengths(field_lengths: &[(&str, u64)]) -> Result<u64> {
646 let first = field_lengths.first().map_or(0, |&(_, len)| len);
647 if field_lengths.iter().all(|&(_, len)| len == first) {
648 return Ok(first);
649 }
650 let columns = field_lengths
651 .iter()
652 .map(|(name, len)| format!("{name}={len}"))
653 .collect::<Vec<_>>()
654 .join(", ");
655 Err(Error::invalid_input(format!(
656 "cannot read columns of differing lengths together ({columns}); \
657 read each column (or equal-length group) separately"
658 )))
659}
660
661impl FileReader {
662 pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
663 Self {
664 core: self.core.with_scheduler(scheduler),
665 metadata: self.metadata.clone(),
666 }
667 }
668
669 pub fn with_io_stats(
677 &self,
678 stats: Arc<dyn lance_core::utils::io_stats::IoStatsRecorder>,
679 ) -> Self {
680 match self.core.scheduler.with_io_stats(stats) {
681 Some(scheduler) => self.with_scheduler(scheduler),
682 None => self.clone(),
683 }
684 }
685
686 pub fn num_rows(&self) -> u64 {
687 self.core.num_rows()
688 }
689
690 pub fn column_num_rows(&self, column_index: usize) -> Result<u64> {
699 let column = self
700 .metadata
701 .column_metadatas
702 .get(column_index)
703 .ok_or_else(|| {
704 Error::invalid_input(format!(
705 "column index {} is out of bounds (file has {} columns)",
706 column_index,
707 self.metadata.column_metadatas.len()
708 ))
709 })?;
710 Ok(column.pages.iter().map(|page| page.length).sum())
711 }
712
713 pub fn metadata(&self) -> &Arc<CachedFileMetadata> {
714 &self.metadata
715 }
716
717 fn statistics_from_column_metadata(
718 column_metadatas: &[pbfile::ColumnMetadata],
719 ) -> FileStatistics {
720 let column_stats = column_metadatas
721 .iter()
722 .map(|col_metadata| {
723 let num_pages = col_metadata.pages.len();
724 let size_bytes = col_metadata
725 .pages
726 .iter()
727 .map(|page| page.buffer_sizes.iter().sum::<u64>())
728 .sum::<u64>();
729 ColumnStatistics {
730 num_pages,
731 size_bytes,
732 }
733 })
734 .collect();
735
736 FileStatistics {
737 columns: column_stats,
738 }
739 }
740
741 pub fn file_statistics(&self) -> FileStatistics {
742 Self::statistics_from_column_metadata(&self.metadata().column_metadatas)
743 }
744
745 pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
746 self.core.read_global_buffer(index).await
747 }
748
749 async fn read_tail(scheduler: &FileScheduler) -> Result<(Bytes, u64)> {
750 let file_size = scheduler.reader().size().await? as u64;
751 let begin = if file_size < scheduler.reader().block_size() as u64 {
752 0
753 } else {
754 file_size - scheduler.reader().block_size() as u64
755 };
756 let tail_bytes = scheduler.submit_single(begin..file_size, 0).await?;
757 Ok((tail_bytes, file_size))
758 }
759
760 async fn read_range_from_tail_or_scheduler(
761 tail_bytes: &Bytes,
762 tail_offset: u64,
763 scheduler: &FileScheduler,
764 range: Range<u64>,
765 ) -> Result<Bytes> {
766 let tail_end = tail_offset + tail_bytes.len() as u64;
767 if range.start >= tail_offset && range.end <= tail_end {
768 let rel_start = (range.start - tail_offset) as usize;
769 let rel_end = (range.end - tail_offset) as usize;
770 Ok(tail_bytes.slice(rel_start..rel_end))
771 } else {
772 scheduler.submit_single(range, 0).await
773 }
774 }
775
776 fn retained_global_buffers_from_tail(
777 gbo_table: &[BufferDescriptor],
778 tail_bytes: &Bytes,
779 tail_offset: u64,
780 ) -> BTreeMap<u32, Bytes> {
781 let tail_end = tail_offset + tail_bytes.len() as u64;
782 gbo_table
783 .iter()
784 .enumerate()
785 .skip(1)
786 .filter_map(|(index, buffer)| {
787 let start = buffer.position;
788 let end = buffer.position + buffer.size;
789 if start >= tail_offset && end <= tail_end {
790 let rel_start = (start - tail_offset) as usize;
791 let rel_end = (end - tail_offset) as usize;
792 let bytes = Bytes::copy_from_slice(&tail_bytes[rel_start..rel_end]);
793 Some((index as u32, bytes))
794 } else {
795 None
796 }
797 })
798 .collect()
799 }
800
801 fn decode_footer(footer_bytes: &Bytes) -> Result<Footer> {
804 let len = footer_bytes.len();
805 if len < FOOTER_LEN {
806 return Err(Error::invalid_input(format!(
807 "does not have sufficient data, len: {}, bytes: {:?}",
808 len, footer_bytes
809 )));
810 }
811 let mut cursor = Cursor::new(footer_bytes.slice(len - FOOTER_LEN..));
812
813 let column_meta_start = cursor.read_u64::<LittleEndian>()?;
814 let column_meta_offsets_start = cursor.read_u64::<LittleEndian>()?;
815 let global_buff_offsets_start = cursor.read_u64::<LittleEndian>()?;
816 let num_global_buffers = cursor.read_u32::<LittleEndian>()?;
817 let num_columns = cursor.read_u32::<LittleEndian>()?;
818 let major_version = cursor.read_u16::<LittleEndian>()?;
819 let minor_version = cursor.read_u16::<LittleEndian>()?;
820
821 if major_version == MAJOR_VERSION as u16 && minor_version == MINOR_VERSION as u16 {
822 return Err(Error::version_conflict(
823 "Attempt to use the lance v2 reader to read a legacy file".to_string(),
824 major_version,
825 minor_version,
826 ));
827 }
828
829 let magic_bytes = footer_bytes.slice(len - 4..);
830 if magic_bytes.as_ref() != MAGIC {
831 return Err(Error::invalid_input(format!(
832 "file does not appear to be a Lance file (invalid magic: {:?})",
833 MAGIC
834 )));
835 }
836 Ok(Footer {
837 column_meta_start,
838 column_meta_offsets_start,
839 global_buff_offsets_start,
840 num_global_buffers,
841 num_columns,
842 major_version,
843 minor_version,
844 })
845 }
846
847 fn read_all_column_metadata(
849 column_metadata_bytes: Bytes,
850 footer: &Footer,
851 ) -> Result<Vec<pbfile::ColumnMetadata>> {
852 let column_metadata_start = footer.column_meta_start;
853 let cmo_table_size = 16 * footer.num_columns as usize;
855 if column_metadata_bytes.len() < cmo_table_size {
856 return Err(Error::invalid_input(format!(
857 "column metadata region has {} bytes but CMO table needs {} bytes for {} columns",
858 column_metadata_bytes.len(),
859 cmo_table_size,
860 footer.num_columns
861 )));
862 }
863 let cmo_table = column_metadata_bytes.slice(column_metadata_bytes.len() - cmo_table_size..);
864 let column_metadata_offsets = Self::decode_cmo_table(cmo_table, footer)?;
865
866 column_metadata_offsets
867 .iter()
868 .map(|(position, length)| {
869 let normalized_position = (*position - column_metadata_start) as usize;
870 let normalized_end = normalized_position + (*length as usize);
871 Ok(pbfile::ColumnMetadata::decode(
872 &column_metadata_bytes[normalized_position..normalized_end],
873 )?)
874 })
875 .collect::<Result<Vec<_>>>()
876 }
877
878 fn decode_cmo_table(cmo_table: Bytes, footer: &Footer) -> Result<Arc<[(u64, u64)]>> {
879 let expected_size = 16 * footer.num_columns as usize;
880 if cmo_table.len() != expected_size {
881 return Err(Error::invalid_input(format!(
882 "column metadata offset table has {} bytes but expected {} bytes for {} columns",
883 cmo_table.len(),
884 expected_size,
885 footer.num_columns
886 )));
887 }
888
889 let mut offsets = Vec::with_capacity(footer.num_columns as usize);
890 for col_idx in 0..footer.num_columns {
891 let offset = (col_idx * 16) as usize;
892 let position = LittleEndian::read_u64(&cmo_table[offset..offset + 8]);
893 let length = LittleEndian::read_u64(&cmo_table[offset + 8..offset + 16]);
894 let end = position.checked_add(length).ok_or_else(|| {
895 Error::invalid_input(format!(
896 "column metadata range overflows for column index {}, position={}, length={}",
897 col_idx, position, length
898 ))
899 })?;
900 if position < footer.column_meta_start || end > footer.column_meta_offsets_start {
901 return Err(Error::invalid_input(format!(
902 "column metadata range for column index {} is outside metadata region: position={}, length={}, metadata_start={}, cmo_start={}",
903 col_idx,
904 position,
905 length,
906 footer.column_meta_start,
907 footer.column_meta_offsets_start
908 )));
909 }
910 offsets.push((position, length));
911 }
912
913 Ok(Arc::from(offsets))
914 }
915
916 async fn optimistic_tail_read(
917 data: &Bytes,
918 start_pos: u64,
919 scheduler: &FileScheduler,
920 file_len: u64,
921 ) -> Result<Bytes> {
922 let num_bytes_needed = (file_len - start_pos) as usize;
923 if data.len() >= num_bytes_needed {
924 Ok(data.slice((data.len() - num_bytes_needed)..))
925 } else {
926 let num_bytes_missing = (num_bytes_needed - data.len()) as u64;
927 let start = file_len - num_bytes_needed as u64;
928 let missing_bytes = scheduler
929 .submit_single(start..start + num_bytes_missing, 0)
930 .await?;
931 let mut combined = BytesMut::with_capacity(data.len() + num_bytes_missing as usize);
932 combined.extend(missing_bytes);
933 combined.extend(data);
934 Ok(combined.freeze())
935 }
936 }
937
938 fn do_decode_gbo_table(
939 gbo_bytes: &Bytes,
940 footer: &Footer,
941 version: LanceFileVersion,
942 ) -> Result<Vec<BufferDescriptor>> {
943 let mut global_bufs_cursor = Cursor::new(gbo_bytes);
944
945 let mut global_buffers = Vec::with_capacity(footer.num_global_buffers as usize);
946 for _ in 0..footer.num_global_buffers {
947 let buf_pos = global_bufs_cursor.read_u64::<LittleEndian>()?;
948 assert!(
949 version < LanceFileVersion::V2_1 || buf_pos % PAGE_BUFFER_ALIGNMENT as u64 == 0
950 );
951 let buf_size = global_bufs_cursor.read_u64::<LittleEndian>()?;
952 global_buffers.push(BufferDescriptor {
953 position: buf_pos,
954 size: buf_size,
955 });
956 }
957
958 Ok(global_buffers)
959 }
960
961 async fn decode_gbo_table(
962 tail_bytes: &Bytes,
963 file_len: u64,
964 scheduler: &FileScheduler,
965 footer: &Footer,
966 version: LanceFileVersion,
967 ) -> Result<Vec<BufferDescriptor>> {
968 let gbo_bytes = Self::optimistic_tail_read(
971 tail_bytes,
972 footer.global_buff_offsets_start,
973 scheduler,
974 file_len,
975 )
976 .await?;
977 Self::do_decode_gbo_table(&gbo_bytes, footer, version)
978 }
979
980 fn decode_schema(schema_bytes: Bytes) -> Result<(u64, lance_core::datatypes::Schema)> {
981 let file_descriptor = pb::FileDescriptor::decode(schema_bytes)?;
982 let pb_schema = file_descriptor.schema.unwrap();
983 let num_rows = file_descriptor.length;
984 let fields_with_meta = FieldsWithMeta {
985 fields: Fields(pb_schema.fields),
986 metadata: pb_schema.metadata,
987 };
988 let schema = Schema::try_from(fields_with_meta)?;
989 Ok((num_rows, schema))
990 }
991
992 pub async fn read_all_metadata(scheduler: &FileScheduler) -> Result<CachedFileMetadata> {
1006 let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1008 let tail_offset = file_len - tail_bytes.len() as u64;
1009 let footer = Self::decode_footer(&tail_bytes)?;
1010
1011 let file_version = LanceFileVersion::try_from_major_minor(
1012 footer.major_version as u32,
1013 footer.minor_version as u32,
1014 )?;
1015
1016 let gbo_table =
1017 Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1018 if gbo_table.is_empty() {
1019 return Err(Error::internal(
1020 "File did not contain any global buffers, schema expected".to_string(),
1021 ));
1022 }
1023 let schema_start = gbo_table[0].position;
1024 let schema_size = gbo_table[0].size;
1025
1026 let num_footer_bytes = file_len - schema_start;
1027
1028 let all_metadata_bytes =
1031 Self::optimistic_tail_read(&tail_bytes, schema_start, scheduler, file_len).await?;
1032
1033 let schema_bytes = all_metadata_bytes.slice(0..schema_size as usize);
1034 let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1035
1036 let column_metadata_start = (footer.column_meta_start - schema_start) as usize;
1039 let column_metadata_end = (footer.global_buff_offsets_start - schema_start) as usize;
1040 let column_metadata_bytes =
1041 all_metadata_bytes.slice(column_metadata_start..column_metadata_end);
1042 let column_metadatas = Self::read_all_column_metadata(column_metadata_bytes, &footer)?;
1043
1044 let num_global_buffer_bytes = gbo_table.iter().map(|buf| buf.size).sum::<u64>();
1045 let num_data_bytes = footer.column_meta_start - num_global_buffer_bytes;
1046 let num_column_metadata_bytes = footer.global_buff_offsets_start - footer.column_meta_start;
1047
1048 let column_infos = Self::meta_to_col_infos(column_metadatas.as_slice(), file_version);
1049
1050 let retained_global_buffers =
1056 Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1057
1058 Ok(CachedFileMetadata {
1059 file_schema: Arc::new(schema),
1060 column_metadatas,
1061 column_infos,
1062 num_rows,
1063 num_data_bytes,
1064 num_column_metadata_bytes,
1065 num_global_buffer_bytes,
1066 num_footer_bytes,
1067 file_buffers: gbo_table,
1068 major_version: footer.major_version,
1069 minor_version: footer.minor_version,
1070 file_size_bytes: file_len,
1071 retained_global_buffers,
1072 })
1073 }
1074
1075 async fn read_metadata_index_with_known_schema(
1076 scheduler: &FileScheduler,
1077 known_schema: Option<(Arc<Schema>, u64)>,
1078 ) -> Result<FileMetadataIndex> {
1079 let (tail_bytes, file_len) = Self::read_tail(scheduler).await?;
1080 let tail_offset = file_len - tail_bytes.len() as u64;
1081 let footer = Self::decode_footer(&tail_bytes)?;
1082
1083 let file_version = LanceFileVersion::try_from_major_minor(
1084 footer.major_version as u32,
1085 footer.minor_version as u32,
1086 )?;
1087
1088 let gbo_table =
1089 Self::decode_gbo_table(&tail_bytes, file_len, scheduler, &footer, file_version).await?;
1090 if gbo_table.is_empty() {
1091 return Err(Error::internal(
1092 "File did not contain any global buffers, schema expected".to_string(),
1093 ));
1094 }
1095 let (file_schema, num_rows) = match known_schema {
1096 Some((file_schema, num_rows)) => (file_schema, num_rows),
1097 None => {
1098 let schema_buffer = &gbo_table[0];
1099 let schema_bytes = Self::read_range_from_tail_or_scheduler(
1100 &tail_bytes,
1101 tail_offset,
1102 scheduler,
1103 schema_buffer.position..schema_buffer.position + schema_buffer.size,
1104 )
1105 .await?;
1106 let (num_rows, schema) = Self::decode_schema(schema_bytes)?;
1107 (Arc::new(schema), num_rows)
1108 }
1109 };
1110
1111 let cmo_table = Self::read_range_from_tail_or_scheduler(
1112 &tail_bytes,
1113 tail_offset,
1114 scheduler,
1115 footer.column_meta_offsets_start..footer.global_buff_offsets_start,
1116 )
1117 .await?;
1118 let column_metadata_offsets = Self::decode_cmo_table(cmo_table, &footer)?;
1119
1120 let retained_global_buffers =
1121 Self::retained_global_buffers_from_tail(&gbo_table, &tail_bytes, tail_offset);
1122
1123 Ok(FileMetadataIndex {
1124 file_schema,
1125 num_rows,
1126 file_buffers: gbo_table,
1127 column_metadata_offsets,
1128 num_columns: footer.num_columns,
1129 version: file_version,
1130 file_size_bytes: file_len,
1131 retained_global_buffers,
1132 })
1133 }
1134
1135 pub async fn read_metadata_index(scheduler: &FileScheduler) -> Result<FileMetadataIndex> {
1141 Self::read_metadata_index_with_known_schema(scheduler, None).await
1142 }
1143
1144 pub async fn read_metadata_index_with_schema(
1149 scheduler: &FileScheduler,
1150 file_schema: Arc<Schema>,
1151 num_rows: u64,
1152 ) -> Result<FileMetadataIndex> {
1153 Self::read_metadata_index_with_known_schema(scheduler, Some((file_schema, num_rows))).await
1154 }
1155
1156 fn fetch_encoding<M: Default + Name + Sized>(encoding: &pbfile::Encoding) -> M {
1157 match &encoding.location {
1158 Some(pbfile::encoding::Location::Indirect(_)) => todo!(),
1159 Some(pbfile::encoding::Location::Direct(encoding)) => {
1160 let encoding_buf = Bytes::from(encoding.encoding.clone());
1161 let encoding_any = prost_types::Any::decode(encoding_buf).unwrap();
1162 encoding_any.to_msg::<M>().unwrap()
1163 }
1164 Some(pbfile::encoding::Location::None(_)) => panic!(),
1165 None => panic!(),
1166 }
1167 }
1168
1169 fn meta_to_col_infos(
1170 column_metadatas: &[pbfile::ColumnMetadata],
1171 file_version: LanceFileVersion,
1172 ) -> Vec<Arc<ColumnInfo>> {
1173 column_metadatas
1174 .iter()
1175 .enumerate()
1176 .map(|(col_idx, col_meta)| {
1177 Self::meta_to_col_info(col_idx as u32, col_meta, file_version)
1178 })
1179 .collect::<Vec<_>>()
1180 }
1181
1182 fn meta_to_col_info(
1183 col_idx: u32,
1184 col_meta: &pbfile::ColumnMetadata,
1185 file_version: LanceFileVersion,
1186 ) -> Arc<ColumnInfo> {
1187 let page_infos = col_meta
1188 .pages
1189 .iter()
1190 .map(|page| {
1191 let num_rows = page.length;
1192 let encoding = match file_version {
1193 LanceFileVersion::V2_0 => {
1194 PageEncoding::Legacy(Self::fetch_encoding::<pbenc::ArrayEncoding>(
1195 page.encoding.as_ref().unwrap(),
1196 ))
1197 }
1198 _ => PageEncoding::Structural(Self::fetch_encoding::<pbenc21::PageLayout>(
1199 page.encoding.as_ref().unwrap(),
1200 )),
1201 };
1202 let buffer_offsets_and_sizes = Arc::from(
1203 page.buffer_offsets
1204 .iter()
1205 .zip(page.buffer_sizes.iter())
1206 .map(|(offset, size)| {
1207 assert!(
1209 file_version < LanceFileVersion::V2_1
1210 || offset % PAGE_BUFFER_ALIGNMENT as u64 == 0
1211 );
1212 (*offset, *size)
1213 })
1214 .collect::<Vec<_>>(),
1215 );
1216 PageInfo {
1217 buffer_offsets_and_sizes,
1218 encoding,
1219 num_rows,
1220 priority: page.priority,
1221 }
1222 })
1223 .collect::<Vec<_>>();
1224 let buffer_offsets_and_sizes = Arc::from(
1225 col_meta
1226 .buffer_offsets
1227 .iter()
1228 .zip(col_meta.buffer_sizes.iter())
1229 .map(|(offset, size)| (*offset, *size))
1230 .collect::<Vec<_>>(),
1231 );
1232 Arc::new(ColumnInfo {
1233 index: col_idx,
1234 page_infos: Arc::from(page_infos),
1235 buffer_offsets_and_sizes,
1236 encoding: Self::fetch_encoding(col_meta.encoding.as_ref().unwrap()),
1237 })
1238 }
1239
1240 fn validate_projection(
1241 projection: &ReaderProjection,
1242 metadata: &CachedFileMetadata,
1243 ) -> Result<()> {
1244 if projection.schema.fields.is_empty() {
1245 return Err(Error::invalid_input(
1246 "Attempt to read zero columns from the file, at least one column must be specified"
1247 .to_string(),
1248 ));
1249 }
1250 let mut column_indices_seen = BTreeSet::new();
1251 for column_index in &projection.column_indices {
1252 if !column_indices_seen.insert(*column_index) {
1253 return Err(Error::invalid_input(format!(
1254 "The projection specified the column index {} more than once",
1255 column_index
1256 )));
1257 }
1258 if *column_index >= metadata.column_infos.len() as u32 {
1259 return Err(Error::invalid_input(format!(
1260 "The projection specified the column index {} but there are only {} columns in the file",
1261 column_index,
1262 metadata.column_infos.len()
1263 )));
1264 }
1265 }
1266 Ok(())
1267 }
1268
1269 pub async fn try_open(
1276 scheduler: FileScheduler,
1277 base_projection: Option<ReaderProjection>,
1278 decoder_plugins: Arc<DecoderPlugins>,
1279 cache: &LanceCache,
1280 options: FileReaderOptions,
1281 ) -> Result<Self> {
1282 let file_metadata = Arc::new(Self::read_all_metadata(&scheduler).await?);
1283 let path = scheduler.reader().path().clone();
1284
1285 let encodings_io =
1287 LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
1288
1289 Self::try_open_with_file_metadata(
1290 Arc::new(encodings_io),
1291 path,
1292 base_projection,
1293 decoder_plugins,
1294 file_metadata,
1295 cache,
1296 options,
1297 )
1298 .await
1299 }
1300
1301 pub async fn try_open_with_file_metadata(
1308 scheduler: Arc<dyn EncodingsIo>,
1309 path: Path,
1310 base_projection: Option<ReaderProjection>,
1311 decoder_plugins: Arc<DecoderPlugins>,
1312 file_metadata: Arc<CachedFileMetadata>,
1313 cache: &LanceCache,
1314 options: FileReaderOptions,
1315 ) -> Result<Self> {
1316 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
1317 let core = FileReadCore::try_new(
1318 scheduler,
1319 base_projection,
1320 decoder_plugins,
1321 FileMetadataProvider::Full(file_metadata.clone()),
1322 cache,
1323 options,
1324 )?;
1325 Ok(Self {
1326 core,
1327 metadata: file_metadata,
1328 })
1329 }
1330
1331 fn collect_columns_from_projection(
1345 &self,
1346 _projection: &ReaderProjection,
1347 ) -> Result<Vec<Arc<ColumnInfo>>> {
1348 Ok(self.metadata.column_infos.clone())
1349 }
1350
1351 #[allow(clippy::too_many_arguments)]
1352 async fn do_read_range(
1353 column_infos: Vec<Arc<ColumnInfo>>,
1354 io: Arc<dyn EncodingsIo>,
1355 cache: Arc<LanceCache>,
1356 num_rows: u64,
1357 decoder_plugins: Arc<DecoderPlugins>,
1358 range: Range<u64>,
1359 batch_size: u32,
1360 projection: ReaderProjection,
1361 filter: FilterExpression,
1362 decoder_config: DecoderConfig,
1363 batch_size_bytes: Option<u64>,
1364 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1365 debug!(
1366 "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1367 range,
1368 batch_size,
1369 num_rows,
1370 column_infos.len(),
1371 projection.schema.fields.len(),
1372 );
1373
1374 let config = SchedulerDecoderConfig {
1375 batch_size,
1376 cache,
1377 decoder_plugins,
1378 io,
1379 decoder_config,
1380 batch_size_bytes,
1381 };
1382
1383 let requested_rows = RequestedRows::Ranges(vec![range]);
1384
1385 schedule_and_decode(
1386 column_infos,
1387 requested_rows,
1388 filter,
1389 projection.column_indices,
1390 projection.schema,
1391 config,
1392 )
1393 .await
1394 }
1395
1396 #[allow(clippy::too_many_arguments)]
1397 async fn do_take_rows(
1398 column_infos: Vec<Arc<ColumnInfo>>,
1399 io: Arc<dyn EncodingsIo>,
1400 cache: Arc<LanceCache>,
1401 decoder_plugins: Arc<DecoderPlugins>,
1402 indices: Vec<u64>,
1403 batch_size: u32,
1404 projection: ReaderProjection,
1405 filter: FilterExpression,
1406 decoder_config: DecoderConfig,
1407 batch_size_bytes: Option<u64>,
1408 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1409 debug!(
1410 "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1411 indices.len(),
1412 indices[0],
1413 indices[indices.len() - 1],
1414 batch_size,
1415 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1416 );
1417
1418 let config = SchedulerDecoderConfig {
1419 batch_size,
1420 cache,
1421 decoder_plugins,
1422 io,
1423 decoder_config,
1424 batch_size_bytes,
1425 };
1426
1427 let requested_rows = RequestedRows::Indices(indices);
1428
1429 schedule_and_decode(
1430 column_infos,
1431 requested_rows,
1432 filter,
1433 projection.column_indices,
1434 projection.schema,
1435 config,
1436 )
1437 .await
1438 }
1439
1440 #[allow(clippy::too_many_arguments)]
1441 async fn do_read_ranges(
1442 column_infos: Vec<Arc<ColumnInfo>>,
1443 io: Arc<dyn EncodingsIo>,
1444 cache: Arc<LanceCache>,
1445 decoder_plugins: Arc<DecoderPlugins>,
1446 ranges: Vec<Range<u64>>,
1447 batch_size: u32,
1448 projection: ReaderProjection,
1449 filter: FilterExpression,
1450 decoder_config: DecoderConfig,
1451 batch_size_bytes: Option<u64>,
1452 ) -> Result<BoxStream<'static, ReadBatchTask>> {
1453 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1454 debug!(
1455 "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1456 ranges.len(),
1457 num_rows,
1458 ranges[0].start,
1459 ranges[ranges.len() - 1].end,
1460 batch_size,
1461 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1462 );
1463
1464 let config = SchedulerDecoderConfig {
1465 batch_size,
1466 cache,
1467 decoder_plugins,
1468 io,
1469 decoder_config,
1470 batch_size_bytes,
1471 };
1472
1473 let requested_rows = RequestedRows::Ranges(ranges);
1474
1475 schedule_and_decode(
1476 column_infos,
1477 requested_rows,
1478 filter,
1479 projection.column_indices,
1480 projection.schema,
1481 config,
1482 )
1483 .await
1484 }
1485
1486 pub async fn read_tasks(
1510 &self,
1511 params: ReadBatchParams,
1512 batch_size: u32,
1513 projection: Option<ReaderProjection>,
1514 filter: FilterExpression,
1515 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
1516 self.core
1517 .read_tasks(params, batch_size, projection, filter)
1518 .await
1519 }
1520
1521 pub async fn read_stream_projected(
1551 &self,
1552 params: ReadBatchParams,
1553 batch_size: u32,
1554 batch_readahead: u32,
1555 projection: ReaderProjection,
1556 filter: FilterExpression,
1557 ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1558 let arrow_schema = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
1559 let tasks_stream = self
1560 .read_tasks(params, batch_size, Some(projection), filter)
1561 .await?;
1562 let batch_stream = tasks_stream
1563 .map(|task| task.task)
1564 .buffered(batch_readahead as usize)
1565 .boxed();
1566 Ok(Box::pin(RecordBatchStreamAdapter::new(
1567 arrow_schema,
1568 batch_stream,
1569 )))
1570 }
1571
1572 fn take_rows_blocking(
1573 &self,
1574 indices: Vec<u64>,
1575 batch_size: u32,
1576 projection: ReaderProjection,
1577 filter: FilterExpression,
1578 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1579 let column_infos = self.collect_columns_from_projection(&projection)?;
1580 debug!(
1581 "Taking {} rows spread across range {}..{} with batch_size {} from columns {:?}",
1582 indices.len(),
1583 indices[0],
1584 indices[indices.len() - 1],
1585 batch_size,
1586 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1587 );
1588
1589 let config = SchedulerDecoderConfig {
1590 batch_size,
1591 cache: self.core.cache.clone(),
1592 decoder_plugins: self.core.decoder_plugins.clone(),
1593 io: self.core.scheduler.clone(),
1594 decoder_config: self.core.options.decoder_config.clone(),
1595 batch_size_bytes: self.core.options.batch_size_bytes,
1596 };
1597
1598 let requested_rows = RequestedRows::Indices(indices);
1599
1600 schedule_and_decode_blocking(
1601 column_infos,
1602 requested_rows,
1603 filter,
1604 projection.column_indices,
1605 projection.schema,
1606 config,
1607 )
1608 }
1609
1610 fn read_ranges_blocking(
1611 &self,
1612 ranges: Vec<Range<u64>>,
1613 batch_size: u32,
1614 projection: ReaderProjection,
1615 filter: FilterExpression,
1616 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1617 let column_infos = self.collect_columns_from_projection(&projection)?;
1618 let num_rows = ranges.iter().map(|r| r.end - r.start).sum::<u64>();
1619 debug!(
1620 "Taking {} ranges ({} rows) spread across range {}..{} with batch_size {} from columns {:?}",
1621 ranges.len(),
1622 num_rows,
1623 ranges[0].start,
1624 ranges[ranges.len() - 1].end,
1625 batch_size,
1626 column_infos.iter().map(|ci| ci.index).collect::<Vec<_>>()
1627 );
1628
1629 let config = SchedulerDecoderConfig {
1630 batch_size,
1631 cache: self.core.cache.clone(),
1632 decoder_plugins: self.core.decoder_plugins.clone(),
1633 io: self.core.scheduler.clone(),
1634 decoder_config: self.core.options.decoder_config.clone(),
1635 batch_size_bytes: self.core.options.batch_size_bytes,
1636 };
1637
1638 let requested_rows = RequestedRows::Ranges(ranges);
1639
1640 schedule_and_decode_blocking(
1641 column_infos,
1642 requested_rows,
1643 filter,
1644 projection.column_indices,
1645 projection.schema,
1646 config,
1647 )
1648 }
1649
1650 fn read_range_blocking(
1651 &self,
1652 range: Range<u64>,
1653 batch_size: u32,
1654 projection: ReaderProjection,
1655 filter: FilterExpression,
1656 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1657 let column_infos = self.collect_columns_from_projection(&projection)?;
1658 let num_rows = self.core.num_rows();
1659
1660 debug!(
1661 "Reading range {:?} with batch_size {} from file with {} rows and {} columns into schema with {} columns",
1662 range,
1663 batch_size,
1664 num_rows,
1665 column_infos.len(),
1666 projection.schema.fields.len(),
1667 );
1668
1669 let config = SchedulerDecoderConfig {
1670 batch_size,
1671 cache: self.core.cache.clone(),
1672 decoder_plugins: self.core.decoder_plugins.clone(),
1673 io: self.core.scheduler.clone(),
1674 decoder_config: self.core.options.decoder_config.clone(),
1675 batch_size_bytes: self.core.options.batch_size_bytes,
1676 };
1677
1678 let requested_rows = RequestedRows::Ranges(vec![range]);
1679
1680 schedule_and_decode_blocking(
1681 column_infos,
1682 requested_rows,
1683 filter,
1684 projection.column_indices,
1685 projection.schema,
1686 config,
1687 )
1688 }
1689
1690 pub fn read_stream_projected_blocking(
1702 &self,
1703 params: ReadBatchParams,
1704 batch_size: u32,
1705 projection: Option<ReaderProjection>,
1706 filter: FilterExpression,
1707 ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
1708 let projection = projection.unwrap_or_else(|| self.core.base_projection.clone());
1709 Self::validate_projection(&projection, &self.metadata)?;
1710 let prepared = PreparedProjection {
1716 column_infos: self.metadata.column_infos.clone(),
1717 decoder_projection: projection.clone(),
1718 };
1719 let read_len = self.core.prepared_read_length(&prepared)?;
1720 let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
1721 if bound > read_len || (bound == read_len && inclusive) {
1722 Err(Error::invalid_input(format!(
1723 "cannot read {params:?} from columns with {read_len} rows"
1724 )))
1725 } else {
1726 Ok(())
1727 }
1728 };
1729 match ¶ms {
1730 ReadBatchParams::Indices(indices) => {
1731 for idx in indices {
1732 match idx {
1733 None => {
1734 return Err(Error::invalid_input("Null value in indices array"));
1735 }
1736 Some(idx) => {
1737 verify_bound(¶ms, idx as u64, true)?;
1738 }
1739 }
1740 }
1741 let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
1742 self.take_rows_blocking(indices, batch_size, projection, filter)
1743 }
1744 ReadBatchParams::Range(range) => {
1745 verify_bound(¶ms, range.end as u64, false)?;
1746 self.read_range_blocking(
1747 range.start as u64..range.end as u64,
1748 batch_size,
1749 projection,
1750 filter,
1751 )
1752 }
1753 ReadBatchParams::Ranges(ranges) => {
1754 let mut ranges_u64 = Vec::with_capacity(ranges.len());
1755 for range in ranges.as_ref() {
1756 verify_bound(¶ms, range.end, false)?;
1757 ranges_u64.push(range.start..range.end);
1758 }
1759 self.read_ranges_blocking(ranges_u64, batch_size, projection, filter)
1760 }
1761 ReadBatchParams::RangeFrom(range) => {
1762 verify_bound(¶ms, range.start as u64, true)?;
1763 self.read_range_blocking(
1764 range.start as u64..read_len,
1765 batch_size,
1766 projection,
1767 filter,
1768 )
1769 }
1770 ReadBatchParams::RangeTo(range) => {
1771 verify_bound(¶ms, range.end as u64, false)?;
1772 self.read_range_blocking(0..range.end as u64, batch_size, projection, filter)
1773 }
1774 ReadBatchParams::RangeFull => {
1775 self.read_range_blocking(0..read_len, batch_size, projection, filter)
1776 }
1777 }
1778 }
1779
1780 pub async fn read_stream(
1792 &self,
1793 params: ReadBatchParams,
1794 batch_size: u32,
1795 batch_readahead: u32,
1796 filter: FilterExpression,
1797 ) -> Result<Pin<Box<dyn RecordBatchStream>>> {
1798 self.read_stream_projected(
1799 params,
1800 batch_size,
1801 batch_readahead,
1802 self.core.base_projection.clone(),
1803 filter,
1804 )
1805 .await
1806 }
1807
1808 pub fn schema(&self) -> &Arc<Schema> {
1809 self.core.schema()
1810 }
1811}
1812
1813impl FileMetadataProvider {
1814 fn version(&self) -> LanceFileVersion {
1815 match self {
1816 Self::Full(metadata) => metadata.version(),
1817 Self::Indexed(metadata_index) => metadata_index.version,
1818 }
1819 }
1820
1821 fn num_rows(&self) -> u64 {
1822 match self {
1823 Self::Full(metadata) => metadata.num_rows,
1824 Self::Indexed(metadata_index) => metadata_index.num_rows,
1825 }
1826 }
1827
1828 fn schema(&self) -> &Arc<Schema> {
1829 match self {
1830 Self::Full(metadata) => &metadata.file_schema,
1831 Self::Indexed(metadata_index) => &metadata_index.file_schema,
1832 }
1833 }
1834
1835 fn file_buffers(&self) -> &Vec<BufferDescriptor> {
1836 match self {
1837 Self::Full(metadata) => &metadata.file_buffers,
1838 Self::Indexed(metadata_index) => &metadata_index.file_buffers,
1839 }
1840 }
1841
1842 fn retained_global_buffers(&self) -> &BTreeMap<u32, Bytes> {
1843 match self {
1844 Self::Full(metadata) => &metadata.retained_global_buffers,
1845 Self::Indexed(metadata_index) => &metadata_index.retained_global_buffers,
1846 }
1847 }
1848
1849 fn file_statistics(&self) -> Option<FileStatistics> {
1850 let metadata = match self {
1851 Self::Full(metadata) => metadata,
1852 Self::Indexed(_) => return None,
1853 };
1854 Some(FileReader::statistics_from_column_metadata(
1855 &metadata.column_metadatas,
1856 ))
1857 }
1858
1859 fn supports_indexed_projection(
1860 projection: &ReaderProjection,
1861 version: LanceFileVersion,
1862 ) -> bool {
1863 if version < LanceFileVersion::V2_1 || projection.schema.fields.is_empty() {
1864 return false;
1865 }
1866
1867 projection
1868 .schema
1869 .fields
1870 .iter()
1871 .try_fold(0usize, |count, field| {
1872 count.checked_add(indexed_projection_column_count(field)?)
1873 })
1874 == Some(projection.column_indices.len())
1875 }
1876
1877 fn validate_indexed_projection(
1878 projection: &ReaderProjection,
1879 metadata_index: &FileMetadataIndex,
1880 ) -> Result<()> {
1881 if projection.schema.fields.is_empty() {
1882 return Err(Error::invalid_input(
1883 "Attempt to read zero columns from the file, at least one column must be specified"
1884 .to_string(),
1885 ));
1886 }
1887 let mut column_indices_seen = BTreeSet::new();
1888 for column_index in &projection.column_indices {
1889 if !column_indices_seen.insert(*column_index) {
1890 return Err(Error::invalid_input(format!(
1891 "The projection specified the column index {} more than once",
1892 column_index
1893 )));
1894 }
1895 if *column_index >= metadata_index.num_columns {
1896 return Err(Error::invalid_input(format!(
1897 "The projection specified the column index {} but there are only {} columns in the file",
1898 column_index, metadata_index.num_columns
1899 )));
1900 }
1901 }
1902 if !Self::supports_indexed_projection(projection, metadata_index.version) {
1903 return Err(Error::not_supported(format!(
1904 "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",
1905 metadata_index.version,
1906 projection.schema.fields.len(),
1907 projection.column_indices.len()
1908 )));
1909 }
1910 Ok(())
1911 }
1912
1913 fn validate_projection(&self, projection: &ReaderProjection) -> Result<()> {
1914 match self {
1915 Self::Full(metadata) => FileReader::validate_projection(projection, metadata),
1916 Self::Indexed(metadata_index) => {
1917 Self::validate_indexed_projection(projection, metadata_index)
1918 }
1919 }
1920 }
1921
1922 fn column_metadata_range(
1923 metadata_index: &FileMetadataIndex,
1924 column_index: u32,
1925 ) -> Result<Range<u64>> {
1926 let (position, length) = metadata_index
1927 .column_metadata_offsets
1928 .get(column_index as usize)
1929 .copied()
1930 .ok_or_else(|| {
1931 Error::invalid_input(format!(
1932 "The projection specified the column index {} but there are only {} columns in the file",
1933 column_index, metadata_index.num_columns
1934 ))
1935 })?;
1936 let end = position.checked_add(length).ok_or_else(|| {
1937 Error::invalid_input(format!(
1938 "column metadata range overflows for column index {}, position={}, length={}",
1939 column_index, position, length
1940 ))
1941 })?;
1942 Ok(position..end)
1943 }
1944
1945 async fn load_indexed_column_infos(
1946 metadata_index: &FileMetadataIndex,
1947 io: &Arc<dyn EncodingsIo>,
1948 cache: &Arc<LanceCache>,
1949 column_indices: &[u32],
1950 ) -> Result<Vec<Arc<ColumnInfo>>> {
1951 let mut column_infos = vec![None; column_indices.len()];
1952 let mut missing_columns = Vec::new();
1953
1954 for (result_index, column_index) in column_indices.iter().copied().enumerate() {
1955 let cache_key = ColumnMetadataCacheKey { column_index };
1956 if let Some(cached) = cache.get_with_key(&cache_key).await {
1957 column_infos[result_index] = Some(cached.column_info.clone());
1958 } else {
1959 let range = Self::column_metadata_range(metadata_index, column_index)?;
1960 missing_columns.push((result_index, column_index, range));
1961 }
1962 }
1963
1964 missing_columns.sort_by_key(|(_, _, range)| range.start);
1965 if !missing_columns.is_empty() {
1966 let ranges = missing_columns
1967 .iter()
1968 .map(|(_, _, range)| range.clone())
1969 .collect::<Vec<_>>();
1970 let metadata_bytes = io.submit_request(ranges, 0).await?;
1971 for ((result_index, column_index, _), bytes) in
1972 missing_columns.into_iter().zip(metadata_bytes)
1973 {
1974 let column_metadata = pbfile::ColumnMetadata::decode(bytes)?;
1975 let column_info = FileReader::meta_to_col_info(
1976 column_index,
1977 &column_metadata,
1978 metadata_index.version,
1979 );
1980 let cached = Arc::new(CachedColumnMetadata {
1981 column_metadata,
1982 column_info: column_info.clone(),
1983 });
1984 let cache_key = ColumnMetadataCacheKey { column_index };
1985 cache.insert_with_key(&cache_key, cached).await;
1986 column_infos[result_index] = Some(column_info);
1987 }
1988 }
1989
1990 column_infos
1991 .into_iter()
1992 .enumerate()
1993 .map(|(idx, column_info)| {
1994 column_info.ok_or_else(|| {
1995 Error::internal(format!(
1996 "lazy metadata loader did not load requested projection column at position {}",
1997 idx
1998 ))
1999 })
2000 })
2001 .collect()
2002 }
2003
2004 async fn prepare_projection(
2005 &self,
2006 projection: &ReaderProjection,
2007 io: &Arc<dyn EncodingsIo>,
2008 cache: &Arc<LanceCache>,
2009 ) -> Result<PreparedProjection> {
2010 self.validate_projection(projection)?;
2011 match self {
2012 Self::Full(metadata) => Ok(PreparedProjection {
2013 column_infos: metadata.column_infos.clone(),
2014 decoder_projection: projection.clone(),
2015 }),
2016 Self::Indexed(metadata_index) => {
2017 let column_infos = Self::load_indexed_column_infos(
2018 metadata_index,
2019 io,
2020 cache,
2021 &projection.column_indices,
2022 )
2023 .await?;
2024 let decoder_projection = ReaderProjection {
2025 schema: projection.schema.clone(),
2026 column_indices: (0..projection.column_indices.len())
2027 .map(|idx| idx as u32)
2028 .collect(),
2029 };
2030 Ok(PreparedProjection {
2031 column_infos,
2032 decoder_projection,
2033 })
2034 }
2035 }
2036 }
2037}
2038
2039impl FileReadCore {
2040 fn try_new(
2041 scheduler: Arc<dyn EncodingsIo>,
2042 base_projection: Option<ReaderProjection>,
2043 decoder_plugins: Arc<DecoderPlugins>,
2044 metadata_provider: FileMetadataProvider,
2045 cache: Arc<LanceCache>,
2046 options: FileReaderOptions,
2047 ) -> Result<Self> {
2048 if let Some(base_projection) = base_projection.as_ref() {
2049 metadata_provider.validate_projection(base_projection)?;
2050 }
2051 let base_projection = base_projection.unwrap_or(ReaderProjection::from_whole_schema(
2052 metadata_provider.schema().as_ref(),
2053 metadata_provider.version(),
2054 ));
2055 Ok(Self {
2056 scheduler,
2057 base_projection,
2058 metadata_provider,
2059 decoder_plugins,
2060 cache,
2061 options,
2062 })
2063 }
2064
2065 fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2066 Self {
2067 scheduler,
2068 base_projection: self.base_projection.clone(),
2069 metadata_provider: self.metadata_provider.clone(),
2070 decoder_plugins: self.decoder_plugins.clone(),
2071 cache: self.cache.clone(),
2072 options: self.options.clone(),
2073 }
2074 }
2075
2076 fn version(&self) -> LanceFileVersion {
2077 self.metadata_provider.version()
2078 }
2079
2080 fn num_rows(&self) -> u64 {
2081 self.metadata_provider.num_rows()
2082 }
2083
2084 fn schema(&self) -> &Arc<Schema> {
2085 self.metadata_provider.schema()
2086 }
2087
2088 async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2089 let file_buffers = self.metadata_provider.file_buffers();
2090 let buffer_desc = file_buffers.get(index as usize).ok_or_else(|| {
2091 Error::invalid_input(format!(
2092 "request for global buffer at index {} but there were only {} global buffers in the file",
2093 index,
2094 file_buffers.len()
2095 ))
2096 })?;
2097
2098 if let Some(bytes) = self.metadata_provider.retained_global_buffers().get(&index) {
2099 return Ok(bytes.clone());
2100 }
2101
2102 let bytes = self
2103 .scheduler
2104 .submit_request(
2105 vec![buffer_desc.position..buffer_desc.position + buffer_desc.size],
2106 0,
2107 )
2108 .await?;
2109 bytes.into_iter().next().ok_or_else(|| {
2110 Error::internal(format!(
2111 "global buffer read for index {} returned no bytes",
2112 index
2113 ))
2114 })
2115 }
2116
2117 fn prepared_read_length(&self, prepared: &PreparedProjection) -> Result<u64> {
2131 let is_structural = self.version() >= LanceFileVersion::V2_1;
2132 let column_infos = &prepared.column_infos;
2133 let column_len = |column: usize| -> Result<u64> {
2134 let info = column_infos.get(column).ok_or_else(|| {
2135 Error::invalid_input(format!(
2136 "projection references column index {} but only {} columns are available",
2137 column,
2138 column_infos.len()
2139 ))
2140 })?;
2141 Ok(info.page_infos.iter().map(|page| page.num_rows).sum())
2142 };
2143 let column_indices = &prepared.decoder_projection.column_indices;
2144 let fields = &prepared.decoder_projection.schema.fields;
2145 let mut cursor = 0usize;
2146 let mut field_lengths = Vec::with_capacity(fields.len());
2147 for field in fields {
2148 let rows = validate_field_length(
2149 field,
2150 is_structural,
2151 true,
2152 column_indices,
2153 &mut cursor,
2154 &column_len,
2155 )?;
2156 field_lengths.push((field.name.as_str(), rows));
2157 }
2158 if cursor != column_indices.len() {
2159 return Err(Error::invalid_input(format!(
2160 "projection supplied {} column indices but its fields require {}",
2161 column_indices.len(),
2162 cursor
2163 )));
2164 }
2165 verify_uniform_lengths(&field_lengths)
2166 }
2167
2168 async fn read_range(
2169 &self,
2170 range: Range<u64>,
2171 batch_size: u32,
2172 prepared: PreparedProjection,
2173 filter: FilterExpression,
2174 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2175 FileReader::do_read_range(
2176 prepared.column_infos,
2177 self.scheduler.clone(),
2178 self.cache.clone(),
2179 self.num_rows(),
2180 self.decoder_plugins.clone(),
2181 range,
2182 batch_size,
2183 prepared.decoder_projection,
2184 filter,
2185 self.options.decoder_config.clone(),
2186 self.options.batch_size_bytes,
2187 )
2188 .await
2189 }
2190
2191 async fn take_rows(
2192 &self,
2193 indices: Vec<u64>,
2194 batch_size: u32,
2195 prepared: PreparedProjection,
2196 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2197 FileReader::do_take_rows(
2198 prepared.column_infos,
2199 self.scheduler.clone(),
2200 self.cache.clone(),
2201 self.decoder_plugins.clone(),
2202 indices,
2203 batch_size,
2204 prepared.decoder_projection,
2205 FilterExpression::no_filter(),
2206 self.options.decoder_config.clone(),
2207 self.options.batch_size_bytes,
2208 )
2209 .await
2210 }
2211
2212 async fn read_ranges(
2213 &self,
2214 ranges: Vec<Range<u64>>,
2215 batch_size: u32,
2216 prepared: PreparedProjection,
2217 filter: FilterExpression,
2218 ) -> Result<BoxStream<'static, ReadBatchTask>> {
2219 FileReader::do_read_ranges(
2220 prepared.column_infos,
2221 self.scheduler.clone(),
2222 self.cache.clone(),
2223 self.decoder_plugins.clone(),
2224 ranges,
2225 batch_size,
2226 prepared.decoder_projection,
2227 filter,
2228 self.options.decoder_config.clone(),
2229 self.options.batch_size_bytes,
2230 )
2231 .await
2232 }
2233
2234 async fn read_tasks(
2235 &self,
2236 params: ReadBatchParams,
2237 batch_size: u32,
2238 projection: Option<ReaderProjection>,
2239 filter: FilterExpression,
2240 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2241 let projection = projection.unwrap_or_else(|| self.base_projection.clone());
2242 let prepared = self
2243 .metadata_provider
2244 .prepare_projection(&projection, &self.scheduler, &self.cache)
2245 .await?;
2246 let read_len = self.prepared_read_length(&prepared)?;
2253 let verify_bound = |params: &ReadBatchParams, bound: u64, inclusive: bool| {
2254 if bound > read_len || (bound == read_len && inclusive) {
2255 Err(Error::invalid_input(format!(
2256 "cannot read {params:?} from columns with {read_len} rows"
2257 )))
2258 } else {
2259 Ok(())
2260 }
2261 };
2262 match ¶ms {
2263 ReadBatchParams::Indices(indices) => {
2264 for idx in indices {
2265 match idx {
2266 None => {
2267 return Err(Error::invalid_input("Null value in indices array"));
2268 }
2269 Some(idx) => {
2270 verify_bound(¶ms, idx as u64, true)?;
2271 }
2272 }
2273 }
2274 let indices = indices.iter().map(|idx| idx.unwrap() as u64).collect();
2275 self.take_rows(indices, batch_size, prepared).await
2276 }
2277 ReadBatchParams::Range(range) => {
2278 verify_bound(¶ms, range.end as u64, false)?;
2279 self.read_range(
2280 range.start as u64..range.end as u64,
2281 batch_size,
2282 prepared,
2283 filter,
2284 )
2285 .await
2286 }
2287 ReadBatchParams::Ranges(ranges) => {
2288 let mut ranges_u64 = Vec::with_capacity(ranges.len());
2289 for range in ranges.as_ref() {
2290 verify_bound(¶ms, range.end, false)?;
2291 ranges_u64.push(range.start..range.end);
2292 }
2293 self.read_ranges(ranges_u64, batch_size, prepared, filter)
2294 .await
2295 }
2296 ReadBatchParams::RangeFrom(range) => {
2297 verify_bound(¶ms, range.start as u64, true)?;
2298 self.read_range(range.start as u64..read_len, batch_size, prepared, filter)
2299 .await
2300 }
2301 ReadBatchParams::RangeTo(range) => {
2302 verify_bound(¶ms, range.end as u64, false)?;
2303 self.read_range(0..range.end as u64, batch_size, prepared, filter)
2304 .await
2305 }
2306 ReadBatchParams::RangeFull => {
2307 self.read_range(0..read_len, batch_size, prepared, filter)
2308 .await
2309 }
2310 }
2311 }
2312}
2313
2314impl ProjectedFileReader {
2315 pub async fn try_open(
2320 scheduler: FileScheduler,
2321 base_projection: Option<ReaderProjection>,
2322 decoder_plugins: Arc<DecoderPlugins>,
2323 cache: &LanceCache,
2324 options: FileReaderOptions,
2325 ) -> Result<Self> {
2326 let base_projection = Self::require_indexed_base_projection(base_projection)?;
2327 let metadata_index = Arc::new(FileReader::read_metadata_index(&scheduler).await?);
2328 let path = scheduler.reader().path().clone();
2329 let encodings_io =
2330 LanceEncodingsIo::new(scheduler).with_read_chunk_size(options.read_chunk_size);
2331 Self::try_open_with_metadata_index(
2332 Arc::new(encodings_io),
2333 path,
2334 Some(base_projection),
2335 decoder_plugins,
2336 metadata_index,
2337 cache,
2338 options,
2339 )
2340 .await
2341 }
2342
2343 pub async fn try_open_with_metadata_index(
2349 scheduler: Arc<dyn EncodingsIo>,
2350 path: Path,
2351 base_projection: Option<ReaderProjection>,
2352 decoder_plugins: Arc<DecoderPlugins>,
2353 metadata_index: Arc<FileMetadataIndex>,
2354 cache: &LanceCache,
2355 options: FileReaderOptions,
2356 ) -> Result<Self> {
2357 let base_projection = Self::require_indexed_base_projection(base_projection)?;
2358 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2359 let core = FileReadCore::try_new(
2360 scheduler,
2361 Some(base_projection),
2362 decoder_plugins,
2363 FileMetadataProvider::Indexed(metadata_index),
2364 cache,
2365 options,
2366 )?;
2367 Ok(Self { core })
2368 }
2369
2370 fn require_indexed_base_projection(
2371 base_projection: Option<ReaderProjection>,
2372 ) -> Result<ReaderProjection> {
2373 base_projection.ok_or_else(|| {
2374 Error::invalid_input("ProjectedFileReader requires an explicit base projection")
2375 })
2376 }
2377
2378 pub async fn try_open_with_file_metadata(
2379 scheduler: Arc<dyn EncodingsIo>,
2380 path: Path,
2381 base_projection: Option<ReaderProjection>,
2382 decoder_plugins: Arc<DecoderPlugins>,
2383 file_metadata: Arc<CachedFileMetadata>,
2384 cache: &LanceCache,
2385 options: FileReaderOptions,
2386 ) -> Result<Self> {
2387 let cache = Arc::new(cache.with_key_prefix(path.as_ref()));
2388 let core = FileReadCore::try_new(
2389 scheduler,
2390 base_projection,
2391 decoder_plugins,
2392 FileMetadataProvider::Full(file_metadata),
2393 cache,
2394 options,
2395 )?;
2396 Ok(Self { core })
2397 }
2398
2399 pub fn supports_projection(projection: &ReaderProjection, version: LanceFileVersion) -> bool {
2401 FileMetadataProvider::supports_indexed_projection(projection, version)
2402 }
2403
2404 pub fn with_scheduler(&self, scheduler: Arc<dyn EncodingsIo>) -> Self {
2406 Self {
2407 core: self.core.with_scheduler(scheduler),
2408 }
2409 }
2410
2411 pub fn version(&self) -> LanceFileVersion {
2413 self.core.version()
2414 }
2415
2416 pub fn num_rows(&self) -> u64 {
2418 self.core.num_rows()
2419 }
2420
2421 pub fn schema(&self) -> &Arc<Schema> {
2423 self.core.schema()
2424 }
2425
2426 pub fn file_statistics(&self) -> Option<FileStatistics> {
2428 self.core.metadata_provider.file_statistics()
2429 }
2430
2431 #[cfg(test)]
2432 fn metadata_index(&self) -> Option<&Arc<FileMetadataIndex>> {
2433 match &self.core.metadata_provider {
2434 FileMetadataProvider::Indexed(metadata_index) => Some(metadata_index),
2435 FileMetadataProvider::Full(_) => None,
2436 }
2437 }
2438
2439 pub async fn read_global_buffer(&self, index: u32) -> Result<Bytes> {
2441 self.core.read_global_buffer(index).await
2442 }
2443
2444 pub async fn read_tasks(
2446 &self,
2447 params: ReadBatchParams,
2448 batch_size: u32,
2449 projection: Option<ReaderProjection>,
2450 filter: FilterExpression,
2451 ) -> Result<Pin<Box<dyn Stream<Item = ReadBatchTask> + Send>>> {
2452 self.core
2453 .read_tasks(params, batch_size, projection, filter)
2454 .await
2455 }
2456}
2457
2458pub fn describe_encoding(page: &pbfile::column_metadata::Page) -> String {
2460 if let Some(encoding) = &page.encoding {
2461 if let Some(style) = &encoding.location {
2462 match style {
2463 pbfile::encoding::Location::Indirect(indirect) => {
2464 format!(
2465 "IndirectEncoding(pos={},size={})",
2466 indirect.buffer_location, indirect.buffer_length
2467 )
2468 }
2469 pbfile::encoding::Location::Direct(direct) => {
2470 let encoding_any =
2471 prost_types::Any::decode(Bytes::from(direct.encoding.clone()))
2472 .expect("failed to deserialize encoding as protobuf");
2473 if encoding_any.type_url == "/lance.encodings.ArrayEncoding" {
2474 let encoding = encoding_any.to_msg::<pbenc::ArrayEncoding>();
2475 match encoding {
2476 Ok(encoding) => {
2477 format!("{:#?}", encoding)
2478 }
2479 Err(err) => {
2480 format!("Unsupported(decode_err={})", err)
2481 }
2482 }
2483 } else if encoding_any.type_url == "/lance.encodings21.PageLayout" {
2484 let encoding = encoding_any.to_msg::<pbenc21::PageLayout>();
2485 match encoding {
2486 Ok(encoding) => {
2487 format!("{:#?}", encoding)
2488 }
2489 Err(err) => {
2490 format!("Unsupported(decode_err={})", err)
2491 }
2492 }
2493 } else {
2494 format!("Unrecognized(type_url={})", encoding_any.type_url)
2495 }
2496 }
2497 pbfile::encoding::Location::None(_) => "NoEncodingDescription".to_string(),
2498 }
2499 } else {
2500 "MISSING STYLE".to_string()
2501 }
2502 } else {
2503 "MISSING".to_string()
2504 }
2505}
2506
2507pub trait EncodedBatchReaderExt {
2508 fn try_from_mini_lance(
2509 bytes: Bytes,
2510 schema: &Schema,
2511 version: LanceFileVersion,
2512 ) -> Result<Self>
2513 where
2514 Self: Sized;
2515 fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2516 where
2517 Self: Sized;
2518}
2519
2520impl EncodedBatchReaderExt for EncodedBatch {
2521 fn try_from_mini_lance(
2522 bytes: Bytes,
2523 schema: &Schema,
2524 file_version: LanceFileVersion,
2525 ) -> Result<Self>
2526 where
2527 Self: Sized,
2528 {
2529 let projection = ReaderProjection::from_whole_schema(schema, file_version);
2530 let footer = FileReader::decode_footer(&bytes)?;
2531
2532 let column_metadata_start = footer.column_meta_start as usize;
2535 let column_metadata_end = footer.global_buff_offsets_start as usize;
2536 let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2537 let column_metadatas =
2538 FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2539
2540 let file_version = LanceFileVersion::try_from_major_minor(
2541 footer.major_version as u32,
2542 footer.minor_version as u32,
2543 )?;
2544
2545 let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version);
2546
2547 Ok(Self {
2548 data: bytes,
2549 num_rows: page_table
2550 .first()
2551 .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2552 .unwrap_or(0),
2553 page_table,
2554 top_level_columns: projection.column_indices,
2555 schema: Arc::new(schema.clone()),
2556 })
2557 }
2558
2559 fn try_from_self_described_lance(bytes: Bytes) -> Result<Self>
2560 where
2561 Self: Sized,
2562 {
2563 let footer = FileReader::decode_footer(&bytes)?;
2564 let file_version = LanceFileVersion::try_from_major_minor(
2565 footer.major_version as u32,
2566 footer.minor_version as u32,
2567 )?;
2568
2569 let gbo_table = FileReader::do_decode_gbo_table(
2570 &bytes.slice(footer.global_buff_offsets_start as usize..),
2571 &footer,
2572 file_version,
2573 )?;
2574 if gbo_table.is_empty() {
2575 return Err(Error::internal(
2576 "File did not contain any global buffers, schema expected".to_string(),
2577 ));
2578 }
2579 let schema_start = gbo_table[0].position as usize;
2580 let schema_size = gbo_table[0].size as usize;
2581
2582 let schema_bytes = bytes.slice(schema_start..(schema_start + schema_size));
2583 let (_, schema) = FileReader::decode_schema(schema_bytes)?;
2584 let projection = ReaderProjection::from_whole_schema(&schema, file_version);
2585
2586 let column_metadata_start = footer.column_meta_start as usize;
2589 let column_metadata_end = footer.global_buff_offsets_start as usize;
2590 let column_metadata_bytes = bytes.slice(column_metadata_start..column_metadata_end);
2591 let column_metadatas =
2592 FileReader::read_all_column_metadata(column_metadata_bytes, &footer)?;
2593
2594 let page_table = FileReader::meta_to_col_infos(&column_metadatas, file_version);
2595
2596 Ok(Self {
2597 data: bytes,
2598 num_rows: page_table
2599 .first()
2600 .map(|col| col.page_infos.iter().map(|page| page.num_rows).sum::<u64>())
2601 .unwrap_or(0),
2602 page_table,
2603 top_level_columns: projection.column_indices,
2604 schema: Arc::new(schema),
2605 })
2606 }
2607}
2608
2609#[cfg(test)]
2610mod tests {
2611 use std::{
2612 collections::{BTreeMap, HashMap},
2613 pin::Pin,
2614 sync::Arc,
2615 };
2616
2617 use arrow_array::{
2618 RecordBatch, UInt32Array,
2619 types::{Float64Type, Int32Type},
2620 };
2621 use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
2622 use bytes::Bytes;
2623 use futures::{StreamExt, prelude::stream::TryStreamExt};
2624 use lance_arrow::{BLOB_META_KEY, RecordBatchExt};
2625 use lance_core::{ArrowResult, datatypes::Schema};
2626 use lance_datagen::{ArrayGeneratorExt, BatchCount, ByteCount, RowCount, array, gen_batch};
2627 use lance_encoding::{
2628 decoder::{
2629 DecodeBatchScheduler, DecoderPlugins, FilterExpression, ReadBatchTask, decode_batch,
2630 },
2631 encoder::{EncodedBatch, EncodingOptions, default_encoding_strategy, encode_batch},
2632 version::LanceFileVersion,
2633 };
2634 use lance_io::{stream::RecordBatchStream, utils::CachedFileSize};
2635 use log::debug;
2636 use rstest::rstest;
2637 use tokio::sync::mpsc;
2638
2639 use crate::reader::{
2640 EncodedBatchReaderExt, FileReader, FileReaderOptions, ProjectedFileReader,
2641 ReaderProjection, validate_field_length, verify_uniform_lengths,
2642 };
2643 use crate::testing::{FsFixture, WrittenFile, test_cache, write_lance_file};
2644 use crate::writer::{EncodedBatchWriteExt, FileWriter, FileWriterOptions};
2645 use lance_encoding::decoder::DecoderConfig;
2646
2647 async fn create_some_file(fs: &FsFixture, version: LanceFileVersion) -> WrittenFile {
2648 let location_type = DataType::Struct(Fields::from(vec![
2649 Field::new("x", DataType::Float64, true),
2650 Field::new("y", DataType::Float64, true),
2651 ]));
2652 let categories_type = DataType::List(Arc::new(Field::new("item", DataType::Utf8, true)));
2653
2654 let mut reader = gen_batch()
2655 .col("score", array::rand::<Float64Type>())
2656 .col("location", array::rand_type(&location_type))
2657 .col("categories", array::rand_type(&categories_type))
2658 .col("binary", array::rand_type(&DataType::Binary));
2659 if version <= LanceFileVersion::V2_0 {
2660 reader = reader.col("large_bin", array::rand_type(&DataType::LargeBinary));
2661 }
2662 let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2663
2664 write_lance_file(
2665 reader,
2666 fs,
2667 FileWriterOptions {
2668 format_version: Some(version),
2669 ..Default::default()
2670 },
2671 )
2672 .await
2673 }
2674
2675 async fn create_wide_direct_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2676 let mut reader = gen_batch();
2677 for column_idx in 0..num_columns {
2678 reader = reader.col(format!("c{column_idx}"), array::step::<Int32Type>());
2679 }
2680 let reader = reader.into_reader_rows(RowCount::from(1000), BatchCount::from(100));
2681
2682 write_lance_file(
2683 reader,
2684 fs,
2685 FileWriterOptions {
2686 format_version: Some(LanceFileVersion::V2_1),
2687 ..Default::default()
2688 },
2689 )
2690 .await
2691 }
2692
2693 async fn create_wide_fixed_size_list_file(fs: &FsFixture, num_columns: usize) -> WrittenFile {
2694 let data_type =
2695 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4);
2696 let mut reader = gen_batch();
2697 for column_idx in 0..num_columns {
2698 reader = reader.col(
2699 format!("c{column_idx}"),
2700 array::rand_type(&data_type).with_random_nulls(0.1),
2701 );
2702 }
2703 let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2704
2705 write_lance_file(
2706 reader,
2707 fs,
2708 FileWriterOptions {
2709 format_version: Some(LanceFileVersion::V2_1),
2710 ..Default::default()
2711 },
2712 )
2713 .await
2714 }
2715
2716 async fn create_wide_structural_file(fs: &FsFixture, num_groups: usize) -> WrittenFile {
2717 let struct_type = DataType::Struct(Fields::from(vec![
2718 Field::new("x", DataType::Int32, true),
2719 Field::new("y", DataType::Int32, true),
2720 ]));
2721 let list_type = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
2722 let mut reader = gen_batch();
2723 for group_idx in 0..num_groups {
2724 reader = reader
2725 .col(
2726 format!("s{group_idx}"),
2727 array::rand_type(&struct_type).with_random_nulls(0.5),
2728 )
2729 .col(
2730 format!("l{group_idx}"),
2731 array::rand_type(&list_type).with_random_nulls(0.5),
2732 );
2733 }
2734 let reader = reader.into_reader_rows(RowCount::from(64), BatchCount::from(4));
2735
2736 write_lance_file(
2737 reader,
2738 fs,
2739 FileWriterOptions {
2740 format_version: Some(LanceFileVersion::V2_1),
2741 ..Default::default()
2742 },
2743 )
2744 .await
2745 }
2746
2747 type Transformer = Box<dyn Fn(&RecordBatch) -> RecordBatch>;
2748
2749 async fn verify_expected(
2750 expected: &[RecordBatch],
2751 mut actual: Pin<Box<dyn RecordBatchStream>>,
2752 read_size: u32,
2753 transform: Option<Transformer>,
2754 ) {
2755 let mut remaining = expected.iter().map(|batch| batch.num_rows()).sum::<usize>() as u32;
2756 let mut expected_iter = expected.iter().map(|batch| {
2757 if let Some(transform) = &transform {
2758 transform(batch)
2759 } else {
2760 batch.clone()
2761 }
2762 });
2763 let mut next_expected = expected_iter.next().unwrap().clone();
2764 while let Some(actual) = actual.next().await {
2765 let mut actual = actual.unwrap();
2766 let mut rows_to_verify = actual.num_rows() as u32;
2767 let expected_length = remaining.min(read_size);
2768 assert_eq!(expected_length, rows_to_verify);
2769
2770 while rows_to_verify > 0 {
2771 let next_slice_len = (next_expected.num_rows() as u32).min(rows_to_verify);
2772 assert_eq!(
2773 next_expected.slice(0, next_slice_len as usize),
2774 actual.slice(0, next_slice_len as usize)
2775 );
2776 remaining -= next_slice_len;
2777 rows_to_verify -= next_slice_len;
2778 if remaining > 0 {
2779 if next_slice_len == next_expected.num_rows() as u32 {
2780 next_expected = expected_iter.next().unwrap().clone();
2781 } else {
2782 next_expected = next_expected.slice(
2783 next_slice_len as usize,
2784 next_expected.num_rows() - next_slice_len as usize,
2785 );
2786 }
2787 }
2788 if rows_to_verify > 0 {
2789 actual = actual.slice(
2790 next_slice_len as usize,
2791 actual.num_rows() - next_slice_len as usize,
2792 );
2793 }
2794 }
2795 }
2796 assert_eq!(remaining, 0);
2797 }
2798
2799 async fn collect_read_tasks(
2800 tasks: Pin<Box<dyn futures::Stream<Item = ReadBatchTask> + Send>>,
2801 readahead: usize,
2802 ) -> Vec<RecordBatch> {
2803 tasks
2804 .map(|task| task.task)
2805 .buffered(readahead)
2806 .try_collect::<Vec<_>>()
2807 .await
2808 .unwrap()
2809 }
2810
2811 #[tokio::test]
2812 async fn test_round_trip() {
2813 let fs = FsFixture::default();
2814
2815 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
2816
2817 for read_size in [32, 1024, 1024 * 1024] {
2818 let file_scheduler = fs
2819 .scheduler
2820 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2821 .await
2822 .unwrap();
2823 let file_reader = FileReader::try_open(
2824 file_scheduler,
2825 None,
2826 Arc::<DecoderPlugins>::default(),
2827 &test_cache(),
2828 FileReaderOptions::default(),
2829 )
2830 .await
2831 .unwrap();
2832
2833 let schema = file_reader.schema();
2834 assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
2835
2836 let batch_stream = file_reader
2837 .read_stream(
2838 lance_io::ReadBatchParams::RangeFull,
2839 read_size,
2840 16,
2841 FilterExpression::no_filter(),
2842 )
2843 .await
2844 .unwrap();
2845
2846 verify_expected(&data, batch_stream, read_size, None).await;
2847 }
2848 }
2849
2850 #[rstest]
2851 #[test_log::test(tokio::test)]
2852 async fn test_encoded_batch_round_trip(
2853 #[values(LanceFileVersion::V2_0)] version: LanceFileVersion,
2855 ) {
2856 let data = gen_batch()
2857 .col("x", array::rand::<Int32Type>())
2858 .col("y", array::rand_utf8(ByteCount::from(16), false))
2859 .into_batch_rows(RowCount::from(10000))
2860 .unwrap();
2861
2862 let lance_schema = Arc::new(Schema::try_from(data.schema().as_ref()).unwrap());
2863
2864 let encoding_options = EncodingOptions {
2865 cache_bytes_per_column: 4096,
2866 max_page_bytes: 32 * 1024 * 1024,
2867 keep_original_array: true,
2868 buffer_alignment: 64,
2869 version,
2870 };
2871
2872 let encoding_strategy = default_encoding_strategy(version);
2873
2874 let encoded_batch = encode_batch(
2875 &data,
2876 lance_schema.clone(),
2877 encoding_strategy.as_ref(),
2878 &encoding_options,
2879 )
2880 .await
2881 .unwrap();
2882
2883 let bytes = encoded_batch.try_to_self_described_lance(version).unwrap();
2885
2886 let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap();
2887
2888 let decoded = decode_batch(
2889 &decoded_batch,
2890 &FilterExpression::no_filter(),
2891 Arc::<DecoderPlugins>::default(),
2892 false,
2893 version,
2894 None,
2895 )
2896 .await
2897 .unwrap();
2898
2899 assert_eq!(data, decoded);
2900
2901 let bytes = encoded_batch.try_to_mini_lance(version).unwrap();
2903 let decoded_batch =
2904 EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref(), LanceFileVersion::V2_0)
2905 .unwrap();
2906 let decoded = decode_batch(
2907 &decoded_batch,
2908 &FilterExpression::no_filter(),
2909 Arc::<DecoderPlugins>::default(),
2910 false,
2911 version,
2912 None,
2913 )
2914 .await
2915 .unwrap();
2916
2917 assert_eq!(data, decoded);
2918 }
2919
2920 #[rstest]
2921 #[test_log::test(tokio::test)]
2922 async fn test_projection(
2923 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
2924 version: LanceFileVersion,
2925 ) {
2926 let fs = FsFixture::default();
2927
2928 let written_file = create_some_file(&fs, version).await;
2929 let file_scheduler = fs
2930 .scheduler
2931 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2932 .await
2933 .unwrap();
2934
2935 let field_id_mapping = written_file
2936 .field_id_mapping
2937 .iter()
2938 .copied()
2939 .collect::<BTreeMap<_, _>>();
2940
2941 let empty_projection = ReaderProjection {
2942 column_indices: Vec::default(),
2943 schema: Arc::new(Schema::default()),
2944 };
2945
2946 for columns in [
2947 vec!["score"],
2948 vec!["location"],
2949 vec!["categories"],
2950 vec!["score.x"],
2951 vec!["score", "categories"],
2952 vec!["score", "location"],
2953 vec!["location", "categories"],
2954 vec!["score.y", "location", "categories"],
2955 ] {
2956 debug!("Testing round trip with projection {:?}", columns);
2957 for use_field_ids in [true, false] {
2958 let file_reader = FileReader::try_open(
2960 file_scheduler.clone(),
2961 None,
2962 Arc::<DecoderPlugins>::default(),
2963 &test_cache(),
2964 FileReaderOptions::default(),
2965 )
2966 .await
2967 .unwrap();
2968
2969 let projected_schema = written_file.schema.project(&columns).unwrap();
2970 let projection = if use_field_ids {
2971 ReaderProjection::from_field_ids(
2972 file_reader.metadata.version(),
2973 &projected_schema,
2974 &field_id_mapping,
2975 )
2976 .unwrap()
2977 } else {
2978 ReaderProjection::from_column_names(
2979 file_reader.metadata.version(),
2980 &written_file.schema,
2981 &columns,
2982 )
2983 .unwrap()
2984 };
2985
2986 let batch_stream = file_reader
2987 .read_stream_projected(
2988 lance_io::ReadBatchParams::RangeFull,
2989 1024,
2990 16,
2991 projection.clone(),
2992 FilterExpression::no_filter(),
2993 )
2994 .await
2995 .unwrap();
2996
2997 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
2998 verify_expected(
2999 &written_file.data,
3000 batch_stream,
3001 1024,
3002 Some(Box::new(move |batch: &RecordBatch| {
3003 batch.project_by_schema(&projection_arrow).unwrap()
3004 })),
3005 )
3006 .await;
3007
3008 let file_reader = FileReader::try_open(
3010 file_scheduler.clone(),
3011 Some(projection.clone()),
3012 Arc::<DecoderPlugins>::default(),
3013 &test_cache(),
3014 FileReaderOptions::default(),
3015 )
3016 .await
3017 .unwrap();
3018
3019 let batch_stream = file_reader
3020 .read_stream(
3021 lance_io::ReadBatchParams::RangeFull,
3022 1024,
3023 16,
3024 FilterExpression::no_filter(),
3025 )
3026 .await
3027 .unwrap();
3028
3029 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3030 verify_expected(
3031 &written_file.data,
3032 batch_stream,
3033 1024,
3034 Some(Box::new(move |batch: &RecordBatch| {
3035 batch.project_by_schema(&projection_arrow).unwrap()
3036 })),
3037 )
3038 .await;
3039
3040 assert!(
3041 file_reader
3042 .read_stream_projected(
3043 lance_io::ReadBatchParams::RangeFull,
3044 1024,
3045 16,
3046 empty_projection.clone(),
3047 FilterExpression::no_filter(),
3048 )
3049 .await
3050 .is_err()
3051 );
3052 }
3053 }
3054
3055 assert!(
3056 FileReader::try_open(
3057 file_scheduler.clone(),
3058 Some(empty_projection),
3059 Arc::<DecoderPlugins>::default(),
3060 &test_cache(),
3061 FileReaderOptions::default(),
3062 )
3063 .await
3064 .is_err()
3065 );
3066
3067 let arrow_schema = ArrowSchema::new(vec![
3068 Field::new("x", DataType::Int32, true),
3069 Field::new("y", DataType::Int32, true),
3070 ]);
3071 let schema = Schema::try_from(&arrow_schema).unwrap();
3072
3073 let projection_with_dupes = ReaderProjection {
3074 column_indices: vec![0, 0],
3075 schema: Arc::new(schema),
3076 };
3077
3078 assert!(
3079 FileReader::try_open(
3080 file_scheduler.clone(),
3081 Some(projection_with_dupes),
3082 Arc::<DecoderPlugins>::default(),
3083 &test_cache(),
3084 FileReaderOptions::default(),
3085 )
3086 .await
3087 .is_err()
3088 );
3089 }
3090
3091 #[tokio::test]
3092 async fn test_lazy_reader_direct_projection_matches_eager_reader() {
3093 let fs = FsFixture::default();
3094 let written_file = create_wide_direct_file(&fs, 16).await;
3095
3096 let file_scheduler = fs
3097 .scheduler
3098 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3099 .await
3100 .unwrap();
3101 let projection = ReaderProjection::from_column_names(
3102 LanceFileVersion::V2_1,
3103 &written_file.schema,
3104 &["c10"],
3105 )
3106 .unwrap();
3107
3108 let eager_reader = FileReader::try_open(
3109 file_scheduler.clone(),
3110 None,
3111 Arc::<DecoderPlugins>::default(),
3112 &test_cache(),
3113 FileReaderOptions::default(),
3114 )
3115 .await
3116 .unwrap();
3117 let expected = eager_reader
3118 .read_stream_projected(
3119 lance_io::ReadBatchParams::RangeFull,
3120 127,
3121 16,
3122 projection.clone(),
3123 FilterExpression::no_filter(),
3124 )
3125 .await
3126 .unwrap()
3127 .try_collect::<Vec<_>>()
3128 .await
3129 .unwrap();
3130
3131 let cache = test_cache();
3132 let lazy_reader = ProjectedFileReader::try_open(
3133 file_scheduler,
3134 Some(projection.clone()),
3135 Arc::<DecoderPlugins>::default(),
3136 &cache,
3137 FileReaderOptions::default(),
3138 )
3139 .await
3140 .unwrap();
3141 let tasks = lazy_reader
3142 .read_tasks(
3143 lance_io::ReadBatchParams::RangeFull,
3144 127,
3145 None,
3146 FilterExpression::no_filter(),
3147 )
3148 .await
3149 .unwrap();
3150 let actual = collect_read_tasks(tasks, 16).await;
3151
3152 assert_eq!(expected, actual);
3153 }
3154
3155 #[tokio::test]
3156 async fn test_lazy_reader_loads_only_requested_column_metadata() {
3157 let fs = FsFixture::default();
3158 let written_file = create_wide_direct_file(&fs, 512).await;
3159
3160 let projection = ReaderProjection::from_column_names(
3161 LanceFileVersion::V2_1,
3162 &written_file.schema,
3163 &["c0"],
3164 )
3165 .unwrap();
3166 let file_scheduler = fs
3167 .scheduler
3168 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3169 .await
3170 .unwrap();
3171 let lazy_reader = ProjectedFileReader::try_open(
3172 file_scheduler,
3173 Some(projection.clone()),
3174 Arc::<DecoderPlugins>::default(),
3175 &test_cache(),
3176 FileReaderOptions::default(),
3177 )
3178 .await
3179 .unwrap();
3180 let selected_column = projection.column_indices[0] as usize;
3181 let requested_metadata_bytes = lazy_reader
3182 .metadata_index()
3183 .unwrap()
3184 .column_metadata_offsets[selected_column]
3185 .1;
3186 let total_metadata_bytes = lazy_reader
3187 .metadata_index()
3188 .unwrap()
3189 .column_metadata_offsets
3190 .iter()
3191 .map(|(_, length)| *length)
3192 .sum::<u64>();
3193 assert!(
3194 total_metadata_bytes > 8 * fs.object_store.block_size() as u64,
3195 "test file metadata is too small to prove lazy loading: {total_metadata_bytes} bytes"
3196 );
3197
3198 fs.object_store.io_stats_incremental();
3199 let tasks = lazy_reader
3200 .read_tasks(
3201 lance_io::ReadBatchParams::Range(0..0),
3202 1024,
3203 Some(projection.clone()),
3204 FilterExpression::no_filter(),
3205 )
3206 .await
3207 .unwrap();
3208 let batches = collect_read_tasks(tasks, 1).await;
3209 assert!(batches.is_empty());
3210
3211 let stats = fs.object_store.io_stats_incremental();
3212 assert!(
3213 stats.read_bytes < total_metadata_bytes / 2,
3214 "lazy read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3215 stats.read_bytes,
3216 requested_metadata_bytes,
3217 total_metadata_bytes
3218 );
3219
3220 fs.object_store.io_stats_incremental();
3221 let tasks = lazy_reader
3222 .read_tasks(
3223 lance_io::ReadBatchParams::Range(0..0),
3224 1024,
3225 Some(projection),
3226 FilterExpression::no_filter(),
3227 )
3228 .await
3229 .unwrap();
3230 let batches = collect_read_tasks(tasks, 1).await;
3231 assert!(batches.is_empty());
3232
3233 let stats = fs.object_store.io_stats_incremental();
3234 assert_eq!(
3235 stats.read_iops, 0,
3236 "cached column metadata should avoid repeat metadata I/O"
3237 );
3238 assert_eq!(
3239 stats.read_bytes, 0,
3240 "cached column metadata should avoid repeat metadata reads"
3241 );
3242 }
3243
3244 async fn assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3245 fs: &FsFixture,
3246 projection: ReaderProjection,
3247 shape: &str,
3248 ) -> Vec<RecordBatch> {
3249 let file_scheduler = fs
3250 .scheduler
3251 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3252 .await
3253 .unwrap();
3254 let eager_reader = FileReader::try_open(
3255 file_scheduler.clone(),
3256 None,
3257 Arc::<DecoderPlugins>::default(),
3258 &test_cache(),
3259 FileReaderOptions::default(),
3260 )
3261 .await
3262 .unwrap();
3263 let expected = eager_reader
3264 .read_stream_projected(
3265 lance_io::ReadBatchParams::RangeFull,
3266 127,
3267 16,
3268 projection.clone(),
3269 FilterExpression::no_filter(),
3270 )
3271 .await
3272 .unwrap()
3273 .try_collect::<Vec<_>>()
3274 .await
3275 .unwrap();
3276
3277 let cache = test_cache();
3278 let lazy_reader = ProjectedFileReader::try_open(
3279 file_scheduler,
3280 Some(projection.clone()),
3281 Arc::<DecoderPlugins>::default(),
3282 &cache,
3283 FileReaderOptions::default(),
3284 )
3285 .await
3286 .unwrap();
3287 let metadata_index = lazy_reader.metadata_index().unwrap();
3288 let requested_metadata_bytes = projection
3289 .column_indices
3290 .iter()
3291 .map(|column_index| metadata_index.column_metadata_offsets[*column_index as usize].1)
3292 .sum::<u64>();
3293 let total_metadata_bytes = metadata_index
3294 .column_metadata_offsets
3295 .iter()
3296 .map(|(_, length)| *length)
3297 .sum::<u64>();
3298 assert!(total_metadata_bytes > requested_metadata_bytes * 8);
3299
3300 fs.object_store.io_stats_incremental();
3301 let tasks = lazy_reader
3302 .read_tasks(
3303 lance_io::ReadBatchParams::Range(0..0),
3304 127,
3305 None,
3306 FilterExpression::no_filter(),
3307 )
3308 .await
3309 .unwrap();
3310 assert!(collect_read_tasks(tasks, 1).await.is_empty());
3311 let metadata_stats = fs.object_store.io_stats_incremental();
3312 assert!(
3313 metadata_stats.read_bytes < total_metadata_bytes / 2,
3314 "lazy {shape} read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3315 metadata_stats.read_bytes,
3316 requested_metadata_bytes,
3317 total_metadata_bytes
3318 );
3319
3320 let tasks = lazy_reader
3321 .read_tasks(
3322 lance_io::ReadBatchParams::RangeFull,
3323 127,
3324 None,
3325 FilterExpression::no_filter(),
3326 )
3327 .await
3328 .unwrap();
3329 let actual = collect_read_tasks(tasks, 16).await;
3330 assert_eq!(expected, actual);
3331 actual
3332 }
3333
3334 #[tokio::test]
3335 async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() {
3336 let fs = FsFixture::default();
3337 let written_file = create_wide_fixed_size_list_file(&fs, 512).await;
3338 let projection = ReaderProjection::from_column_names(
3339 LanceFileVersion::V2_1,
3340 &written_file.schema,
3341 &["c17", "c509"],
3342 )
3343 .unwrap();
3344 assert!(ProjectedFileReader::supports_projection(
3345 &projection,
3346 LanceFileVersion::V2_1
3347 ));
3348 assert!(!ProjectedFileReader::supports_projection(
3349 &projection,
3350 LanceFileVersion::V2_0
3351 ));
3352 assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3353 &fs,
3354 projection,
3355 "fixed-size-list",
3356 )
3357 .await;
3358 }
3359
3360 #[tokio::test]
3361 async fn test_lazy_reader_nested_projection_compacts_physical_columns() {
3362 let fs = FsFixture::default();
3363 let written_file = create_wide_structural_file(&fs, 128).await;
3364 let projection = ReaderProjection::from_column_names(
3365 LanceFileVersion::V2_1,
3366 &written_file.schema,
3367 &["s97.y", "l4", "s3"],
3368 )
3369 .unwrap();
3370
3371 assert_eq!(
3372 projection
3373 .schema
3374 .fields
3375 .iter()
3376 .map(|field| field.name.as_str())
3377 .collect::<Vec<_>>(),
3378 vec!["s97", "l4", "s3"]
3379 );
3380 assert_eq!(projection.schema.fields[0].children.len(), 1);
3381 assert_eq!(projection.schema.fields[0].children[0].name, "y");
3382 assert_eq!(projection.schema.fields[2].children.len(), 2);
3383 assert_eq!(projection.column_indices.len(), 4);
3384 assert!(
3385 projection
3386 .column_indices
3387 .windows(2)
3388 .any(|indices| indices[0] > indices[1]),
3389 "the projection must reorder physical columns to exercise compact remapping"
3390 );
3391 assert!(ProjectedFileReader::supports_projection(
3392 &projection,
3393 LanceFileVersion::V2_1
3394 ));
3395 let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3396 &fs, projection, "nested",
3397 )
3398 .await;
3399 assert!(
3400 actual
3401 .iter()
3402 .flat_map(|batch| batch.columns())
3403 .any(|column| column.null_count() > 0),
3404 "the structural projection must exercise nullable arrays"
3405 );
3406 }
3407
3408 #[rstest]
3409 #[case::before_metadata_region(90, 5)]
3410 #[case::after_metadata_region(190, 20)]
3411 fn test_decode_cmo_table_rejects_out_of_range_offsets(
3412 #[case] position: u64,
3413 #[case] length: u64,
3414 ) {
3415 let mut cmo_table = [0; 16];
3416 cmo_table[0..8].copy_from_slice(&position.to_le_bytes());
3417 cmo_table[8..16].copy_from_slice(&length.to_le_bytes());
3418 let footer = super::Footer {
3419 column_meta_start: 100,
3420 column_meta_offsets_start: 200,
3421 global_buff_offsets_start: 200,
3422 num_global_buffers: 0,
3423 num_columns: 1,
3424 major_version: 2,
3425 minor_version: 1,
3426 };
3427
3428 let err = FileReader::decode_cmo_table(Bytes::copy_from_slice(&cmo_table), &footer)
3429 .expect_err("out-of-range CMO entries must be rejected");
3430 assert!(
3431 matches!(err, lance_core::Error::InvalidInput { .. }),
3432 "expected InvalidInput, got {err:?}"
3433 );
3434 }
3435
3436 #[rstest]
3437 #[case::blob(BLOB_META_KEY)]
3438 #[case::packed_struct("lance-encoding:packed")]
3439 #[tokio::test]
3440 async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) {
3441 let fs = FsFixture::default();
3442 let written_file = create_some_file(&fs, LanceFileVersion::V2_1).await;
3443
3444 let ordinary_projection = ReaderProjection::from_column_names(
3445 LanceFileVersion::V2_1,
3446 &written_file.schema,
3447 &["location.x"],
3448 )
3449 .unwrap();
3450 assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1);
3451 assert!(ProjectedFileReader::supports_projection(
3452 &ordinary_projection,
3453 LanceFileVersion::V2_1
3454 ));
3455
3456 let file_scheduler = fs
3457 .scheduler
3458 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3459 .await
3460 .unwrap();
3461 let err = ProjectedFileReader::try_open(
3462 file_scheduler.clone(),
3463 None,
3464 Arc::<DecoderPlugins>::default(),
3465 &test_cache(),
3466 FileReaderOptions::default(),
3467 )
3468 .await
3469 .unwrap_err();
3470 assert!(
3471 matches!(err, lance_core::Error::InvalidInput { .. }),
3472 "expected InvalidInput, got {err:?}"
3473 );
3474
3475 let mut projection = ordinary_projection;
3476 Arc::make_mut(&mut projection.schema).fields[0]
3477 .metadata
3478 .insert(metadata_key.to_string(), "true".to_string());
3479 assert!(!ProjectedFileReader::supports_projection(
3480 &projection,
3481 LanceFileVersion::V2_1
3482 ));
3483
3484 let err = ProjectedFileReader::try_open(
3485 file_scheduler,
3486 Some(projection),
3487 Arc::<DecoderPlugins>::default(),
3488 &test_cache(),
3489 FileReaderOptions::default(),
3490 )
3491 .await
3492 .unwrap_err();
3493 assert!(
3494 matches!(err, lance_core::Error::NotSupported { .. }),
3495 "expected NotSupported for {metadata_key}, got {err:?}"
3496 );
3497 }
3498
3499 #[tokio::test]
3506 async fn test_lazy_reader_validates_unequal_length_projection() {
3507 use arrow_array::Int32Array;
3508 use lance_io::ReadBatchParams;
3509
3510 let arrow_schema = Arc::new(ArrowSchema::new(vec![
3511 Field::new("a", DataType::Int32, true),
3512 Field::new("c", DataType::Int32, true),
3513 ]));
3514 let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
3515
3516 let fs = FsFixture::default();
3517 let options = FileWriterOptions {
3518 format_version: Some(LanceFileVersion::V2_1),
3519 ..Default::default()
3520 };
3521 let mut writer = FileWriter::try_new(
3522 fs.object_store.create(&fs.tmp_path).await.unwrap(),
3523 lance_schema.clone(),
3524 options,
3525 )
3526 .unwrap();
3527 writer
3529 .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])))
3530 .await
3531 .unwrap();
3532 writer
3533 .write_column(1, Arc::new(Int32Array::from(vec![100])))
3534 .await
3535 .unwrap();
3536 writer.finish().await.unwrap();
3537
3538 let file_scheduler = fs
3539 .scheduler
3540 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3541 .await
3542 .unwrap();
3543 let cache = test_cache();
3544 let open_indexed = |names: &[&str]| {
3545 let projection =
3546 ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names)
3547 .unwrap();
3548 ProjectedFileReader::try_open(
3549 file_scheduler.clone(),
3550 Some(projection),
3551 Arc::<DecoderPlugins>::default(),
3552 &cache,
3553 FileReaderOptions::default(),
3554 )
3555 };
3556
3557 let lazy = open_indexed(&["a", "c"]).await.unwrap();
3560 let err = match lazy
3563 .read_tasks(
3564 ReadBatchParams::RangeFull,
3565 1024,
3566 None,
3567 FilterExpression::no_filter(),
3568 )
3569 .await
3570 {
3571 Ok(_) => panic!("expected the mismatched-length projection to be rejected"),
3572 Err(e) => e.to_string(),
3573 };
3574 assert!(
3575 err.contains("a=5") && err.contains("c=1"),
3576 "error should name each column's length, got: {err}"
3577 );
3578
3579 let lazy = open_indexed(&["c"]).await.unwrap();
3582 let tasks = lazy
3583 .read_tasks(
3584 ReadBatchParams::RangeFull,
3585 1024,
3586 None,
3587 FilterExpression::no_filter(),
3588 )
3589 .await
3590 .unwrap();
3591 let batches = collect_read_tasks(tasks, 16).await;
3592 let values: Vec<Option<i32>> = batches
3593 .iter()
3594 .flat_map(|b| {
3595 b.column(0)
3596 .as_any()
3597 .downcast_ref::<Int32Array>()
3598 .unwrap()
3599 .iter()
3600 .collect::<Vec<_>>()
3601 })
3602 .collect();
3603 assert_eq!(values, vec![Some(100)]);
3604 }
3605
3606 #[test_log::test(tokio::test)]
3607 async fn test_compressing_buffer() {
3608 let fs = FsFixture::default();
3609
3610 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
3611 let file_scheduler = fs
3612 .scheduler
3613 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3614 .await
3615 .unwrap();
3616
3617 let file_reader = FileReader::try_open(
3619 file_scheduler.clone(),
3620 None,
3621 Arc::<DecoderPlugins>::default(),
3622 &test_cache(),
3623 FileReaderOptions::default(),
3624 )
3625 .await
3626 .unwrap();
3627
3628 let mut projection = written_file.schema.project(&["score"]).unwrap();
3629 for field in projection.fields.iter_mut() {
3630 field
3631 .metadata
3632 .insert("lance:compression".to_string(), "zstd".to_string());
3633 }
3634 let projection = ReaderProjection {
3635 column_indices: projection.fields.iter().map(|f| f.id as u32).collect(),
3636 schema: Arc::new(projection),
3637 };
3638
3639 let batch_stream = file_reader
3640 .read_stream_projected(
3641 lance_io::ReadBatchParams::RangeFull,
3642 1024,
3643 16,
3644 projection.clone(),
3645 FilterExpression::no_filter(),
3646 )
3647 .await
3648 .unwrap();
3649
3650 let projection_arrow = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
3651 verify_expected(
3652 &written_file.data,
3653 batch_stream,
3654 1024,
3655 Some(Box::new(move |batch: &RecordBatch| {
3656 batch.project_by_schema(&projection_arrow).unwrap()
3657 })),
3658 )
3659 .await;
3660 }
3661
3662 #[tokio::test]
3663 async fn test_read_all() {
3664 let fs = FsFixture::default();
3665 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3666 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3667
3668 let file_scheduler = fs
3669 .scheduler
3670 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3671 .await
3672 .unwrap();
3673 let file_reader = FileReader::try_open(
3674 file_scheduler.clone(),
3675 None,
3676 Arc::<DecoderPlugins>::default(),
3677 &test_cache(),
3678 FileReaderOptions::default(),
3679 )
3680 .await
3681 .unwrap();
3682
3683 let batches = file_reader
3684 .read_stream(
3685 lance_io::ReadBatchParams::RangeFull,
3686 total_rows as u32,
3687 16,
3688 FilterExpression::no_filter(),
3689 )
3690 .await
3691 .unwrap()
3692 .try_collect::<Vec<_>>()
3693 .await
3694 .unwrap();
3695 assert_eq!(batches.len(), 1);
3696 assert_eq!(batches[0].num_rows(), total_rows);
3697 }
3698
3699 #[rstest]
3700 #[tokio::test]
3701 async fn test_blocking_take(
3702 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
3703 version: LanceFileVersion,
3704 ) {
3705 let fs = FsFixture::default();
3706 let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await;
3707 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3708
3709 let file_scheduler = fs
3710 .scheduler
3711 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3712 .await
3713 .unwrap();
3714 let file_reader = FileReader::try_open(
3715 file_scheduler.clone(),
3716 Some(ReaderProjection::from_column_names(version, &schema, &["score"]).unwrap()),
3717 Arc::<DecoderPlugins>::default(),
3718 &test_cache(),
3719 FileReaderOptions::default(),
3720 )
3721 .await
3722 .unwrap();
3723
3724 let batches = tokio::task::spawn_blocking(move || {
3725 file_reader
3726 .read_stream_projected_blocking(
3727 lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![0, 1, 2, 3, 4])),
3728 total_rows as u32,
3729 None,
3730 FilterExpression::no_filter(),
3731 )
3732 .unwrap()
3733 .collect::<ArrowResult<Vec<_>>>()
3734 .unwrap()
3735 })
3736 .await
3737 .unwrap();
3738
3739 assert_eq!(batches.len(), 1);
3740 assert_eq!(batches[0].num_rows(), 5);
3741 assert_eq!(batches[0].num_columns(), 1);
3742 }
3743
3744 #[tokio::test(flavor = "multi_thread")]
3745 async fn test_drop_in_progress() {
3746 let fs = FsFixture::default();
3747 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3748 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3749
3750 let file_scheduler = fs
3751 .scheduler
3752 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3753 .await
3754 .unwrap();
3755 let file_reader = FileReader::try_open(
3756 file_scheduler.clone(),
3757 None,
3758 Arc::<DecoderPlugins>::default(),
3759 &test_cache(),
3760 FileReaderOptions::default(),
3761 )
3762 .await
3763 .unwrap();
3764
3765 let mut batches = file_reader
3766 .read_stream(
3767 lance_io::ReadBatchParams::RangeFull,
3768 (total_rows / 10) as u32,
3769 16,
3770 FilterExpression::no_filter(),
3771 )
3772 .await
3773 .unwrap();
3774
3775 drop(file_reader);
3776
3777 let batch = batches.next().await.unwrap().unwrap();
3778 assert!(batch.num_rows() > 0);
3779
3780 drop(batches);
3782 }
3783
3784 #[tokio::test]
3785 async fn drop_while_scheduling() {
3786 let fs = FsFixture::default();
3796 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
3797 let total_rows = written_file
3798 .data
3799 .iter()
3800 .map(|batch| batch.num_rows())
3801 .sum::<usize>();
3802
3803 let file_scheduler = fs
3804 .scheduler
3805 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3806 .await
3807 .unwrap();
3808 let file_reader = FileReader::try_open(
3809 file_scheduler.clone(),
3810 None,
3811 Arc::<DecoderPlugins>::default(),
3812 &test_cache(),
3813 FileReaderOptions::default(),
3814 )
3815 .await
3816 .unwrap();
3817
3818 let projection =
3819 ReaderProjection::from_whole_schema(&written_file.schema, LanceFileVersion::V2_0);
3820 let column_infos = file_reader
3821 .collect_columns_from_projection(&projection)
3822 .unwrap();
3823 let mut decode_scheduler = DecodeBatchScheduler::try_new(
3824 &projection.schema,
3825 &projection.column_indices,
3826 &column_infos,
3827 &vec![],
3828 total_rows as u64,
3829 Arc::<DecoderPlugins>::default(),
3830 file_reader.core.scheduler.clone(),
3831 test_cache(),
3832 &FilterExpression::no_filter(),
3833 &DecoderConfig::default(),
3834 )
3835 .await
3836 .unwrap();
3837
3838 let range = 0..total_rows as u64;
3839
3840 let (tx, rx) = mpsc::unbounded_channel();
3841
3842 drop(rx);
3844
3845 decode_scheduler.schedule_range(
3847 range,
3848 &FilterExpression::no_filter(),
3849 tx,
3850 file_reader.core.scheduler.clone(),
3851 )
3852 }
3853
3854 #[tokio::test]
3855 async fn test_read_empty_range() {
3856 let fs = FsFixture::default();
3857 create_some_file(&fs, LanceFileVersion::V2_0).await;
3858
3859 let file_scheduler = fs
3860 .scheduler
3861 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3862 .await
3863 .unwrap();
3864 let file_reader = FileReader::try_open(
3865 file_scheduler.clone(),
3866 None,
3867 Arc::<DecoderPlugins>::default(),
3868 &test_cache(),
3869 FileReaderOptions::default(),
3870 )
3871 .await
3872 .unwrap();
3873
3874 let batches = file_reader
3876 .read_stream(
3877 lance_io::ReadBatchParams::Range(0..0),
3878 1024,
3879 16,
3880 FilterExpression::no_filter(),
3881 )
3882 .await
3883 .unwrap()
3884 .try_collect::<Vec<_>>()
3885 .await
3886 .unwrap();
3887
3888 assert_eq!(batches.len(), 0);
3889
3890 let batches = file_reader
3892 .read_stream(
3893 lance_io::ReadBatchParams::Ranges(Arc::new([0..1, 2..2])),
3894 1024,
3895 16,
3896 FilterExpression::no_filter(),
3897 )
3898 .await
3899 .unwrap()
3900 .try_collect::<Vec<_>>()
3901 .await
3902 .unwrap();
3903 assert_eq!(batches.len(), 1);
3904 }
3905
3906 async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) {
3907 let lance_schema =
3908 lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new(
3909 "foo",
3910 DataType::Int32,
3911 true,
3912 )]))
3913 .unwrap();
3914
3915 let mut file_writer = FileWriter::try_new(
3916 fs.object_store.create(&fs.tmp_path).await.unwrap(),
3917 lance_schema,
3918 FileWriterOptions::default(),
3919 )
3920 .unwrap();
3921
3922 let buf_index = file_writer.add_global_buffer(buffer).await.unwrap();
3923 assert_eq!(buf_index, 1);
3924
3925 file_writer.finish().await.unwrap();
3926 }
3927
3928 #[rstest]
3932 #[case::within_tail_window(true)]
3933 #[case::outside_tail_window(false)]
3934 #[tokio::test]
3935 async fn test_read_global_buffer(#[case] within_window: bool) {
3936 let fs = FsFixture::default();
3937
3938 let block_size = fs.object_store.block_size();
3939 let buffer = if within_window {
3940 Bytes::from_static(b"hello")
3941 } else {
3942 Bytes::from(vec![7u8; 2 * block_size])
3943 };
3944 let expected_read_iops = if within_window { 0 } else { 1 };
3945
3946 write_file_with_global_buffer(&fs, buffer.clone()).await;
3947
3948 let file_scheduler = fs
3949 .scheduler
3950 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3951 .await
3952 .unwrap();
3953 let file_reader = FileReader::try_open(
3954 file_scheduler,
3955 None,
3956 Arc::<DecoderPlugins>::default(),
3957 &test_cache(),
3958 FileReaderOptions::default(),
3959 )
3960 .await
3961 .unwrap();
3962
3963 let retained = &file_reader.metadata().retained_global_buffers;
3966 assert!(!retained.contains_key(&0), "schema must not be retained");
3967 assert_eq!(retained.contains_key(&1), within_window);
3968
3969 fs.object_store.io_stats_incremental();
3971
3972 let buf = file_reader.read_global_buffer(1).await.unwrap();
3973 assert_eq!(buf, buffer);
3974
3975 let stats = fs.object_store.io_stats_incremental();
3976 assert_eq!(stats.read_iops, expected_read_iops);
3977 }
3978
3979 #[tokio::test]
3982 async fn test_read_global_buffer_no_user_buffers() {
3983 let fs = FsFixture::default();
3984 create_some_file(&fs, LanceFileVersion::V2_1).await;
3985
3986 let file_scheduler = fs
3987 .scheduler
3988 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3989 .await
3990 .unwrap();
3991 let file_reader = FileReader::try_open(
3992 file_scheduler,
3993 None,
3994 Arc::<DecoderPlugins>::default(),
3995 &test_cache(),
3996 FileReaderOptions::default(),
3997 )
3998 .await
3999 .unwrap();
4000
4001 let metadata = file_reader.metadata();
4002 assert_eq!(metadata.file_buffers.len(), 1, "expected only the schema");
4003 assert!(
4004 metadata.retained_global_buffers.is_empty(),
4005 "a file with no user global buffers must retain nothing"
4006 );
4007 }
4008
4009 #[rstest]
4010 #[tokio::test]
4011 async fn test_deep_size_of_includes_column_metadata(
4012 #[values(
4013 LanceFileVersion::V2_0,
4014 LanceFileVersion::V2_1,
4015 LanceFileVersion::V2_2,
4016 LanceFileVersion::V2_3
4017 )]
4018 version: LanceFileVersion,
4019 ) {
4020 use lance_core::deepsize::DeepSizeOf;
4025
4026 let fs = FsFixture::default();
4027 let _written = create_some_file(&fs, version).await;
4028 let cache = test_cache();
4029 let file_scheduler = fs
4030 .scheduler
4031 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4032 .await
4033 .unwrap();
4034 let file_reader = FileReader::try_open(
4035 file_scheduler,
4036 None,
4037 Arc::<DecoderPlugins>::default(),
4038 &cache,
4039 FileReaderOptions::default(),
4040 )
4041 .await
4042 .unwrap();
4043
4044 let metadata = file_reader.metadata();
4045 let deep_size = metadata.deep_size_of();
4046
4047 assert!(
4053 deep_size > 1024,
4054 "deep_size_of ({deep_size}) is suspiciously small — \
4055 column_metadatas and column_infos may not be accounted for"
4056 );
4057
4058 assert!(
4060 !metadata.column_metadatas.is_empty(),
4061 "Expected non-empty column_metadatas"
4062 );
4063
4064 let num_columns = metadata.column_metadatas.len();
4067 assert!(
4068 deep_size > num_columns * 50,
4069 "deep_size_of ({deep_size}) should scale with column count ({num_columns})"
4070 );
4071 }
4072
4073 #[tokio::test]
4074 async fn test_read_global_buffer_out_of_range() {
4075 let fs = FsFixture::default();
4076
4077 write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
4078
4079 let file_scheduler = fs
4080 .scheduler
4081 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4082 .await
4083 .unwrap();
4084 let file_reader = FileReader::try_open(
4085 file_scheduler,
4086 None,
4087 Arc::<DecoderPlugins>::default(),
4088 &test_cache(),
4089 FileReaderOptions::default(),
4090 )
4091 .await
4092 .unwrap();
4093
4094 let err = file_reader.read_global_buffer(2).await.unwrap_err();
4097 assert!(
4098 matches!(err, lance_core::Error::InvalidInput { .. }),
4099 "expected InvalidInput, got: {err:?}"
4100 );
4101 let msg = err.to_string();
4102 assert!(msg.contains('2'), "error should mention the index: {msg}");
4103 }
4104
4105 #[rstest]
4110 fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) {
4111 let run = |dt: DataType, indices: &[u32], lengths: Vec<u64>| -> lance_core::Result<u64> {
4112 let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]);
4113 let schema = Schema::try_from(&arrow).unwrap();
4114 let column_len = |c: usize| Ok(lengths[c]);
4115 let mut cursor = 0usize;
4116 let mut field_lengths = Vec::new();
4117 for field in &schema.fields {
4118 let rows = validate_field_length(
4119 field,
4120 is_structural,
4121 true,
4122 indices,
4123 &mut cursor,
4124 &column_len,
4125 )?;
4126 field_lengths.push((field.name.as_str(), rows));
4127 }
4128 verify_uniform_lengths(&field_lengths)
4129 };
4130
4131 let struct_ty = || {
4132 DataType::Struct(Fields::from(vec![
4133 Field::new("a", DataType::Int32, true),
4134 Field::new("b", DataType::Int32, true),
4135 ]))
4136 };
4137
4138 let (indices, equal, unequal): (&[u32], Vec<u64>, Vec<u64>) = if is_structural {
4141 (&[0, 1], vec![5, 5], vec![5, 3])
4142 } else {
4143 (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3])
4144 };
4145
4146 assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5);
4147
4148 let err = run(struct_ty(), indices, unequal).unwrap_err();
4149 let msg = err.to_string();
4150 assert!(
4151 msg.contains("differing lengths") && msg.contains('b'),
4152 "expected a child-length error naming 'b', got: {msg}"
4153 );
4154 }
4155
4156 #[test]
4157 fn test_validate_v2_0_unloaded_blob_projection_is_opaque() {
4158 let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]);
4159 let arrow = ArrowSchema::new(vec![
4160 Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata),
4161 ]);
4162 let mut schema = Schema::try_from(&arrow).unwrap();
4163 schema.fields[0].unloaded_mut();
4164 let projection = ReaderProjection {
4165 schema: Arc::new(schema),
4166 column_indices: vec![0],
4167 };
4168 let column_len = |column: usize| {
4169 assert_eq!(column, 0);
4170 Ok(3)
4171 };
4172 let mut cursor = 0usize;
4173
4174 let rows = validate_field_length(
4175 &projection.schema.fields[0],
4176 false,
4177 true,
4178 &projection.column_indices,
4179 &mut cursor,
4180 &column_len,
4181 )
4182 .unwrap();
4183
4184 assert_eq!(rows, 3);
4185 assert_eq!(cursor, 1);
4186 }
4187
4188 #[test]
4189 fn test_validate_length_list_and_empty_struct() {
4190 let validate = |dt: DataType,
4191 is_structural: bool,
4192 indices: &[u32],
4193 lengths: Vec<u64>|
4194 -> lance_core::Result<u64> {
4195 let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]);
4196 let schema = Schema::try_from(&arrow).unwrap();
4197 let column_len = |c: usize| Ok(lengths[c]);
4198 let mut cursor = 0usize;
4199 validate_field_length(
4200 &schema.fields[0],
4201 is_structural,
4202 true,
4203 indices,
4204 &mut cursor,
4205 &column_len,
4206 )
4207 };
4208
4209 let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
4214 assert_eq!(
4215 validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(),
4216 5
4217 );
4218 assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5);
4219
4220 let list_of_struct = DataType::List(Arc::new(Field::new(
4227 "item",
4228 DataType::Struct(Fields::from(vec![
4229 Field::new("a", DataType::Int32, true),
4230 Field::new("b", DataType::Int32, true),
4231 ])),
4232 true,
4233 )));
4234 assert_eq!(
4235 validate(
4236 list_of_struct.clone(),
4237 false,
4238 &[0, 1, 2, 3],
4239 vec![6, 6, 29, 29]
4240 )
4241 .unwrap(),
4242 6
4243 );
4244 assert_eq!(
4246 validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(),
4247 29
4248 );
4249
4250 let empty_struct = DataType::Struct(Fields::empty());
4252 assert_eq!(
4253 validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(),
4254 9
4255 );
4256 assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9);
4257 }
4258}