ObjectStoreReadAt to optionally accept scheme for more accurate perf hints (#2343)

Author: onursatici

bench-vortex/benches/random_access.rs

@@ -42,7 +42,9 @@ fn random_access_vortex(c: &mut Criterion) {
         b.to_async(Runtime::new().unwrap()).iter_with_setup(
             || (local_fs.clone(), local_fs_path.clone(), indices.clone()),
             |(fs, path, indices)| async {
-                take_vortex_object_store(fs, path, indices).await.unwrap()
+                take_vortex_object_store(object_store::ObjectStoreScheme::Local, fs, path, indices)
+                    .await
+                    .unwrap()
             },
         )
     });
@@ -70,7 +72,14 @@ fn random_access_vortex(c: &mut Criterion) {
             b.to_async(Runtime::new().unwrap()).iter_with_setup(
                 || (r2_fs.clone(), r2_path.clone(), indices.clone()),
                 |(fs, path, indices)| async {
-                    take_vortex_object_store(fs, path, indices).await.unwrap()
+                    take_vortex_object_store(
+                        object_store::ObjectStoreScheme::AmazonS3,
+                        fs,
+                        path,
+                        indices,
+                    )
+                    .await
+                    .unwrap()
                 },
             )
         });

bench-vortex/src/reader.rs

@@ -14,7 +14,7 @@ use arrow_select::take::take_record_batch;
 use futures::stream;
 use itertools::Itertools;
 use log::info;
-use object_store::ObjectStore;
+use object_store::{ObjectStore, ObjectStoreScheme};
 use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder};
 use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
 use parquet::arrow::ParquetRecordBatchStreamBuilder;
