Extract Listing URI logic into ListingTableUri structure

Author: tustvold

ballista/rust/core/src/serde/logical_plan/mod.rs

@@ -27,7 +27,9 @@ use datafusion::datasource::file_format::avro::AvroFormat;
 use datafusion::datasource::file_format::csv::CsvFormat;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion::datasource::file_format::FileFormat;
-use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig};
+use datafusion::datasource::listing::{
+    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
 use datafusion::logical_plan::plan::{
     Aggregate, EmptyRelation, Filter, Join, Projection, Sort, SubqueryAlias, Window,
 };
@@ -227,24 +229,19 @@ impl AsLogicalPlan for LogicalPlanNode {
                     target_partitions: scan.target_partitions as usize,
                 };
 
-                let object_store = ctx
-                    .runtime_env()
-                    .object_store(scan.path.as_str())
-                    .map_err(|e| {
+                let path = ListingTableUrl::parse(scan.path.as_str())?;
+
+                let object_store =
+                    ctx.runtime_env().object_store(&path).map_err(|e| {
                         BallistaError::NotImplemented(format!(
                             "No object store is registered for path {}: {:?}",
                             scan.path, e
                         ))
-                    })?
-                    .0;
+                    })?;
 
-                println!(
-                    "Found object store {:?} for path {}",
-                    object_store,
-                    scan.path.as_str()
-                );
+                println!("Found object store {:?} for path {}", object_store, path);
 
-                let config = ListingTableConfig::new(object_store, scan.path.as_str())
+                let config = ListingTableConfig::new(object_store, path)
                     .with_listing_options(options)
                     .with_schema(Arc::new(schema));
 
@@ -1043,6 +1040,7 @@ mod roundtrip_tests {
     use async_trait::async_trait;
     use core::panic;
     use datafusion::common::DFSchemaRef;
+    use datafusion::datasource::listing::ListingTableUrl;
     use datafusion::logical_plan::source_as_provider;
     use datafusion::{
         arrow::datatypes::{DataType, Field, Schema},
@@ -1163,7 +1161,7 @@ mod roundtrip_tests {
             vec![col("c1") + col("c2"), Expr::Literal((4.0).into())];
 
         let plan = std::sync::Arc::new(
-            test_scan_csv("employee.csv", Some(vec![3, 4]))
+            test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
                 .await?
                 .sort(vec![col("salary")])?
                 .build()?,
@@ -1235,13 +1233,13 @@ mod roundtrip_tests {
 
     #[tokio::test]
     async fn roundtrip_analyze() -> Result<()> {
-        let verbose_plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let verbose_plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .sort(vec![col("salary")])?
             .explain(true, true)?
             .build()?;
 
-        let plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .sort(vec![col("salary")])?
             .explain(false, true)?
@@ -1256,13 +1254,13 @@ mod roundtrip_tests {
 
     #[tokio::test]
     async fn roundtrip_explain() -> Result<()> {
-        let verbose_plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let verbose_plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .sort(vec![col("salary")])?
             .explain(true, false)?
             .build()?;
 
-        let plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .sort(vec![col("salary")])?
             .explain(false, false)?
@@ -1277,11 +1275,11 @@ mod roundtrip_tests {
 
     #[tokio::test]
     async fn roundtrip_join() -> Result<()> {
-        let scan_plan = test_scan_csv("employee1", Some(vec![0, 3, 4]))
+        let scan_plan = test_scan_csv("file:///employee1", Some(vec![0, 3, 4]))
             .await?
             .build()?;
 
-        let plan = test_scan_csv("employee2", Some(vec![0, 3, 4]))
+        let plan = test_scan_csv("file:///employee2", Some(vec![0, 3, 4]))
             .await?
             .join(&scan_plan, JoinType::Inner, (vec!["id"], vec!["id"]))?
             .build()?;
@@ -1292,7 +1290,7 @@ mod roundtrip_tests {
 
     #[tokio::test]
     async fn roundtrip_sort() -> Result<()> {
-        let plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .sort(vec![col("salary")])?
             .build()?;
@@ -1316,7 +1314,7 @@ mod roundtrip_tests {
 
     #[tokio::test]
     async fn roundtrip_logical_plan() -> Result<()> {
-        let plan = test_scan_csv("employee.csv", Some(vec![3, 4]))
+        let plan = test_scan_csv("file:///employee.csv", Some(vec![3, 4]))
             .await?
             .aggregate(vec![col("state")], vec![max(col("salary"))])?
             .build()?;
@@ -1336,9 +1334,10 @@ mod roundtrip_tests {
         ctx.runtime_env()
             .register_object_store("test", custom_object_store.clone());
 
-        let (os, uri) = ctx.runtime_env().object_store("test://foo.csv")?;
+        let uri = ListingTableUrl::parse("test://foo.csv")?;
+        let os = ctx.runtime_env().object_store(&uri)?;
         assert_eq!("TestObjectStore", &format!("{:?}", os));
-        assert_eq!("foo.csv", uri);
+        assert_eq!("foo.csv", uri.to_string());
 
         let schema = test_schema();
         let plan = ctx

ballista/rust/core/src/serde/physical_plan/mod.rs

@@ -1008,7 +1008,7 @@ impl AsExecutionPlan for PhysicalPlanNode {
 
 fn decode_scan_config(
     proto: &protobuf::FileScanExecConf,
-    runtime: &RuntimeEnv,
+    _runtime: &RuntimeEnv,
 ) -> Result<FileScanConfig, BallistaError> {
     let schema = Arc::new(convert_required!(proto.schema)?);
     let projection = proto
@@ -1029,14 +1029,10 @@ fn decode_scan_config(
         .map(|f| f.try_into())
         .collect::<Result<Vec<_>, _>>()?;
 
-    let object_store = if let Some(file) = file_groups.get(0).and_then(|h| h.get(0)) {
-        runtime.object_store(file.file_meta.path())?.0
-    } else {
-        Arc::new(LocalFileSystem {})
-    };
+    // TODO: This will not roundtrip object storage correctly as it loses the scheme
 
     Ok(FileScanConfig {
-        object_store,
+        object_store: Arc::new(LocalFileSystem {}),
         file_schema: schema,
         file_groups,
         statistics,

datafusion-examples/examples/flight_server.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use arrow_flight::SchemaAsIpc;
 use datafusion::datafusion_data_access::object_store::local::LocalFileSystem;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
-use datafusion::datasource::listing::ListingOptions;
+use datafusion::datasource::listing::{ListingOptions, ListingTableUrl};
 use futures::Stream;
 use tonic::transport::Server;
 use tonic::{Request, Response, Status, Streaming};
@@ -68,9 +68,10 @@ impl FlightService for FlightServiceImpl {
         let request = request.into_inner();
 
         let listing_options = ListingOptions::new(Arc::new(ParquetFormat::default()));
+        let url = ListingTableUrl::parse(&request.path[0])?;
 
         let schema = listing_options
-            .infer_schema(Arc::new(LocalFileSystem {}), &request.path[0])
+            .infer_schema(Arc::new(LocalFileSystem {}), url)
             .await
             .unwrap();
 

datafusion/core/Cargo.toml

@@ -66,7 +66,9 @@ datafusion-jit = { path = "../jit", version = "8.0.0", optional = true }
 datafusion-physical-expr = { path = "../physical-expr", version = "8.0.0" }
 datafusion-row = { path = "../row", version = "8.0.0" }
 futures = "0.3"
+glob = "0.3.0"
 hashbrown = { version = "0.12", features = ["raw"] }
+itertools = "0.10"
 lazy_static = { version = "^1.4.0" }
 log = "^0.4"
 num-traits = { version = "0.2", optional = true }
@@ -84,6 +86,7 @@ sqlparser = "0.17"
 tempfile = "3"
 tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
 tokio-stream = "0.1"
+url = "2.2"
 uuid = { version = "1.0", features = ["v4"] }
 
 [dev-dependencies]

datafusion/core/benches/sort_limit_query_sql.rs

@@ -20,7 +20,9 @@ extern crate criterion;
 use criterion::Criterion;
 use datafusion::datafusion_data_access::object_store::local::LocalFileSystem;
 use datafusion::datasource::file_format::csv::CsvFormat;
-use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig};
+use datafusion::datasource::listing::{
+    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
 
 use parking_lot::Mutex;
 use std::sync::Arc;
@@ -64,11 +66,12 @@ fn create_context() -> Arc<Mutex<SessionContext>> {
     let testdata = datafusion::test_util::arrow_test_data();
 
     let path = format!("{}/csv/aggregate_test_100.csv", testdata);
+    let url = ListingTableUrl::parse(path).unwrap();
 
     // create CSV data source
     let listing_options = ListingOptions::new(Arc::new(CsvFormat::default()));
 
-    let config = ListingTableConfig::new(Arc::new(LocalFileSystem {}), &path)
+    let config = ListingTableConfig::new(Arc::new(LocalFileSystem {}), url)
         .with_listing_options(listing_options)
         .with_schema(schema);
 

datafusion/core/src/catalog/schema.rs

@@ -23,7 +23,7 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::datasource::listing::{ListingTable, ListingTableConfig};
+use crate::datasource::listing::{ListingTable, ListingTableConfig, ListingTableUrl};
 use crate::datasource::object_store_registry::ObjectStoreRegistry;
 use crate::datasource::TableProvider;
 use crate::error::{DataFusionError, Result};
@@ -157,10 +157,7 @@ impl ObjectStoreSchemaProvider {
     }
 
     /// Retrieves a `ObjectStore` instance by scheme
-    pub fn object_store<'a>(
-        &self,
-        uri: &'a str,
-    ) -> Result<(Arc<dyn ObjectStore>, &'a str)> {
+    pub fn object_store(&self, uri: &ListingTableUrl) -> Result<Arc<dyn ObjectStore>> {
         self.object_store_registry
             .lock()
             .get_by_uri(uri)
@@ -173,13 +170,13 @@ impl ObjectStoreSchemaProvider {
     pub async fn register_listing_table(
         &self,
         name: &str,
-        uri: &str,
+        uri: ListingTableUrl,
         config: Option<ListingTableConfig>,
     ) -> Result<()> {
         let config = match config {
             Some(cfg) => cfg,
             None => {
-                let (object_store, _path) = self.object_store(uri)?;
+                let object_store = self.object_store(&uri)?;
                 ListingTableConfig::new(object_store, uri).infer().await?
             }
         };
@@ -255,6 +252,7 @@ mod tests {
     use crate::datasource::empty::EmptyTable;
     use crate::execution::context::SessionContext;
 
+    use crate::datasource::listing::ListingTableUrl;
     use futures::StreamExt;
 
     #[tokio::test]
@@ -280,12 +278,13 @@ mod tests {
     async fn test_schema_register_listing_table() {
         let testdata = crate::test_util::parquet_test_data();
         let filename = format!("{}/{}", testdata, "alltypes_plain.parquet");
+        let uri = ListingTableUrl::parse(filename).unwrap();
 
         let schema = ObjectStoreSchemaProvider::new();
         let _store = schema.register_object_store("test", Arc::new(LocalFileSystem {}));
 
         schema
-            .register_listing_table("alltypes_plain", &filename, None)
+            .register_listing_table("alltypes_plain", uri, None)
             .await
             .unwrap();
 
@@ -338,8 +337,9 @@ mod tests {
                 || file == OsStr::new("alltypes_plain.parquet")
             {
                 let name = path.file_stem().unwrap().to_str().unwrap();
+                let path = ListingTableUrl::parse(&sized_file.path).unwrap();
                 schema
-                    .register_listing_table(name, &sized_file.path, None)
+                    .register_listing_table(name, path, None)
                     .await
                     .unwrap();
             }
@@ -360,17 +360,18 @@ mod tests {
     async fn test_schema_register_same_listing_table() {
         let testdata = crate::test_util::parquet_test_data();
         let filename = format!("{}/{}", testdata, "alltypes_plain.parquet");
+        let uri = ListingTableUrl::parse(filename).unwrap();
 
         let schema = ObjectStoreSchemaProvider::new();
         let _store = schema.register_object_store("test", Arc::new(LocalFileSystem {}));
 
         schema
-            .register_listing_table("alltypes_plain", &filename, None)
+            .register_listing_table("alltypes_plain", uri.clone(), None)
             .await
             .unwrap();
 
         schema
-            .register_listing_table("alltypes_plain", &filename, None)
+            .register_listing_table("alltypes_plain", uri, None)
             .await
             .unwrap();
     }