Merge remote-tracking branch 'apache/master' into alamb/consolidate_parquet_tests

Author: alamb

.github/pull_request_template.md

@@ -19,6 +19,16 @@ Closes #.
 There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
 -->
 
+# Are these changes tested?
+
+<!--
+We typically require tests for all PRs in order to:
+1. Prevent the code from being accidentally broken by subsequent changes
+2. Serve as another way to document the expected behavior of the code
+
+If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
+-->
+
 # Are there any user-facing changes?
 
 <!--

.github/workflows/rust.yml

@@ -477,3 +477,36 @@ jobs:
           # ignore ./Cargo.toml because putting workspaces in multi-line lists make it easy to read
           ci/scripts/rust_toml_fmt.sh
           git diff --exit-code
+
+  config-docs-check:
+    name: check configs.md is up-to-date
+    needs: [linux-build-lib]
+    runs-on: ubuntu-latest
+    container:
+      image: amd64/rust
+      env:
+        # Disable full debug symbol generation to speed up CI build and keep memory down
+        # "1" means line tables only, which is useful for panic tracebacks.
+        RUSTFLAGS: "-C debuginfo=1"
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          submodules: true
+      - name: Cache Cargo
+        uses: actions/cache@v3
+        with:
+          path: /github/home/.cargo
+          # this key equals the ones on `linux-build-lib` for re-use
+          key: cargo-cache-
+      - name: Setup Rust toolchain
+        uses: ./.github/actions/setup-builder
+        with:
+          rust-version: stable
+      - uses: actions/setup-node@v3
+        with:
+          node-version: "14"
+      - name: Check if configs.md has been modified
+        run: |
+          # If you encounter an error, run './dev/update_config_docs.sh' and commit
+          ./dev/update_config_docs.sh
+          git diff --exit-code

benchmarks/src/bin/parquet_filter_pushdown.rs

@@ -21,6 +21,7 @@ use datafusion::logical_expr::{lit, or, Expr};
 use datafusion::optimizer::utils::disjunction;
 use datafusion::physical_plan::collect;
 use datafusion::prelude::{col, SessionConfig, SessionContext};
+use parquet::file::properties::WriterProperties;
 use parquet_test_utils::{ParquetScanOptions, TestParquetFile};
 use std::path::PathBuf;
 use std::time::Instant;
@@ -73,7 +74,19 @@ async fn main() -> Result<()> {
 
     let path = opt.path.join("logs.parquet");
 
-    let test_file = gen_data(path, opt.scale_factor, opt.page_size, opt.row_group_size)?;
+    let mut props_builder = WriterProperties::builder();
+
+    if let Some(s) = opt.page_size {
+        props_builder = props_builder
+            .set_data_pagesize_limit(s)
+            .set_write_batch_size(s);
+    }
+
+    if let Some(s) = opt.row_group_size {
+        props_builder = props_builder.set_max_row_group_size(s);
+    }
+
+    let test_file = gen_data(path, opt.scale_factor, props_builder.build())?;
 
     run_benchmarks(&mut ctx, &test_file, opt.iterations, opt.debug).await?;
 
@@ -137,14 +150,9 @@ async fn run_benchmarks(
             println!("Using scan options {:?}", scan_options);
             for i in 0..iterations {
                 let start = Instant::now();
-                let rows = exec_scan(
-                    ctx,
-                    test_file,
-                    filter_expr.clone(),
-                    scan_options.clone(),
-                    debug,
-                )
-                .await?;
+                let rows =
+                    exec_scan(ctx, test_file, filter_expr.clone(), *scan_options, debug)
+                        .await?;
                 println!(
                     "Iteration {} returned {} rows in {} ms",
                     i,
@@ -179,17 +187,11 @@ async fn exec_scan(
 fn gen_data(
     path: PathBuf,
     scale_factor: f32,
-    page_size: Option<usize>,
-    row_group_size: Option<usize>,
+    props: WriterProperties,
 ) -> Result<TestParquetFile> {
     let generator = AccessLogGenerator::new();
 
     let num_batches = 100_f32 * scale_factor;
 
-    TestParquetFile::try_new(
-        path,
-        generator.take(num_batches as usize),
-        page_size,
-        row_group_size,
-    )
+    TestParquetFile::try_new(path, props, generator.take(num_batches as usize))
 }

datafusion-cli/src/main.rs

@@ -16,6 +16,9 @@
 // under the License.
 
 use clap::Parser;
+use datafusion::datasource::datasource::TableProviderFactory;
+use datafusion::datasource::file_format::file_type::FileType;
+use datafusion::datasource::listing_table_factory::ListingTableFactory;
 use datafusion::datasource::object_store::ObjectStoreRegistry;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::SessionConfig;
@@ -26,6 +29,7 @@ use datafusion_cli::{
     exec, print_format::PrintFormat, print_options::PrintOptions, DATAFUSION_CLI_VERSION,
 };
 use mimalloc::MiMalloc;
+use std::collections::HashMap;
 use std::env;
 use std::path::Path;
 use std::sync::Arc;
@@ -93,7 +97,7 @@ pub async fn main() -> Result<()> {
 
     if let Some(ref path) = args.data_path {
         let p = Path::new(path);
-        env::set_current_dir(&p).unwrap();
+        env::set_current_dir(p).unwrap();
     };
 
     let mut session_config = SessionConfig::from_env().with_information_schema(true);
@@ -105,6 +109,7 @@ pub async fn main() -> Result<()> {
     let runtime_env = create_runtime_env()?;
     let mut ctx =
         SessionContext::with_config_rt(session_config.clone(), Arc::new(runtime_env));
+    ctx.refresh_catalogs().await?;
 
     let mut print_options = PrintOptions {
         format: args.format,
@@ -142,11 +147,31 @@ pub async fn main() -> Result<()> {
 }
 
 fn create_runtime_env() -> Result<RuntimeEnv> {
+    let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
+        HashMap::new();
+    table_factories.insert(
+        "csv".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::CSV)),
+    );
+    table_factories.insert(
+        "parquet".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::PARQUET)),
+    );
+    table_factories.insert(
+        "avro".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::AVRO)),
+    );
+    table_factories.insert(
+        "json".to_string(),
+        Arc::new(ListingTableFactory::new(FileType::JSON)),
+    );
+
     let object_store_provider = DatafusionCliObjectStoreProvider {};
     let object_store_registry =
         ObjectStoreRegistry::new_with_provider(Some(Arc::new(object_store_provider)));
-    let rn_config =
-        RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry));
+    let rn_config = RuntimeConfig::new()
+        .with_object_store_registry(Arc::new(object_store_registry))
+        .with_table_factories(table_factories);
     RuntimeEnv::new(rn_config)
 }
 

