Skip to main content

vortex_datafusion/persistent/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Factory for creating [`VortexReadAt`][vortex::io::VortexReadAt] instances
5//! from [`PartitionedFile`]s.
6
7use std::fmt::Debug;
8use std::sync::Arc;
9
10use datafusion_common::Result as DFResult;
11use datafusion_datasource::PartitionedFile;
12use object_store::ObjectStore;
13use vortex::array::memory::MemorySessionExt;
14use vortex::io::VortexReadAt;
15use vortex::io::object_store::ObjectStoreReadAt;
16use vortex::io::session::RuntimeSessionExt;
17use vortex::session::VortexSession;
18
19/// Factory to create [`VortexReadAt`] instances to read the target file.
20pub trait VortexReaderFactory: Debug + Send + Sync + 'static {
21    /// Create a reader for a target object.
22    fn create_reader(
23        &self,
24        file: &PartitionedFile,
25        session: &VortexSession,
26    ) -> DFResult<Arc<dyn VortexReadAt>>;
27}
28
29/// Default factory, creates [`ObjectStore`] backed readers for files,
30/// works with multiple cloud providers.
31#[derive(Debug)]
32pub struct DefaultVortexReaderFactory {
33    object_store: Arc<dyn ObjectStore>,
34}
35
36impl DefaultVortexReaderFactory {
37    /// Creates new instance
38    pub fn new(object_store: Arc<dyn ObjectStore>) -> Self {
39        Self { object_store }
40    }
41}
42
43impl VortexReaderFactory for DefaultVortexReaderFactory {
44    fn create_reader(
45        &self,
46        file: &PartitionedFile,
47        session: &VortexSession,
48    ) -> DFResult<Arc<dyn VortexReadAt>> {
49        Ok(Arc::new(ObjectStoreReadAt::new_with_allocator(
50            Arc::clone(&self.object_store),
51            file.path().as_ref().into(),
52            session.handle(),
53            session.allocator(),
54        )) as _)
55    }
56}