[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #4042: Extract common parquet testing code to `parquet-test-util` crate

[GitBox <gi...@apache.org>]

alamb commented on code in PR #4042:
URL: https://github.com/apache/arrow-datafusion/pull/4042#discussion_r1010399611


##########
parquet-test-utils/src/lib.rs:
##########
@@ -0,0 +1,214 @@
+// 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.
+
+//! Helpers for writing parquet files and reading them back
+
+use std::fs::File;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use datafusion::arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
+use datafusion::common::ToDFSchema;
+use datafusion::config::{
+    ConfigOptions, OPT_PARQUET_ENABLE_PAGE_INDEX, OPT_PARQUET_PUSHDOWN_FILTERS,
+    OPT_PARQUET_REORDER_FILTERS,
+};
+use datafusion::datasource::listing::{ListingTableUrl, PartitionedFile};
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::error::Result;
+use datafusion::optimizer::simplify_expressions::{ExprSimplifier, SimplifyContext};
+use datafusion::physical_expr::create_physical_expr;
+use datafusion::physical_expr::execution_props::ExecutionProps;
+use datafusion::physical_plan::file_format::{FileScanConfig, ParquetExec};
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::metrics::MetricsSet;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::Expr;
+use object_store::path::Path;
+use object_store::ObjectMeta;
+use parquet::arrow::ArrowWriter;
+use parquet::file::properties::WriterProperties;
+
+///  a ParquetFile that has been created for testing.
+pub struct TestParquetFile {
+    path: PathBuf,
+    schema: SchemaRef,
+    object_store_url: ObjectStoreUrl,
+    object_meta: ObjectMeta,
+}
+
+#[derive(Debug, Clone)]
+pub struct ParquetScanOptions {
+    pub pushdown_filters: bool,
+    pub reorder_filters: bool,
+    pub enable_page_index: bool,
+}
+
+impl TestParquetFile {
+    /// Creates a new parquet file at the specified location
+    pub fn try_new(
+        path: PathBuf,

Review Comment:
   This is a good idea - thank you for the suggestion -- I tried to do this, but it seems as if the object store interface requires the length and canonical path which are currently computed from the path
   
   ```
   
           let size = std::fs::metadata(&path)?.len() as usize;
   
           let canonical_path = path.canonicalize()?;
   ```
   
   Given this code is no worse than today I would like to leave it in this PR -- it  does in fact leave large files around which I will attempt to clean up in #3976 



##########
parquet-test-utils/src/lib.rs:
##########
@@ -0,0 +1,214 @@
+// 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.
+
+//! Helpers for writing parquet files and reading them back
+
+use std::fs::File;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use datafusion::arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
+use datafusion::common::ToDFSchema;
+use datafusion::config::{
+    ConfigOptions, OPT_PARQUET_ENABLE_PAGE_INDEX, OPT_PARQUET_PUSHDOWN_FILTERS,
+    OPT_PARQUET_REORDER_FILTERS,
+};
+use datafusion::datasource::listing::{ListingTableUrl, PartitionedFile};
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::error::Result;
+use datafusion::optimizer::simplify_expressions::{ExprSimplifier, SimplifyContext};
+use datafusion::physical_expr::create_physical_expr;
+use datafusion::physical_expr::execution_props::ExecutionProps;
+use datafusion::physical_plan::file_format::{FileScanConfig, ParquetExec};
+use datafusion::physical_plan::filter::FilterExec;
+use datafusion::physical_plan::metrics::MetricsSet;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::Expr;
+use object_store::path::Path;
+use object_store::ObjectMeta;
+use parquet::arrow::ArrowWriter;
+use parquet::file::properties::WriterProperties;
+
+///  a ParquetFile that has been created for testing.
+pub struct TestParquetFile {
+    path: PathBuf,
+    schema: SchemaRef,
+    object_store_url: ObjectStoreUrl,
+    object_meta: ObjectMeta,
+}
+
+#[derive(Debug, Clone)]
+pub struct ParquetScanOptions {
+    pub pushdown_filters: bool,
+    pub reorder_filters: bool,
+    pub enable_page_index: bool,
+}
+
+impl TestParquetFile {
+    /// Creates a new parquet file at the specified location
+    pub fn try_new(
+        path: PathBuf,
+        batches: impl IntoIterator<Item = RecordBatch>,
+        page_size: Option<usize>,
+        row_group_size: Option<usize>,
+    ) -> Result<Self> {
+        let file = File::create(&path).unwrap();
+
+        let mut props_builder = WriterProperties::builder();
+
+        if let Some(s) = page_size {
+            props_builder = props_builder
+                .set_data_pagesize_limit(s)
+                .set_write_batch_size(s);
+        }
+
+        if let Some(s) = row_group_size {
+            props_builder = props_builder.set_max_row_group_size(s);
+        }
+
+        let mut batches = batches.into_iter();
+        let first_batch = batches.next().expect("need at least one record batch");
+        let schema = first_batch.schema();
+
+        let mut writer =
+            ArrowWriter::try_new(file, schema.clone(), Some(props_builder.build()))
+                .unwrap();
+
+        writer.write(&first_batch).unwrap();
+        writer.flush()?;
+        let mut num_rows = first_batch.num_rows();
+
+        for batch in batches {
+            writer.write(&batch).unwrap();
+            writer.flush()?;
+            num_rows += batch.num_rows();
+        }
+        writer.close().unwrap();
+
+        println!("Generated test dataset with {} rows", num_rows);
+
+        let size = std::fs::metadata(&path)?.len() as usize;
+
+        let canonical_path = path.canonicalize()?;
+
+        let object_store_url =
+            ListingTableUrl::parse(canonical_path.to_str().unwrap_or_default())?
+                .object_store();
+
+        let object_meta = ObjectMeta {
+            location: Path::parse(canonical_path.to_str().unwrap_or_default())?,
+            last_modified: Default::default(),
+            size,
+        };
+
+        Ok(Self {
+            path,
+            schema,
+            object_store_url,
+            object_meta,
+        })
+    }
+
+    /// return a `ParquetExec` and `FilterExec` with the specified options to scan this parquet file.
+    ///
+    /// This returns the same plan that DataFusion will make with a pushed down predicate followed by a filter:
+    ///
+    /// ```text
+    /// (FilterExec)
+    ///   (ParquetExec)
+    /// ```
+    pub async fn create_scan(
+        &self,
+        filter: Expr,
+        scan_options: ParquetScanOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let ParquetScanOptions {
+            pushdown_filters,
+            reorder_filters,
+            enable_page_index,
+        } = scan_options;
+
+        let mut config_options = ConfigOptions::new();
+        config_options.set_bool(OPT_PARQUET_PUSHDOWN_FILTERS, pushdown_filters);
+        config_options.set_bool(OPT_PARQUET_REORDER_FILTERS, reorder_filters);
+        config_options.set_bool(OPT_PARQUET_ENABLE_PAGE_INDEX, enable_page_index);
+
+        let scan_config = FileScanConfig {
+            object_store_url: self.object_store_url.clone(),
+            file_schema: self.schema.clone(),
+            file_groups: vec![vec![PartitionedFile {
+                object_meta: self.object_meta.clone(),
+                partition_values: vec![],
+                range: None,
+                extensions: None,
+            }]],
+            statistics: Default::default(),
+            projection: None,
+            limit: None,
+            table_partition_cols: vec![],
+            config_options: config_options.into_shareable(),
+        };
+
+        let df_schema = self.schema.clone().to_dfschema_ref()?;
+
+        // run coercion on the filters to coerce types etc.
+        let props = ExecutionProps::new();
+        let context = SimplifyContext::new(&props).with_schema(df_schema.clone());
+        let simplifier = ExprSimplifier::new(context);
+        let filter = simplifier.coerce(filter, df_schema.clone()).unwrap();
+
+        let physical_filter_expr = create_physical_expr(
+            &filter,
+            &df_schema,
+            self.schema.as_ref(),
+            &ExecutionProps::default(),
+        )?;
+
+        let parquet_exec = Arc::new(ParquetExec::new(scan_config, Some(filter), None));
+
+        let exec = Arc::new(FilterExec::try_new(physical_filter_expr, parquet_exec)?);
+
+        Ok(exec)
+    }
+
+    /// Retrieve metrics from the parquet exec returned from `create_scan`
+    ///
+    /// Recursively searches for ParquetExec and returns the metrics
+    /// on the first one it finds
+    pub fn parquet_metrics(&self, plan: Arc<dyn ExecutionPlan>) -> Option<MetricsSet> {

Review Comment:
   No -- I am still in the object oriented frame of mind -- will remove



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org