datafusion-cli/src/object_storage.rs

@@ -141,7 +141,13 @@ mod tests {
         assert!(err.to_string().contains("Generic S3 error: Missing region"));
 
         env::set_var("AWS_REGION", "us-east-1");
-        assert!(provider.get_by_url(&Url::from_str(s3).unwrap()).is_ok());
+        let url = Url::from_str(s3).expect("Unable to parse s3 url");
+        let res = provider.get_by_url(&url);
+        let msg = match res {
+            Err(e) => format!("{}", e),
+            Ok(_) => "".to_string()
+        };
+        assert_eq!("".to_string(), msg); // Fail with error message
         env::remove_var("AWS_REGION");
     }
 }

datafusion-cli/src/print_format.rs

@@ -107,9 +107,9 @@ mod tests {
         let batch = RecordBatch::try_new(
             schema,
             vec![
-                Arc::new(Int32Array::from_slice(&[1, 2, 3])),
-                Arc::new(Int32Array::from_slice(&[4, 5, 6])),
-                Arc::new(Int32Array::from_slice(&[7, 8, 9])),
+                Arc::new(Int32Array::from_slice([1, 2, 3])),
+                Arc::new(Int32Array::from_slice([4, 5, 6])),
+                Arc::new(Int32Array::from_slice([7, 8, 9])),
             ],
         )
         .unwrap();
@@ -137,9 +137,9 @@ mod tests {
         let batch = RecordBatch::try_new(
             schema,
             vec![
-                Arc::new(Int32Array::from_slice(&[1, 2, 3])),
-                Arc::new(Int32Array::from_slice(&[4, 5, 6])),
-                Arc::new(Int32Array::from_slice(&[7, 8, 9])),
+                Arc::new(Int32Array::from_slice([1, 2, 3])),
+                Arc::new(Int32Array::from_slice([4, 5, 6])),
+                Arc::new(Int32Array::from_slice([7, 8, 9])),
             ],
         )
         .unwrap();

datafusion-examples/README.md

@@ -44,4 +44,4 @@ Run `git submodule update --init` to init test files.
 
 ## Distributed
 
-- [`flight-client.rs`](examples/flight-client.rs) and [`flight-server.rs`](examples/flight-server.rs): Run DataFusion as a standalone process and execute SQL queries from a client using the Flight protocol.
+- [`flight_client.rs`](examples/flight_client.rs) and [`flight_server.rs`](examples/flight_server.rs): Run DataFusion as a standalone process and execute SQL queries from a client using the Flight protocol.

datafusion/core/src/catalog/listing_schema.rs

@@ -19,8 +19,10 @@
 use crate::catalog::schema::SchemaProvider;
 use crate::datasource::datasource::TableProviderFactory;
 use crate::datasource::TableProvider;
-use datafusion_common::DataFusionError;
+use crate::execution::context::SessionState;
+use datafusion_common::{context, DataFusionError};
 use futures::TryStreamExt;
+use itertools::Itertools;
 use object_store::ObjectStore;
 use std::any::Any;
 use std::collections::{HashMap, HashSet};
@@ -72,7 +74,7 @@ impl ListingSchemaProvider {
     }
 
     /// Reload table information from ObjectStore
-    pub async fn refresh(&self) -> datafusion_common::Result<()> {
+    pub async fn refresh(&self, state: &SessionState) -> datafusion_common::Result<()> {
         let entries: Vec<_> = self
             .store
             .list(Some(&self.path))
@@ -100,13 +102,20 @@ impl ListingSchemaProvider {
                 .ok_or_else(|| {
                     DataFusionError::Internal("Cannot parse file name!".to_string())
                 })?;
-            let table_name = table.to_str().ok_or_else(|| {
+            let table_name = file_name.split('.').collect_vec()[0];
+            let table_path = table.to_str().ok_or_else(|| {
                 DataFusionError::Internal("Cannot parse file name!".to_string())
             })?;
-            if !self.table_exist(file_name) {
-                let table_name = format!("{}/{}", self.authority, table_name);
-                let provider = self.factory.create(table_name.as_str()).await?;
-                let _ = self.register_table(file_name.to_string(), provider.clone())?;
+            if !self.table_exist(table_name) {
+                let table_url = format!("{}/{}", self.authority, table_path);
+                let provider = self
+                    .factory
+                    .create(state, table_url.as_str())
+                    .await
+                    .map_err(|e| {
+                        context!(format!("Could not create table for {}", table_url), e)
+                    })?;
+                let _ = self.register_table(table_name.to_string(), provider.clone())?;
             }
         }
         Ok(())

datafusion/core/src/datasource/datasource.rs

@@ -92,5 +92,9 @@ pub trait TableProvider: Sync + Send {
 #[async_trait]
 pub trait TableProviderFactory: Sync + Send {
     /// Create a TableProvider with the given url
-    async fn create(&self, url: &str) -> Result<Arc<dyn TableProvider>>;
+    async fn create(
+        &self,
+        ctx: &SessionState,
+        url: &str,
+    ) -> Result<Arc<dyn TableProvider>>;
 }

datafusion/core/src/datasource/listing_table_factory.rs

@@ -0,0 +1,79 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Factory for creating ListingTables with default options
+
+use crate::datasource::datasource::TableProviderFactory;
+use crate::datasource::file_format::avro::AvroFormat;
+use crate::datasource::file_format::csv::CsvFormat;
+use crate::datasource::file_format::file_type::{FileType, GetExt};
+use crate::datasource::file_format::json::JsonFormat;
+use crate::datasource::file_format::parquet::ParquetFormat;
+use crate::datasource::file_format::FileFormat;
+use crate::datasource::listing::{
+    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
+use crate::datasource::TableProvider;
+use crate::execution::context::SessionState;
+use async_trait::async_trait;
+use std::sync::Arc;
+
+/// A `TableProviderFactory` capable of creating new `ListingTable`s
+pub struct ListingTableFactory {
+    file_type: FileType,
+}
+
+impl ListingTableFactory {
+    /// Creates a new `ListingTableFactory`
+    pub fn new(file_type: FileType) -> Self {
+        Self { file_type }
+    }
+}
+
+#[async_trait]
+impl TableProviderFactory for ListingTableFactory {
+    async fn create(
+        &self,
+        state: &SessionState,
+        url: &str,
+    ) -> datafusion_common::Result<Arc<dyn TableProvider>> {
+        let file_extension = self.file_type.get_ext();
+
+        let file_format: Arc<dyn FileFormat> = match self.file_type {
+            FileType::CSV => Arc::new(CsvFormat::default()),
+            FileType::PARQUET => Arc::new(ParquetFormat::default()),
+            FileType::AVRO => Arc::new(AvroFormat::default()),
+            FileType::JSON => Arc::new(JsonFormat::default()),
+        };
+
+        let options = ListingOptions {
+            format: file_format,
+            collect_stat: true,
+            file_extension: file_extension.to_owned(),
+            target_partitions: 1,
+            table_partition_cols: vec![],
+        };
+
+        let table_path = ListingTableUrl::parse(url)?;
+        let resolved_schema = options.infer_schema(state, &table_path).await?;
+        let config = ListingTableConfig::new(table_path)
+            .with_listing_options(options)
+            .with_schema(resolved_schema);
+        let table = ListingTable::try_new(config)?;
+        Ok(Arc::new(table))
+    }
+}