lance_file/
testing.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_array::{RecordBatch, RecordBatchReader};
7use arrow_schema::ArrowError;
8use futures::TryStreamExt;
9use lance_core::{cache::LanceCache, datatypes::Schema, utils::tempfile::TempObjFile};
10use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
11use lance_io::{
12    ReadBatchParams,
13    object_store::ObjectStore,
14    scheduler::{ScanScheduler, SchedulerConfig},
15    utils::CachedFileSize,
16};
17
18use crate::reader::{FileReader, FileReaderOptions};
19use crate::version::ConcreteFileVersion;
20use crate::{versions, writer::FileWriterOptions};
21
22pub struct FsFixture {
23    pub tmp_path: TempObjFile,
24    pub object_store: Arc<ObjectStore>,
25    pub scheduler: Arc<ScanScheduler>,
26}
27
28impl Default for FsFixture {
29    fn default() -> Self {
30        let tmp_path = TempObjFile::default();
31        let object_store = Arc::new(ObjectStore::local());
32        let scheduler =
33            ScanScheduler::new(object_store.clone(), SchedulerConfig::default_for_testing());
34        Self {
35            object_store,
36            tmp_path,
37            scheduler,
38        }
39    }
40}
41
42pub struct WrittenFile {
43    pub schema: Arc<Schema>,
44    pub data: Vec<RecordBatch>,
45    pub field_id_mapping: Vec<(u32, u32)>,
46}
47
48pub async fn write_lance_file(
49    data: impl RecordBatchReader,
50    fs: &FsFixture,
51    version: ConcreteFileVersion,
52    options: FileWriterOptions,
53) -> WrittenFile {
54    let writer = fs.object_store.create(&fs.tmp_path).await.unwrap();
55
56    let lance_schema = lance_core::datatypes::Schema::try_from(data.schema().as_ref()).unwrap();
57
58    let mut file_writer =
59        versions::create_writer(version, writer, lance_schema.clone(), options).unwrap();
60
61    let data = data
62        .collect::<std::result::Result<Vec<_>, ArrowError>>()
63        .unwrap();
64
65    for batch in &data {
66        file_writer.write_batch(batch).await.unwrap();
67    }
68    let field_id_mapping = file_writer.field_id_to_column_indices().to_vec();
69    file_writer.add_schema_metadata("foo", "bar");
70    file_writer.finish().await.unwrap();
71    WrittenFile {
72        schema: Arc::new(lance_schema),
73        data,
74        field_id_mapping,
75    }
76}
77
78pub fn test_cache() -> Arc<LanceCache> {
79    Arc::new(LanceCache::with_capacity(128 * 1024 * 1024))
80}
81
82pub async fn read_lance_file(
83    fs: &FsFixture,
84    decoder_middleware: Arc<DecoderPlugins>,
85    filter: FilterExpression,
86) -> Vec<RecordBatch> {
87    let file_scheduler = fs
88        .scheduler
89        .open_file(&fs.tmp_path, &CachedFileSize::unknown())
90        .await
91        .unwrap();
92    let file_reader = FileReader::try_open(
93        file_scheduler,
94        None,
95        decoder_middleware,
96        &test_cache(),
97        FileReaderOptions::default(),
98    )
99    .await
100    .unwrap();
101
102    let schema = file_reader.schema();
103    assert_eq!(schema.metadata.get("foo").unwrap(), "bar");
104
105    let batch_stream = file_reader
106        .read_stream(ReadBatchParams::RangeFull, 1024, 16, filter)
107        .await
108        .unwrap();
109
110    batch_stream.try_collect().await.unwrap()
111}
112
113pub async fn count_lance_file(
114    fs: &FsFixture,
115    decoder_middleware: Arc<DecoderPlugins>,
116    filter: FilterExpression,
117) -> usize {
118    read_lance_file(fs, decoder_middleware, filter)
119        .await
120        .iter()
121        .map(|b| b.num_rows())
122        .sum()
123}
lance_file/testing.rs

lance_file/
testing.rs