@@ -112,11 +112,16 @@ async fn take_vortex<T: VortexReadAt + Unpin + 'static>(
 }
 
 pub async fn take_vortex_object_store(
+    scheme: ObjectStoreScheme,
     fs: Arc<dyn ObjectStore>,
     path: object_store::path::Path,
     indices: Buffer<u64>,
 ) -> VortexResult<Array> {
-    take_vortex(ObjectStoreReadAt::new(fs.clone(), path), indices).await
+    take_vortex(
+        ObjectStoreReadAt::new(fs.clone(), path, Some(scheme)),
+        indices,
+    )
+    .await
 }
 
 pub async fn take_vortex_tokio(path: &Path, indices: Buffer<u64>) -> VortexResult<Array> {

pyvortex/src/object_store_urls.rs

@@ -11,7 +11,9 @@ use url::Url;
 use vortex::error::{vortex_bail, VortexResult};
 use vortex::io::ObjectStoreReadAt;
 
-fn better_parse_url(url_str: &str) -> VortexResult<(Box<dyn ObjectStore>, Path)> {
+fn better_parse_url(
+    url_str: &str,
+) -> VortexResult<(ObjectStoreScheme, Box<dyn ObjectStore>, Path)> {
     let url = Url::parse(url_str)?;
 
     let (scheme, path) = ObjectStoreScheme::parse(&url).map_err(object_store::Error::from)?;
@@ -38,10 +40,14 @@ fn better_parse_url(url_str: &str) -> VortexResult<(Box<dyn ObjectStore>, Path)>
         otherwise => vortex_bail!("unrecognized object store scheme: {:?}", otherwise),
     };
 
-    Ok((store, path))
+    Ok((scheme, store, path))
 }
 
 pub async fn vortex_read_at_from_url(url: &str) -> VortexResult<ObjectStoreReadAt> {
-    let (object_store, location) = better_parse_url(url)?;
-    Ok(ObjectStoreReadAt::new(Arc::from(object_store), location))
+    let (scheme, object_store, location) = better_parse_url(url)?;
+    Ok(ObjectStoreReadAt::new(
+        Arc::from(object_store),
+        location,
+        Some(scheme),
+    ))
 }

vortex-datafusion/src/persistent/cache.rs

@@ -80,7 +80,8 @@ impl FileLayoutCache {
     ) -> VortexResult<FileLayout> {
         self.inner
             .try_get_with(Key::from(object), async {
-                let os_read_at = ObjectStoreReadAt::new(object_store, object.location.clone());
+                let os_read_at =
+                    ObjectStoreReadAt::new(object_store, object.location.clone(), None);
                 let vxf = VortexOpenOptions::file(os_read_at)
                     .with_ctx(self.context.clone())
                     .with_file_size(object.size as u64)

vortex-datafusion/src/persistent/execution.rs

@@ -12,6 +12,7 @@ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use itertools::Itertools;
+use object_store::ObjectStoreScheme;
 use vortex_array::ContextRef;
 use vortex_expr::datafusion::convert_expr_to_vortex;
 use vortex_expr::{and, VortexExpr};
@@ -124,9 +125,12 @@ impl ExecutionPlan for VortexExec {
         let object_store = context
             .runtime_env()
             .object_store(&self.file_scan_config.object_store_url)?;
+        let (scheme, _) = ObjectStoreScheme::parse(self.file_scan_config.object_store_url.as_ref())
+            .map_err(object_store::Error::from)?;
 
         let opener = VortexFileOpener::new(
             self.ctx.clone(),
+            scheme,
             object_store,
             self.projection.clone(),
             self.predicate.clone(),

vortex-datafusion/src/persistent/format.rs

@@ -192,7 +192,7 @@ impl FileFormat for VortexFormat {
         table_schema: SchemaRef,
         object: &ObjectMeta,
     ) -> DFResult<Statistics> {
-        let read_at = ObjectStoreReadAt::new(store.clone(), object.location.clone());
+        let read_at = ObjectStoreReadAt::new(store.clone(), object.location.clone(), None);
         let file_layout = self
             .file_layout_cache
             .try_get(object, store.clone())

vortex-datafusion/src/persistent/opener.rs

@@ -4,7 +4,7 @@ use arrow_schema::SchemaRef;
 use datafusion::datasource::physical_plan::{FileMeta, FileOpenFuture, FileOpener};
 use datafusion_common::Result as DFResult;
 use futures::{FutureExt as _, StreamExt};
-use object_store::ObjectStore;
+use object_store::{ObjectStore, ObjectStoreScheme};
 use vortex_array::{ContextRef, IntoArrayVariant};
 use vortex_error::VortexResult;
 use vortex_expr::{ExprRef, VortexExpr};
@@ -16,6 +16,7 @@ use super::cache::FileLayoutCache;
 #[derive(Clone)]
 pub(crate) struct VortexFileOpener {
     pub ctx: ContextRef,
+    pub scheme: ObjectStoreScheme,
     pub object_store: Arc<dyn ObjectStore>,
     pub projection: ExprRef,
     pub filter: Option<ExprRef>,
@@ -25,8 +26,10 @@ pub(crate) struct VortexFileOpener {
 }
 
 impl VortexFileOpener {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         ctx: ContextRef,
+        scheme: ObjectStoreScheme,
         object_store: Arc<dyn ObjectStore>,
         projection: Arc<dyn VortexExpr>,
         filter: Option<Arc<dyn VortexExpr>>,
@@ -36,6 +39,7 @@ impl VortexFileOpener {
     ) -> VortexResult<Self> {
         Ok(Self {
             ctx,
+            scheme,
             object_store,
             projection,
             filter,
@@ -48,8 +52,11 @@ impl VortexFileOpener {
 
 impl FileOpener for VortexFileOpener {
     fn open(&self, file_meta: FileMeta) -> DFResult<FileOpenFuture> {
-        let read_at =
-            ObjectStoreReadAt::new(self.object_store.clone(), file_meta.location().clone());
+        let read_at = ObjectStoreReadAt::new(
+            self.object_store.clone(),
+            file_meta.location().clone(),
+            Some(self.scheme.clone()),
+        );
 
         let filter = self.filter.clone();
         let projection = self.projection.clone();

vortex-io/src/object_store.rs

@@ -7,24 +7,31 @@ use bytes::Bytes;
 use futures_util::StreamExt;
 use object_store::path::Path;
 use object_store::{
-    GetOptions, GetRange, GetResultPayload, MultipartUpload, ObjectStore, PutPayload,
+    GetOptions, GetRange, GetResultPayload, MultipartUpload, ObjectStore, ObjectStoreScheme,
+    PutPayload,
 };
 use vortex_buffer::{Alignment, ByteBuffer, ByteBufferMut};
 use vortex_error::{VortexExpect, VortexResult};
 
-use crate::{IoBuf, VortexReadAt, VortexWrite};
+use crate::{IoBuf, PerformanceHint, VortexReadAt, VortexWrite};
 
 #[derive(Clone)]
 pub struct ObjectStoreReadAt {
     object_store: Arc<dyn ObjectStore>,
     location: Path,
+    scheme: Option<ObjectStoreScheme>,
 }
 
 impl ObjectStoreReadAt {
-    pub fn new(object_store: Arc<dyn ObjectStore>, location: Path) -> Self {
+    pub fn new(
+        object_store: Arc<dyn ObjectStore>,
+        location: Path,
+        scheme: Option<ObjectStoreScheme>,
+    ) -> Self {
         Self {
             object_store,
             location,
+            scheme,
         }
     }
 }
@@ -84,6 +91,13 @@ impl VortexReadAt for ObjectStoreReadAt {
         let location = self.location.clone();
         Ok(object_store.head(&location).await?.size as u64)
     }
+
+    fn performance_hint(&self) -> PerformanceHint {
+        match &self.scheme {
+            Some(ObjectStoreScheme::Local | ObjectStoreScheme::Memory) => PerformanceHint::local(),
+            _ => PerformanceHint::default(),
+        }
+    }
 }
 
 pub struct ObjectStoreWriter {