feat[vortex-datafusion]: use PartitionedFile instead of path in create_reader (#7156)

The motivation for this change is that we have a need to create a custom
VortexReadAt implementation based on some custom values in a
PartitionedFile extension.

I don't think this will break any users because the reader module was
private.

<!--
Thank you for submitting a pull request! We appreciate your time and
effort.

Please make sure to provide enough information so that we can review
your pull
request. The Summary and Testing sections below contain guidance on what
to
include.
-->

<!--
## API Changes

Uncomment this section if there are any user-facing changes.

Consider whether the change affects users in one of the following ways:

1. Breaks public APIs in some way.
2. Changes the underlying behavior of one of the engine integrations.
3. Should some documentation be updated to reflect this change?

If a public API is changed in a breaking manner, make sure to add the
appropriate label. You can run `./scripts/public-api.sh` locally to see
if there
are any public API changes (and this also runs in our CI).
-->

## Testing

<!--
Please describe how this change was tested. Here are some common
categories for
testing in Vortex:

1. Verifying existing behavior is maintained.
2. Verifying new behavior and functionality works correctly.
3. Serialization compatibility (backwards and forwards) should be
maintained or
   explicitly broken.
-->

Signed-off-by: Alfonso Subiotto Marques <alfonso.subiotto@polarsignals.com>

Author: asubiotto

vortex-datafusion/public-api.lock

@@ -18,6 +18,30 @@ pub type vortex_datafusion::metrics::VortexMetricsFinder::Error = core::convert:
 
 pub fn vortex_datafusion::metrics::VortexMetricsFinder::pre_visit(&mut self, plan: &dyn datafusion_physical_plan::execution_plan::ExecutionPlan) -> core::result::Result<bool, Self::Error>
 
+pub mod vortex_datafusion::reader
+
+pub struct vortex_datafusion::reader::DefaultVortexReaderFactory
+
+impl vortex_datafusion::reader::DefaultVortexReaderFactory
+
+pub fn vortex_datafusion::reader::DefaultVortexReaderFactory::new(object_store: alloc::sync::Arc<dyn object_store::ObjectStore>) -> Self
+
+impl core::fmt::Debug for vortex_datafusion::reader::DefaultVortexReaderFactory
+
+pub fn vortex_datafusion::reader::DefaultVortexReaderFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_datafusion::reader::VortexReaderFactory for vortex_datafusion::reader::DefaultVortexReaderFactory
+
+pub fn vortex_datafusion::reader::DefaultVortexReaderFactory::create_reader(&self, file: &datafusion_datasource::PartitionedFile, session: &vortex_session::VortexSession) -> datafusion_common::error::Result<alloc::sync::Arc<dyn vortex_io::read_at::VortexReadAt>>
+
+pub trait vortex_datafusion::reader::VortexReaderFactory: core::fmt::Debug + core::marker::Send + core::marker::Sync + 'static
+
+pub fn vortex_datafusion::reader::VortexReaderFactory::create_reader(&self, file: &datafusion_datasource::PartitionedFile, session: &vortex_session::VortexSession) -> datafusion_common::error::Result<alloc::sync::Arc<dyn vortex_io::read_at::VortexReadAt>>
+
+impl vortex_datafusion::reader::VortexReaderFactory for vortex_datafusion::reader::DefaultVortexReaderFactory
+
+pub fn vortex_datafusion::reader::DefaultVortexReaderFactory::create_reader(&self, file: &datafusion_datasource::PartitionedFile, session: &vortex_session::VortexSession) -> datafusion_common::error::Result<alloc::sync::Arc<dyn vortex_io::read_at::VortexReadAt>>
+
 pub mod vortex_datafusion::v2
 
 pub struct vortex_datafusion::v2::VortexDataSource
@@ -176,7 +200,7 @@ pub fn vortex_datafusion::VortexSource::with_projection_pushdown(self, enabled:
 
 pub fn vortex_datafusion::VortexSource::with_scan_concurrency(self, scan_concurrency: usize) -> Self
 
-pub fn vortex_datafusion::VortexSource::with_vortex_reader_factory(self, vortex_reader_factory: alloc::sync::Arc<dyn vortex_datafusion::persistent::reader::VortexReaderFactory>) -> Self
+pub fn vortex_datafusion::VortexSource::with_vortex_reader_factory(self, vortex_reader_factory: alloc::sync::Arc<dyn vortex_datafusion::reader::VortexReaderFactory>) -> Self
 
 impl core::clone::Clone for vortex_datafusion::VortexSource
 

vortex-datafusion/src/persistent/mod.rs

@@ -7,7 +7,7 @@ mod cache;
 mod format;
 pub mod metrics;
 mod opener;
-mod reader;
+pub mod reader;
 mod sink;
 mod source;
 mod stream;

vortex-datafusion/src/persistent/opener.rs

@@ -110,9 +110,7 @@ impl FileOpener for VortexOpener {
         let mut projection = self.projection.clone();
         let mut filter = self.filter.clone();
 
-        let reader = self
-            .vortex_reader_factory
-            .create_reader(file.path().as_ref(), &session)?;
+        let reader = self.vortex_reader_factory.create_reader(&file, &session)?;
 
         let reader =
             InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone());

vortex-datafusion/src/persistent/reader.rs

@@ -1,10 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Factory for creating [`VortexReadAt`][vortex::io::VortexReadAt] instances
+//! from [`PartitionedFile`]s.
+
 use std::fmt::Debug;
 use std::sync::Arc;
 
 use datafusion_common::Result as DFResult;
+use datafusion_datasource::PartitionedFile;
 use object_store::ObjectStore;
 use vortex::io::VortexReadAt;
 use vortex::io::object_store::ObjectStoreReadAt;
@@ -14,8 +18,11 @@ use vortex::session::VortexSession;
 /// Factory to create [`VortexReadAt`] instances to read the target file.
 pub trait VortexReaderFactory: Debug + Send + Sync + 'static {
     /// Create a reader for a target object.
-    fn create_reader(&self, path: &str, session: &VortexSession)
-    -> DFResult<Arc<dyn VortexReadAt>>;
+    fn create_reader(
+        &self,
+        file: &PartitionedFile,
+        session: &VortexSession,
+    ) -> DFResult<Arc<dyn VortexReadAt>>;
 }
 
 /// Default factory, creates [`ObjectStore`] backed readers for files,
@@ -35,12 +42,12 @@ impl DefaultVortexReaderFactory {
 impl VortexReaderFactory for DefaultVortexReaderFactory {
     fn create_reader(
         &self,
-        path: &str,
+        file: &PartitionedFile,
         session: &VortexSession,
     ) -> DFResult<Arc<dyn VortexReadAt>> {
         Ok(Arc::new(ObjectStoreReadAt::new(
             self.object_store.clone(),
-            path.into(),
+            file.path().as_ref().into(),
             session.handle(),
         )) as _)
     }