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, RecordBatchIterator, 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 async fn read_file_with_mutated_bytes(
2815 version: LanceFileVersion,
2816 batch: RecordBatch,
2817 pattern: &[u8],
2818 patch_offset: usize,
2819 patch: &[u8],
2820 ) -> lance_core::Result<Vec<RecordBatch>> {
2821 let fs = FsFixture::default();
2822 let schema = batch.schema();
2823 write_lance_file(
2824 RecordBatchIterator::new(vec![Ok(batch)], schema),
2825 &fs,
2826 FileWriterOptions {
2827 format_version: Some(version),
2828 ..Default::default()
2829 },
2830 )
2831 .await;
2832
2833 let mut bytes = fs
2834 .object_store
2835 .read_one_all(&fs.tmp_path)
2836 .await
2837 .unwrap()
2838 .to_vec();
2839 let matches = bytes
2840 .windows(pattern.len())
2841 .enumerate()
2842 .filter_map(|(position, window)| (window == pattern).then_some(position))
2843 .collect::<Vec<_>>();
2844 assert_eq!(
2845 matches.len(),
2846 1,
2847 "expected the byte pattern to appear exactly once in the file"
2848 );
2849 let patch_start = matches[0] + patch_offset;
2850 bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch);
2851 fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap();
2852
2853 let file_scheduler = fs
2854 .scheduler
2855 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2856 .await
2857 .unwrap();
2858 let file_reader = FileReader::try_open(
2859 file_scheduler,
2860 None,
2861 Arc::<DecoderPlugins>::default(),
2862 &test_cache(),
2863 FileReaderOptions::default(),
2864 )
2865 .await
2866 .unwrap();
2867 file_reader
2868 .read_stream(
2869 lance_io::ReadBatchParams::RangeFull,
2870 1024,
2871 16,
2872 FilterExpression::no_filter(),
2873 )
2874 .await?
2875 .try_collect::<Vec<_>>()
2876 .await
2877 }
2878
2879 #[rstest]
2889 #[tokio::test]
2890 async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets(
2891 #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
2892 version: LanceFileVersion,
2893 ) {
2894 use arrow_array::{Array, DictionaryArray, Int32Array, StringArray};
2895
2896 let values = StringArray::from(vec!["alpha", "beta", "gamma"]);
2897 let indices = Int32Array::from((0..300).map(|i| i % 3).collect::<Vec<i32>>());
2898 let dictionary = DictionaryArray::new(indices, Arc::new(values));
2899 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
2900 "category",
2901 dictionary.data_type().clone(),
2902 false,
2903 )]));
2904 let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap();
2905
2906 let offsets_tail_pattern = [5_i32, 9, 14]
2913 .iter()
2914 .flat_map(|value| value.to_le_bytes())
2915 .collect::<Vec<u8>>();
2916 let error = read_file_with_mutated_bytes(
2917 version,
2918 batch,
2919 &offsets_tail_pattern,
2920 8,
2921 &100_000_i32.to_le_bytes(),
2922 )
2923 .await
2924 .expect_err("out-of-bounds offsets must fail the read");
2925 let error_message = error.to_string();
2926 assert!(
2927 error_message.contains("corrupt file"),
2928 "expected a corruption error, got: {error_message}"
2929 );
2930 assert!(
2931 error_message.contains("out of bounds"),
2932 "unexpected message: {error_message}"
2933 );
2934 }
2935
2936 #[rstest]
2941 #[tokio::test]
2942 async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets(
2943 #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)]
2944 version: LanceFileVersion,
2945 ) {
2946 use arrow_array::StringArray;
2947
2948 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
2949 "strings",
2950 DataType::Utf8,
2951 false,
2952 )]));
2953 let batch = RecordBatch::try_new(
2954 arrow_schema,
2955 vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))],
2956 )
2957 .unwrap();
2958
2959 let chunk_offsets_pattern = [16_i32, 21, 25, 30]
2964 .iter()
2965 .flat_map(|value| value.to_le_bytes())
2966 .collect::<Vec<u8>>();
2967 let error = read_file_with_mutated_bytes(
2968 version,
2969 batch,
2970 &chunk_offsets_pattern,
2971 12,
2972 &100_000_i32.to_le_bytes(),
2973 )
2974 .await
2975 .expect_err("an out-of-bounds chunk offset must fail the read");
2976 let error_message = error.to_string();
2977 assert!(
2978 error_message.contains("corrupt file"),
2979 "expected a corruption error, got: {error_message}"
2980 );
2981 assert!(
2982 error_message.contains("out of bounds"),
2983 "unexpected message: {error_message}"
2984 );
2985 }
2986
2987 #[tokio::test]
2988 async fn test_round_trip() {
2989 let fs = FsFixture::default();
2990
2991 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
2992
2993 for read_size in [32, 1024, 1024 * 1024] {
2994 let file_scheduler = fs
2995 .scheduler
2996 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
2997 .await
2998 .unwrap();
2999 let file_reader = FileReader::try_open(
3000 file_scheduler,
3001 None,
3002 Arc::<DecoderPlugins>::default(),
3003 &test_cache(),
3004 FileReaderOptions::default(),
3005 )
3006 .await
3007 .unwrap();
3008
3009 let schema = file_reader.schema();
3010 assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
3011
3012 let batch_stream = file_reader
3013 .read_stream(
3014 lance_io::ReadBatchParams::RangeFull,
3015 read_size,
3016 16,
3017 FilterExpression::no_filter(),
3018 )
3019 .await
3020 .unwrap();
3021
3022 verify_expected(&data, batch_stream, read_size, None).await;
3023 }
3024 }
3025
3026 #[rstest]
3027 #[test_log::test(tokio::test)]
3028 async fn test_encoded_batch_round_trip(
3029 #[values(LanceFileVersion::V2_0)] version: LanceFileVersion,
3031 ) {
3032 let data = gen_batch()
3033 .col("x", array::rand::<Int32Type>())
3034 .col("y", array::rand_utf8(ByteCount::from(16), false))
3035 .into_batch_rows(RowCount::from(10000))
3036 .unwrap();
3037
3038 let lance_schema = Arc::new(Schema::try_from(data.schema().as_ref()).unwrap());
3039
3040 let encoding_options = EncodingOptions {
3041 cache_bytes_per_column: 4096,
3042 max_page_bytes: 32 * 1024 * 1024,
3043 keep_original_array: true,
3044 buffer_alignment: 64,
3045 version,
3046 };
3047
3048 let encoding_strategy = default_encoding_strategy(version);
3049
3050 let encoded_batch = encode_batch(
3051 &data,
3052 lance_schema.clone(),
3053 encoding_strategy.as_ref(),
3054 &encoding_options,
3055 )
3056 .await
3057 .unwrap();
3058
3059 let bytes = encoded_batch.try_to_self_described_lance(version).unwrap();
3061
3062 let decoded_batch = EncodedBatch::try_from_self_described_lance(bytes).unwrap();
3063
3064 let decoded = decode_batch(
3065 &decoded_batch,
3066 &FilterExpression::no_filter(),
3067 Arc::<DecoderPlugins>::default(),
3068 false,
3069 version,
3070 None,
3071 )
3072 .await
3073 .unwrap();
3074
3075 assert_eq!(data, decoded);
3076
3077 let bytes = encoded_batch.try_to_mini_lance(version).unwrap();
3079 let decoded_batch =
3080 EncodedBatch::try_from_mini_lance(bytes, lance_schema.as_ref(), LanceFileVersion::V2_0)
3081 .unwrap();
3082 let decoded = decode_batch(
3083 &decoded_batch,
3084 &FilterExpression::no_filter(),
3085 Arc::<DecoderPlugins>::default(),
3086 false,
3087 version,
3088 None,
3089 )
3090 .await
3091 .unwrap();
3092
3093 assert_eq!(data, decoded);
3094 }
3095
3096 #[rstest]
3097 #[test_log::test(tokio::test)]
3098 async fn test_projection(
3099 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
3100 version: LanceFileVersion,
3101 ) {
3102 let fs = FsFixture::default();
3103
3104 let written_file = create_some_file(&fs, version).await;
3105 let file_scheduler = fs
3106 .scheduler
3107 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3108 .await
3109 .unwrap();
3110
3111 let field_id_mapping = written_file
3112 .field_id_mapping
3113 .iter()
3114 .copied()
3115 .collect::<BTreeMap<_, _>>();
3116
3117 let empty_projection = ReaderProjection {
3118 column_indices: Vec::default(),
3119 schema: Arc::new(Schema::default()),
3120 };
3121
3122 for columns in [
3123 vec!["score"],
3124 vec!["location"],
3125 vec!["categories"],
3126 vec!["score.x"],
3127 vec!["score", "categories"],
3128 vec!["score", "location"],
3129 vec!["location", "categories"],
3130 vec!["score.y", "location", "categories"],
3131 ] {
3132 debug!("Testing round trip with projection {:?}", columns);
3133 for use_field_ids in [true, false] {
3134 let file_reader = FileReader::try_open(
3136 file_scheduler.clone(),
3137 None,
3138 Arc::<DecoderPlugins>::default(),
3139 &test_cache(),
3140 FileReaderOptions::default(),
3141 )
3142 .await
3143 .unwrap();
3144
3145 let projected_schema = written_file.schema.project(&columns).unwrap();
3146 let projection = if use_field_ids {
3147 ReaderProjection::from_field_ids(
3148 file_reader.metadata.version(),
3149 &projected_schema,
3150 &field_id_mapping,
3151 )
3152 .unwrap()
3153 } else {
3154 ReaderProjection::from_column_names(
3155 file_reader.metadata.version(),
3156 &written_file.schema,
3157 &columns,
3158 )
3159 .unwrap()
3160 };
3161
3162 let batch_stream = file_reader
3163 .read_stream_projected(
3164 lance_io::ReadBatchParams::RangeFull,
3165 1024,
3166 16,
3167 projection.clone(),
3168 FilterExpression::no_filter(),
3169 )
3170 .await
3171 .unwrap();
3172
3173 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3174 verify_expected(
3175 &written_file.data,
3176 batch_stream,
3177 1024,
3178 Some(Box::new(move |batch: &RecordBatch| {
3179 batch.project_by_schema(&projection_arrow).unwrap()
3180 })),
3181 )
3182 .await;
3183
3184 let file_reader = FileReader::try_open(
3186 file_scheduler.clone(),
3187 Some(projection.clone()),
3188 Arc::<DecoderPlugins>::default(),
3189 &test_cache(),
3190 FileReaderOptions::default(),
3191 )
3192 .await
3193 .unwrap();
3194
3195 let batch_stream = file_reader
3196 .read_stream(
3197 lance_io::ReadBatchParams::RangeFull,
3198 1024,
3199 16,
3200 FilterExpression::no_filter(),
3201 )
3202 .await
3203 .unwrap();
3204
3205 let projection_arrow = ArrowSchema::from(projection.schema.as_ref());
3206 verify_expected(
3207 &written_file.data,
3208 batch_stream,
3209 1024,
3210 Some(Box::new(move |batch: &RecordBatch| {
3211 batch.project_by_schema(&projection_arrow).unwrap()
3212 })),
3213 )
3214 .await;
3215
3216 assert!(
3217 file_reader
3218 .read_stream_projected(
3219 lance_io::ReadBatchParams::RangeFull,
3220 1024,
3221 16,
3222 empty_projection.clone(),
3223 FilterExpression::no_filter(),
3224 )
3225 .await
3226 .is_err()
3227 );
3228 }
3229 }
3230
3231 assert!(
3232 FileReader::try_open(
3233 file_scheduler.clone(),
3234 Some(empty_projection),
3235 Arc::<DecoderPlugins>::default(),
3236 &test_cache(),
3237 FileReaderOptions::default(),
3238 )
3239 .await
3240 .is_err()
3241 );
3242
3243 let arrow_schema = ArrowSchema::new(vec![
3244 Field::new("x", DataType::Int32, true),
3245 Field::new("y", DataType::Int32, true),
3246 ]);
3247 let schema = Schema::try_from(&arrow_schema).unwrap();
3248
3249 let projection_with_dupes = ReaderProjection {
3250 column_indices: vec![0, 0],
3251 schema: Arc::new(schema),
3252 };
3253
3254 assert!(
3255 FileReader::try_open(
3256 file_scheduler.clone(),
3257 Some(projection_with_dupes),
3258 Arc::<DecoderPlugins>::default(),
3259 &test_cache(),
3260 FileReaderOptions::default(),
3261 )
3262 .await
3263 .is_err()
3264 );
3265 }
3266
3267 #[tokio::test]
3268 async fn test_lazy_reader_direct_projection_matches_eager_reader() {
3269 let fs = FsFixture::default();
3270 let written_file = create_wide_direct_file(&fs, 16).await;
3271
3272 let file_scheduler = fs
3273 .scheduler
3274 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3275 .await
3276 .unwrap();
3277 let projection = ReaderProjection::from_column_names(
3278 LanceFileVersion::V2_1,
3279 &written_file.schema,
3280 &["c10"],
3281 )
3282 .unwrap();
3283
3284 let eager_reader = FileReader::try_open(
3285 file_scheduler.clone(),
3286 None,
3287 Arc::<DecoderPlugins>::default(),
3288 &test_cache(),
3289 FileReaderOptions::default(),
3290 )
3291 .await
3292 .unwrap();
3293 let expected = eager_reader
3294 .read_stream_projected(
3295 lance_io::ReadBatchParams::RangeFull,
3296 127,
3297 16,
3298 projection.clone(),
3299 FilterExpression::no_filter(),
3300 )
3301 .await
3302 .unwrap()
3303 .try_collect::<Vec<_>>()
3304 .await
3305 .unwrap();
3306
3307 let cache = test_cache();
3308 let lazy_reader = ProjectedFileReader::try_open(
3309 file_scheduler,
3310 Some(projection.clone()),
3311 Arc::<DecoderPlugins>::default(),
3312 &cache,
3313 FileReaderOptions::default(),
3314 )
3315 .await
3316 .unwrap();
3317 let tasks = lazy_reader
3318 .read_tasks(
3319 lance_io::ReadBatchParams::RangeFull,
3320 127,
3321 None,
3322 FilterExpression::no_filter(),
3323 )
3324 .await
3325 .unwrap();
3326 let actual = collect_read_tasks(tasks, 16).await;
3327
3328 assert_eq!(expected, actual);
3329 }
3330
3331 #[tokio::test]
3332 async fn test_lazy_reader_loads_only_requested_column_metadata() {
3333 let fs = FsFixture::default();
3334 let written_file = create_wide_direct_file(&fs, 512).await;
3335
3336 let projection = ReaderProjection::from_column_names(
3337 LanceFileVersion::V2_1,
3338 &written_file.schema,
3339 &["c0"],
3340 )
3341 .unwrap();
3342 let file_scheduler = fs
3343 .scheduler
3344 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3345 .await
3346 .unwrap();
3347 let lazy_reader = ProjectedFileReader::try_open(
3348 file_scheduler,
3349 Some(projection.clone()),
3350 Arc::<DecoderPlugins>::default(),
3351 &test_cache(),
3352 FileReaderOptions::default(),
3353 )
3354 .await
3355 .unwrap();
3356 let selected_column = projection.column_indices[0] as usize;
3357 let requested_metadata_bytes = lazy_reader
3358 .metadata_index()
3359 .unwrap()
3360 .column_metadata_offsets[selected_column]
3361 .1;
3362 let total_metadata_bytes = lazy_reader
3363 .metadata_index()
3364 .unwrap()
3365 .column_metadata_offsets
3366 .iter()
3367 .map(|(_, length)| *length)
3368 .sum::<u64>();
3369 assert!(
3370 total_metadata_bytes > 8 * fs.object_store.block_size() as u64,
3371 "test file metadata is too small to prove lazy loading: {total_metadata_bytes} bytes"
3372 );
3373
3374 fs.object_store.io_stats_incremental();
3375 let tasks = lazy_reader
3376 .read_tasks(
3377 lance_io::ReadBatchParams::Range(0..0),
3378 1024,
3379 Some(projection.clone()),
3380 FilterExpression::no_filter(),
3381 )
3382 .await
3383 .unwrap();
3384 let batches = collect_read_tasks(tasks, 1).await;
3385 assert!(batches.is_empty());
3386
3387 let stats = fs.object_store.io_stats_incremental();
3388 assert!(
3389 stats.read_bytes < total_metadata_bytes / 2,
3390 "lazy read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3391 stats.read_bytes,
3392 requested_metadata_bytes,
3393 total_metadata_bytes
3394 );
3395
3396 fs.object_store.io_stats_incremental();
3397 let tasks = lazy_reader
3398 .read_tasks(
3399 lance_io::ReadBatchParams::Range(0..0),
3400 1024,
3401 Some(projection),
3402 FilterExpression::no_filter(),
3403 )
3404 .await
3405 .unwrap();
3406 let batches = collect_read_tasks(tasks, 1).await;
3407 assert!(batches.is_empty());
3408
3409 let stats = fs.object_store.io_stats_incremental();
3410 assert_eq!(
3411 stats.read_iops, 0,
3412 "cached column metadata should avoid repeat metadata I/O"
3413 );
3414 assert_eq!(
3415 stats.read_bytes, 0,
3416 "cached column metadata should avoid repeat metadata reads"
3417 );
3418 }
3419
3420 async fn assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3421 fs: &FsFixture,
3422 projection: ReaderProjection,
3423 shape: &str,
3424 ) -> Vec<RecordBatch> {
3425 let file_scheduler = fs
3426 .scheduler
3427 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3428 .await
3429 .unwrap();
3430 let eager_reader = FileReader::try_open(
3431 file_scheduler.clone(),
3432 None,
3433 Arc::<DecoderPlugins>::default(),
3434 &test_cache(),
3435 FileReaderOptions::default(),
3436 )
3437 .await
3438 .unwrap();
3439 let expected = eager_reader
3440 .read_stream_projected(
3441 lance_io::ReadBatchParams::RangeFull,
3442 127,
3443 16,
3444 projection.clone(),
3445 FilterExpression::no_filter(),
3446 )
3447 .await
3448 .unwrap()
3449 .try_collect::<Vec<_>>()
3450 .await
3451 .unwrap();
3452
3453 let cache = test_cache();
3454 let lazy_reader = ProjectedFileReader::try_open(
3455 file_scheduler,
3456 Some(projection.clone()),
3457 Arc::<DecoderPlugins>::default(),
3458 &cache,
3459 FileReaderOptions::default(),
3460 )
3461 .await
3462 .unwrap();
3463 let metadata_index = lazy_reader.metadata_index().unwrap();
3464 let requested_metadata_bytes = projection
3465 .column_indices
3466 .iter()
3467 .map(|column_index| metadata_index.column_metadata_offsets[*column_index as usize].1)
3468 .sum::<u64>();
3469 let total_metadata_bytes = metadata_index
3470 .column_metadata_offsets
3471 .iter()
3472 .map(|(_, length)| *length)
3473 .sum::<u64>();
3474 assert!(total_metadata_bytes > requested_metadata_bytes * 8);
3475
3476 fs.object_store.io_stats_incremental();
3477 let tasks = lazy_reader
3478 .read_tasks(
3479 lance_io::ReadBatchParams::Range(0..0),
3480 127,
3481 None,
3482 FilterExpression::no_filter(),
3483 )
3484 .await
3485 .unwrap();
3486 assert!(collect_read_tasks(tasks, 1).await.is_empty());
3487 let metadata_stats = fs.object_store.io_stats_incremental();
3488 assert!(
3489 metadata_stats.read_bytes < total_metadata_bytes / 2,
3490 "lazy {shape} read fetched too much metadata: read {} bytes, requested column metadata is {} bytes, total column metadata is {} bytes",
3491 metadata_stats.read_bytes,
3492 requested_metadata_bytes,
3493 total_metadata_bytes
3494 );
3495
3496 let tasks = lazy_reader
3497 .read_tasks(
3498 lance_io::ReadBatchParams::RangeFull,
3499 127,
3500 None,
3501 FilterExpression::no_filter(),
3502 )
3503 .await
3504 .unwrap();
3505 let actual = collect_read_tasks(tasks, 16).await;
3506 assert_eq!(expected, actual);
3507 actual
3508 }
3509
3510 #[tokio::test]
3511 async fn test_lazy_reader_fixed_size_list_projection_matches_eager_reader() {
3512 let fs = FsFixture::default();
3513 let written_file = create_wide_fixed_size_list_file(&fs, 512).await;
3514 let projection = ReaderProjection::from_column_names(
3515 LanceFileVersion::V2_1,
3516 &written_file.schema,
3517 &["c17", "c509"],
3518 )
3519 .unwrap();
3520 assert!(ProjectedFileReader::supports_projection(
3521 &projection,
3522 LanceFileVersion::V2_1
3523 ));
3524 assert!(!ProjectedFileReader::supports_projection(
3525 &projection,
3526 LanceFileVersion::V2_0
3527 ));
3528 assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3529 &fs,
3530 projection,
3531 "fixed-size-list",
3532 )
3533 .await;
3534 }
3535
3536 #[tokio::test]
3537 async fn test_lazy_reader_nested_projection_compacts_physical_columns() {
3538 let fs = FsFixture::default();
3539 let written_file = create_wide_structural_file(&fs, 128).await;
3540 let projection = ReaderProjection::from_column_names(
3541 LanceFileVersion::V2_1,
3542 &written_file.schema,
3543 &["s97.y", "l4", "s3"],
3544 )
3545 .unwrap();
3546
3547 assert_eq!(
3548 projection
3549 .schema
3550 .fields
3551 .iter()
3552 .map(|field| field.name.as_str())
3553 .collect::<Vec<_>>(),
3554 vec!["s97", "l4", "s3"]
3555 );
3556 assert_eq!(projection.schema.fields[0].children.len(), 1);
3557 assert_eq!(projection.schema.fields[0].children[0].name, "y");
3558 assert_eq!(projection.schema.fields[2].children.len(), 2);
3559 assert_eq!(projection.column_indices.len(), 4);
3560 assert!(
3561 projection
3562 .column_indices
3563 .windows(2)
3564 .any(|indices| indices[0] > indices[1]),
3565 "the projection must reorder physical columns to exercise compact remapping"
3566 );
3567 assert!(ProjectedFileReader::supports_projection(
3568 &projection,
3569 LanceFileVersion::V2_1
3570 ));
3571 let actual = assert_lazy_projection_matches_eager_and_reads_metadata_subset(
3572 &fs, projection, "nested",
3573 )
3574 .await;
3575 assert!(
3576 actual
3577 .iter()
3578 .flat_map(|batch| batch.columns())
3579 .any(|column| column.null_count() > 0),
3580 "the structural projection must exercise nullable arrays"
3581 );
3582 }
3583
3584 #[rstest]
3585 #[case::before_metadata_region(90, 5)]
3586 #[case::after_metadata_region(190, 20)]
3587 fn test_decode_cmo_table_rejects_out_of_range_offsets(
3588 #[case] position: u64,
3589 #[case] length: u64,
3590 ) {
3591 let mut cmo_table = [0; 16];
3592 cmo_table[0..8].copy_from_slice(&position.to_le_bytes());
3593 cmo_table[8..16].copy_from_slice(&length.to_le_bytes());
3594 let footer = super::Footer {
3595 column_meta_start: 100,
3596 column_meta_offsets_start: 200,
3597 global_buff_offsets_start: 200,
3598 num_global_buffers: 0,
3599 num_columns: 1,
3600 major_version: 2,
3601 minor_version: 1,
3602 };
3603
3604 let err = FileReader::decode_cmo_table(Bytes::copy_from_slice(&cmo_table), &footer)
3605 .expect_err("out-of-range CMO entries must be rejected");
3606 assert!(
3607 matches!(err, lance_core::Error::InvalidInput { .. }),
3608 "expected InvalidInput, got {err:?}"
3609 );
3610 }
3611
3612 #[rstest]
3613 #[case::blob(BLOB_META_KEY)]
3614 #[case::packed_struct("lance-encoding:packed")]
3615 #[tokio::test]
3616 async fn test_lazy_reader_rejects_opaque_projection(#[case] metadata_key: &str) {
3617 let fs = FsFixture::default();
3618 let written_file = create_some_file(&fs, LanceFileVersion::V2_1).await;
3619
3620 let ordinary_projection = ReaderProjection::from_column_names(
3621 LanceFileVersion::V2_1,
3622 &written_file.schema,
3623 &["location.x"],
3624 )
3625 .unwrap();
3626 assert_eq!(ordinary_projection.schema.fields[0].children.len(), 1);
3627 assert!(ProjectedFileReader::supports_projection(
3628 &ordinary_projection,
3629 LanceFileVersion::V2_1
3630 ));
3631
3632 let file_scheduler = fs
3633 .scheduler
3634 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3635 .await
3636 .unwrap();
3637 let err = ProjectedFileReader::try_open(
3638 file_scheduler.clone(),
3639 None,
3640 Arc::<DecoderPlugins>::default(),
3641 &test_cache(),
3642 FileReaderOptions::default(),
3643 )
3644 .await
3645 .unwrap_err();
3646 assert!(
3647 matches!(err, lance_core::Error::InvalidInput { .. }),
3648 "expected InvalidInput, got {err:?}"
3649 );
3650
3651 let mut projection = ordinary_projection;
3652 Arc::make_mut(&mut projection.schema).fields[0]
3653 .metadata
3654 .insert(metadata_key.to_string(), "true".to_string());
3655 assert!(!ProjectedFileReader::supports_projection(
3656 &projection,
3657 LanceFileVersion::V2_1
3658 ));
3659
3660 let err = ProjectedFileReader::try_open(
3661 file_scheduler,
3662 Some(projection),
3663 Arc::<DecoderPlugins>::default(),
3664 &test_cache(),
3665 FileReaderOptions::default(),
3666 )
3667 .await
3668 .unwrap_err();
3669 assert!(
3670 matches!(err, lance_core::Error::NotSupported { .. }),
3671 "expected NotSupported for {metadata_key}, got {err:?}"
3672 );
3673 }
3674
3675 #[tokio::test]
3682 async fn test_lazy_reader_validates_unequal_length_projection() {
3683 use arrow_array::Int32Array;
3684 use lance_io::ReadBatchParams;
3685
3686 let arrow_schema = Arc::new(ArrowSchema::new(vec![
3687 Field::new("a", DataType::Int32, true),
3688 Field::new("c", DataType::Int32, true),
3689 ]));
3690 let lance_schema = Schema::try_from(arrow_schema.as_ref()).unwrap();
3691
3692 let fs = FsFixture::default();
3693 let options = FileWriterOptions {
3694 format_version: Some(LanceFileVersion::V2_1),
3695 ..Default::default()
3696 };
3697 let mut writer = FileWriter::try_new(
3698 fs.object_store.create(&fs.tmp_path).await.unwrap(),
3699 lance_schema.clone(),
3700 options,
3701 )
3702 .unwrap();
3703 writer
3705 .write_column(0, Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])))
3706 .await
3707 .unwrap();
3708 writer
3709 .write_column(1, Arc::new(Int32Array::from(vec![100])))
3710 .await
3711 .unwrap();
3712 writer.finish().await.unwrap();
3713
3714 let file_scheduler = fs
3715 .scheduler
3716 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3717 .await
3718 .unwrap();
3719 let cache = test_cache();
3720 let open_indexed = |names: &[&str]| {
3721 let projection =
3722 ReaderProjection::from_column_names(LanceFileVersion::V2_1, &lance_schema, names)
3723 .unwrap();
3724 ProjectedFileReader::try_open(
3725 file_scheduler.clone(),
3726 Some(projection),
3727 Arc::<DecoderPlugins>::default(),
3728 &cache,
3729 FileReaderOptions::default(),
3730 )
3731 };
3732
3733 let lazy = open_indexed(&["a", "c"]).await.unwrap();
3736 let err = match lazy
3739 .read_tasks(
3740 ReadBatchParams::RangeFull,
3741 1024,
3742 None,
3743 FilterExpression::no_filter(),
3744 )
3745 .await
3746 {
3747 Ok(_) => panic!("expected the mismatched-length projection to be rejected"),
3748 Err(e) => e.to_string(),
3749 };
3750 assert!(
3751 err.contains("a=5") && err.contains("c=1"),
3752 "error should name each column's length, got: {err}"
3753 );
3754
3755 let lazy = open_indexed(&["c"]).await.unwrap();
3758 let tasks = lazy
3759 .read_tasks(
3760 ReadBatchParams::RangeFull,
3761 1024,
3762 None,
3763 FilterExpression::no_filter(),
3764 )
3765 .await
3766 .unwrap();
3767 let batches = collect_read_tasks(tasks, 16).await;
3768 let values: Vec<Option<i32>> = batches
3769 .iter()
3770 .flat_map(|b| {
3771 b.column(0)
3772 .as_any()
3773 .downcast_ref::<Int32Array>()
3774 .unwrap()
3775 .iter()
3776 .collect::<Vec<_>>()
3777 })
3778 .collect();
3779 assert_eq!(values, vec![Some(100)]);
3780 }
3781
3782 #[test_log::test(tokio::test)]
3783 async fn test_compressing_buffer() {
3784 let fs = FsFixture::default();
3785
3786 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
3787 let file_scheduler = fs
3788 .scheduler
3789 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3790 .await
3791 .unwrap();
3792
3793 let file_reader = FileReader::try_open(
3795 file_scheduler.clone(),
3796 None,
3797 Arc::<DecoderPlugins>::default(),
3798 &test_cache(),
3799 FileReaderOptions::default(),
3800 )
3801 .await
3802 .unwrap();
3803
3804 let mut projection = written_file.schema.project(&["score"]).unwrap();
3805 for field in projection.fields.iter_mut() {
3806 field
3807 .metadata
3808 .insert("lance:compression".to_string(), "zstd".to_string());
3809 }
3810 let projection = ReaderProjection {
3811 column_indices: projection.fields.iter().map(|f| f.id as u32).collect(),
3812 schema: Arc::new(projection),
3813 };
3814
3815 let batch_stream = file_reader
3816 .read_stream_projected(
3817 lance_io::ReadBatchParams::RangeFull,
3818 1024,
3819 16,
3820 projection.clone(),
3821 FilterExpression::no_filter(),
3822 )
3823 .await
3824 .unwrap();
3825
3826 let projection_arrow = Arc::new(ArrowSchema::from(projection.schema.as_ref()));
3827 verify_expected(
3828 &written_file.data,
3829 batch_stream,
3830 1024,
3831 Some(Box::new(move |batch: &RecordBatch| {
3832 batch.project_by_schema(&projection_arrow).unwrap()
3833 })),
3834 )
3835 .await;
3836 }
3837
3838 #[tokio::test]
3839 async fn test_read_all() {
3840 let fs = FsFixture::default();
3841 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3842 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3843
3844 let file_scheduler = fs
3845 .scheduler
3846 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3847 .await
3848 .unwrap();
3849 let file_reader = FileReader::try_open(
3850 file_scheduler.clone(),
3851 None,
3852 Arc::<DecoderPlugins>::default(),
3853 &test_cache(),
3854 FileReaderOptions::default(),
3855 )
3856 .await
3857 .unwrap();
3858
3859 let batches = file_reader
3860 .read_stream(
3861 lance_io::ReadBatchParams::RangeFull,
3862 total_rows as u32,
3863 16,
3864 FilterExpression::no_filter(),
3865 )
3866 .await
3867 .unwrap()
3868 .try_collect::<Vec<_>>()
3869 .await
3870 .unwrap();
3871 assert_eq!(batches.len(), 1);
3872 assert_eq!(batches[0].num_rows(), total_rows);
3873 }
3874
3875 #[rstest]
3876 #[tokio::test]
3877 async fn test_blocking_take(
3878 #[values(LanceFileVersion::V2_0, LanceFileVersion::V2_1, LanceFileVersion::V2_2)]
3879 version: LanceFileVersion,
3880 ) {
3881 let fs = FsFixture::default();
3882 let WrittenFile { data, schema, .. } = create_some_file(&fs, version).await;
3883 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3884
3885 let file_scheduler = fs
3886 .scheduler
3887 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3888 .await
3889 .unwrap();
3890 let file_reader = FileReader::try_open(
3891 file_scheduler.clone(),
3892 Some(ReaderProjection::from_column_names(version, &schema, &["score"]).unwrap()),
3893 Arc::<DecoderPlugins>::default(),
3894 &test_cache(),
3895 FileReaderOptions::default(),
3896 )
3897 .await
3898 .unwrap();
3899
3900 let batches = tokio::task::spawn_blocking(move || {
3901 file_reader
3902 .read_stream_projected_blocking(
3903 lance_io::ReadBatchParams::Indices(UInt32Array::from(vec![0, 1, 2, 3, 4])),
3904 total_rows as u32,
3905 None,
3906 FilterExpression::no_filter(),
3907 )
3908 .unwrap()
3909 .collect::<ArrowResult<Vec<_>>>()
3910 .unwrap()
3911 })
3912 .await
3913 .unwrap();
3914
3915 assert_eq!(batches.len(), 1);
3916 assert_eq!(batches[0].num_rows(), 5);
3917 assert_eq!(batches[0].num_columns(), 1);
3918 }
3919
3920 #[tokio::test(flavor = "multi_thread")]
3921 async fn test_drop_in_progress() {
3922 let fs = FsFixture::default();
3923 let WrittenFile { data, .. } = create_some_file(&fs, LanceFileVersion::V2_0).await;
3924 let total_rows = data.iter().map(|batch| batch.num_rows()).sum::<usize>();
3925
3926 let file_scheduler = fs
3927 .scheduler
3928 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3929 .await
3930 .unwrap();
3931 let file_reader = FileReader::try_open(
3932 file_scheduler.clone(),
3933 None,
3934 Arc::<DecoderPlugins>::default(),
3935 &test_cache(),
3936 FileReaderOptions::default(),
3937 )
3938 .await
3939 .unwrap();
3940
3941 let mut batches = file_reader
3942 .read_stream(
3943 lance_io::ReadBatchParams::RangeFull,
3944 (total_rows / 10) as u32,
3945 16,
3946 FilterExpression::no_filter(),
3947 )
3948 .await
3949 .unwrap();
3950
3951 drop(file_reader);
3952
3953 let batch = batches.next().await.unwrap().unwrap();
3954 assert!(batch.num_rows() > 0);
3955
3956 drop(batches);
3958 }
3959
3960 #[tokio::test]
3961 async fn drop_while_scheduling() {
3962 let fs = FsFixture::default();
3972 let written_file = create_some_file(&fs, LanceFileVersion::V2_0).await;
3973 let total_rows = written_file
3974 .data
3975 .iter()
3976 .map(|batch| batch.num_rows())
3977 .sum::<usize>();
3978
3979 let file_scheduler = fs
3980 .scheduler
3981 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
3982 .await
3983 .unwrap();
3984 let file_reader = FileReader::try_open(
3985 file_scheduler.clone(),
3986 None,
3987 Arc::<DecoderPlugins>::default(),
3988 &test_cache(),
3989 FileReaderOptions::default(),
3990 )
3991 .await
3992 .unwrap();
3993
3994 let projection =
3995 ReaderProjection::from_whole_schema(&written_file.schema, LanceFileVersion::V2_0);
3996 let column_infos = file_reader
3997 .collect_columns_from_projection(&projection)
3998 .unwrap();
3999 let mut decode_scheduler = DecodeBatchScheduler::try_new(
4000 &projection.schema,
4001 &projection.column_indices,
4002 &column_infos,
4003 &vec![],
4004 total_rows as u64,
4005 Arc::<DecoderPlugins>::default(),
4006 file_reader.core.scheduler.clone(),
4007 test_cache(),
4008 &FilterExpression::no_filter(),
4009 &DecoderConfig::default(),
4010 )
4011 .await
4012 .unwrap();
4013
4014 let range = 0..total_rows as u64;
4015
4016 let (tx, rx) = mpsc::unbounded_channel();
4017
4018 drop(rx);
4020
4021 decode_scheduler.schedule_range(
4023 range,
4024 &FilterExpression::no_filter(),
4025 tx,
4026 file_reader.core.scheduler.clone(),
4027 )
4028 }
4029
4030 #[tokio::test]
4031 async fn test_read_empty_range() {
4032 let fs = FsFixture::default();
4033 create_some_file(&fs, LanceFileVersion::V2_0).await;
4034
4035 let file_scheduler = fs
4036 .scheduler
4037 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4038 .await
4039 .unwrap();
4040 let file_reader = FileReader::try_open(
4041 file_scheduler.clone(),
4042 None,
4043 Arc::<DecoderPlugins>::default(),
4044 &test_cache(),
4045 FileReaderOptions::default(),
4046 )
4047 .await
4048 .unwrap();
4049
4050 let batches = file_reader
4052 .read_stream(
4053 lance_io::ReadBatchParams::Range(0..0),
4054 1024,
4055 16,
4056 FilterExpression::no_filter(),
4057 )
4058 .await
4059 .unwrap()
4060 .try_collect::<Vec<_>>()
4061 .await
4062 .unwrap();
4063
4064 assert_eq!(batches.len(), 0);
4065
4066 let batches = file_reader
4068 .read_stream(
4069 lance_io::ReadBatchParams::Ranges(Arc::new([0..1, 2..2])),
4070 1024,
4071 16,
4072 FilterExpression::no_filter(),
4073 )
4074 .await
4075 .unwrap()
4076 .try_collect::<Vec<_>>()
4077 .await
4078 .unwrap();
4079 assert_eq!(batches.len(), 1);
4080 }
4081
4082 async fn write_file_with_global_buffer(fs: &FsFixture, buffer: Bytes) {
4083 let lance_schema =
4084 lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![Field::new(
4085 "foo",
4086 DataType::Int32,
4087 true,
4088 )]))
4089 .unwrap();
4090
4091 let mut file_writer = FileWriter::try_new(
4092 fs.object_store.create(&fs.tmp_path).await.unwrap(),
4093 lance_schema,
4094 FileWriterOptions::default(),
4095 )
4096 .unwrap();
4097
4098 let buf_index = file_writer.add_global_buffer(buffer).await.unwrap();
4099 assert_eq!(buf_index, 1);
4100
4101 file_writer.finish().await.unwrap();
4102 }
4103
4104 #[rstest]
4108 #[case::within_tail_window(true)]
4109 #[case::outside_tail_window(false)]
4110 #[tokio::test]
4111 async fn test_read_global_buffer(#[case] within_window: bool) {
4112 let fs = FsFixture::default();
4113
4114 let block_size = fs.object_store.block_size();
4115 let buffer = if within_window {
4116 Bytes::from_static(b"hello")
4117 } else {
4118 Bytes::from(vec![7u8; 2 * block_size])
4119 };
4120 let expected_read_iops = if within_window { 0 } else { 1 };
4121
4122 write_file_with_global_buffer(&fs, buffer.clone()).await;
4123
4124 let file_scheduler = fs
4125 .scheduler
4126 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4127 .await
4128 .unwrap();
4129 let file_reader = FileReader::try_open(
4130 file_scheduler,
4131 None,
4132 Arc::<DecoderPlugins>::default(),
4133 &test_cache(),
4134 FileReaderOptions::default(),
4135 )
4136 .await
4137 .unwrap();
4138
4139 let retained = &file_reader.metadata().retained_global_buffers;
4142 assert!(!retained.contains_key(&0), "schema must not be retained");
4143 assert_eq!(retained.contains_key(&1), within_window);
4144
4145 fs.object_store.io_stats_incremental();
4147
4148 let buf = file_reader.read_global_buffer(1).await.unwrap();
4149 assert_eq!(buf, buffer);
4150
4151 let stats = fs.object_store.io_stats_incremental();
4152 assert_eq!(stats.read_iops, expected_read_iops);
4153 }
4154
4155 #[tokio::test]
4158 async fn test_read_global_buffer_no_user_buffers() {
4159 let fs = FsFixture::default();
4160 create_some_file(&fs, LanceFileVersion::V2_1).await;
4161
4162 let file_scheduler = fs
4163 .scheduler
4164 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4165 .await
4166 .unwrap();
4167 let file_reader = FileReader::try_open(
4168 file_scheduler,
4169 None,
4170 Arc::<DecoderPlugins>::default(),
4171 &test_cache(),
4172 FileReaderOptions::default(),
4173 )
4174 .await
4175 .unwrap();
4176
4177 let metadata = file_reader.metadata();
4178 assert_eq!(metadata.file_buffers.len(), 1, "expected only the schema");
4179 assert!(
4180 metadata.retained_global_buffers.is_empty(),
4181 "a file with no user global buffers must retain nothing"
4182 );
4183 }
4184
4185 #[rstest]
4186 #[tokio::test]
4187 async fn test_deep_size_of_includes_column_metadata(
4188 #[values(
4189 LanceFileVersion::V2_0,
4190 LanceFileVersion::V2_1,
4191 LanceFileVersion::V2_2,
4192 LanceFileVersion::V2_3
4193 )]
4194 version: LanceFileVersion,
4195 ) {
4196 use lance_core::deepsize::DeepSizeOf;
4201
4202 let fs = FsFixture::default();
4203 let _written = create_some_file(&fs, version).await;
4204 let cache = test_cache();
4205 let file_scheduler = fs
4206 .scheduler
4207 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4208 .await
4209 .unwrap();
4210 let file_reader = FileReader::try_open(
4211 file_scheduler,
4212 None,
4213 Arc::<DecoderPlugins>::default(),
4214 &cache,
4215 FileReaderOptions::default(),
4216 )
4217 .await
4218 .unwrap();
4219
4220 let metadata = file_reader.metadata();
4221 let deep_size = metadata.deep_size_of();
4222
4223 assert!(
4229 deep_size > 1024,
4230 "deep_size_of ({deep_size}) is suspiciously small — \
4231 column_metadatas and column_infos may not be accounted for"
4232 );
4233
4234 assert!(
4236 !metadata.column_metadatas.is_empty(),
4237 "Expected non-empty column_metadatas"
4238 );
4239
4240 let num_columns = metadata.column_metadatas.len();
4243 assert!(
4244 deep_size > num_columns * 50,
4245 "deep_size_of ({deep_size}) should scale with column count ({num_columns})"
4246 );
4247 }
4248
4249 #[tokio::test]
4250 async fn test_read_global_buffer_out_of_range() {
4251 let fs = FsFixture::default();
4252
4253 write_file_with_global_buffer(&fs, Bytes::from_static(b"hello")).await;
4254
4255 let file_scheduler = fs
4256 .scheduler
4257 .open_file(&fs.tmp_path, &CachedFileSize::unknown())
4258 .await
4259 .unwrap();
4260 let file_reader = FileReader::try_open(
4261 file_scheduler,
4262 None,
4263 Arc::<DecoderPlugins>::default(),
4264 &test_cache(),
4265 FileReaderOptions::default(),
4266 )
4267 .await
4268 .unwrap();
4269
4270 let err = file_reader.read_global_buffer(2).await.unwrap_err();
4273 assert!(
4274 matches!(err, lance_core::Error::InvalidInput { .. }),
4275 "expected InvalidInput, got: {err:?}"
4276 );
4277 let msg = err.to_string();
4278 assert!(msg.contains('2'), "error should mention the index: {msg}");
4279 }
4280
4281 #[rstest]
4286 fn test_validate_struct_child_lengths(#[values(false, true)] is_structural: bool) {
4287 let run = |dt: DataType, indices: &[u32], lengths: Vec<u64>| -> lance_core::Result<u64> {
4288 let arrow = ArrowSchema::new(vec![Field::new("s", dt, true)]);
4289 let schema = Schema::try_from(&arrow).unwrap();
4290 let column_len = |c: usize| Ok(lengths[c]);
4291 let mut cursor = 0usize;
4292 let mut field_lengths = Vec::new();
4293 for field in &schema.fields {
4294 let rows = validate_field_length(
4295 field,
4296 is_structural,
4297 true,
4298 indices,
4299 &mut cursor,
4300 &column_len,
4301 )?;
4302 field_lengths.push((field.name.as_str(), rows));
4303 }
4304 verify_uniform_lengths(&field_lengths)
4305 };
4306
4307 let struct_ty = || {
4308 DataType::Struct(Fields::from(vec![
4309 Field::new("a", DataType::Int32, true),
4310 Field::new("b", DataType::Int32, true),
4311 ]))
4312 };
4313
4314 let (indices, equal, unequal): (&[u32], Vec<u64>, Vec<u64>) = if is_structural {
4317 (&[0, 1], vec![5, 5], vec![5, 3])
4318 } else {
4319 (&[0, 1, 2], vec![5, 5, 5], vec![5, 5, 3])
4320 };
4321
4322 assert_eq!(run(struct_ty(), indices, equal).unwrap(), 5);
4323
4324 let err = run(struct_ty(), indices, unequal).unwrap_err();
4325 let msg = err.to_string();
4326 assert!(
4327 msg.contains("differing lengths") && msg.contains('b'),
4328 "expected a child-length error naming 'b', got: {msg}"
4329 );
4330 }
4331
4332 #[test]
4333 fn test_validate_v2_0_unloaded_blob_projection_is_opaque() {
4334 let metadata = HashMap::from([(BLOB_META_KEY.to_string(), "true".to_string())]);
4335 let arrow = ArrowSchema::new(vec![
4336 Field::new("blob", DataType::LargeBinary, true).with_metadata(metadata),
4337 ]);
4338 let mut schema = Schema::try_from(&arrow).unwrap();
4339 schema.fields[0].unloaded_mut();
4340 let projection = ReaderProjection {
4341 schema: Arc::new(schema),
4342 column_indices: vec![0],
4343 };
4344 let column_len = |column: usize| {
4345 assert_eq!(column, 0);
4346 Ok(3)
4347 };
4348 let mut cursor = 0usize;
4349
4350 let rows = validate_field_length(
4351 &projection.schema.fields[0],
4352 false,
4353 true,
4354 &projection.column_indices,
4355 &mut cursor,
4356 &column_len,
4357 )
4358 .unwrap();
4359
4360 assert_eq!(rows, 3);
4361 assert_eq!(cursor, 1);
4362 }
4363
4364 #[test]
4365 fn test_validate_length_list_and_empty_struct() {
4366 let validate = |dt: DataType,
4367 is_structural: bool,
4368 indices: &[u32],
4369 lengths: Vec<u64>|
4370 -> lance_core::Result<u64> {
4371 let arrow = ArrowSchema::new(vec![Field::new("f", dt, true)]);
4372 let schema = Schema::try_from(&arrow).unwrap();
4373 let column_len = |c: usize| Ok(lengths[c]);
4374 let mut cursor = 0usize;
4375 validate_field_length(
4376 &schema.fields[0],
4377 is_structural,
4378 true,
4379 indices,
4380 &mut cursor,
4381 &column_len,
4382 )
4383 };
4384
4385 let list_ty = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
4390 assert_eq!(
4391 validate(list_ty.clone(), false, &[0, 1], vec![5, 17]).unwrap(),
4392 5
4393 );
4394 assert_eq!(validate(list_ty, true, &[0], vec![5]).unwrap(), 5);
4395
4396 let list_of_struct = DataType::List(Arc::new(Field::new(
4403 "item",
4404 DataType::Struct(Fields::from(vec![
4405 Field::new("a", DataType::Int32, true),
4406 Field::new("b", DataType::Int32, true),
4407 ])),
4408 true,
4409 )));
4410 assert_eq!(
4411 validate(
4412 list_of_struct.clone(),
4413 false,
4414 &[0, 1, 2, 3],
4415 vec![6, 6, 29, 29]
4416 )
4417 .unwrap(),
4418 6
4419 );
4420 assert_eq!(
4422 validate(list_of_struct, true, &[0, 1], vec![29, 29]).unwrap(),
4423 29
4424 );
4425
4426 let empty_struct = DataType::Struct(Fields::empty());
4428 assert_eq!(
4429 validate(empty_struct.clone(), false, &[0], vec![9]).unwrap(),
4430 9
4431 );
4432 assert_eq!(validate(empty_struct, true, &[0], vec![9]).unwrap(), 9);
4433 }
4434}