You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/10/18 12:14:39 UTC
[GitHub] [arrow-datafusion] rdettai opened a new pull request #1141: [feat] adding partition_values to PartitionedFile
rdettai opened a new pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141
# Which issue does this PR close?
Closes #1139.
# Rationale for this change
Adding capability to parse file partitioning and prune unnecessary files
# What changes are included in this PR?
- add `partition_values: Vec<ScalarValue>` to `PartitionedFile`
- implement pruned_partition_list
- add extra column for the partition dimensions to `execute()` record batch result in file format execution plans
- make the partitioning configuration available through config
# Are there any user-facing changes?
Adding possibility to configure input file partitioning
--
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
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740235574
##########
File path: datafusion/src/datasource/object_store/mod.rs
##########
@@ -70,7 +70,7 @@ pub enum ListEntry {
}
/// The path and size of the file.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
Review comment:
```suggestion
#[derive(Debug, Clone, PartialEq, Eq)]
```
--
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
[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738724739
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -613,33 +614,28 @@ message ScanLimit {
uint32 limit = 1;
}
-message ParquetScanExecNode {
+message FileScanExecConf {
repeated FileGroup file_groups = 1;
Schema schema = 2;
- uint32 batch_size = 4;
- repeated uint32 projection = 6;
- ScanLimit limit = 7;
- Statistics statistics = 8;
+ uint32 batch_size = 3;
+ repeated uint32 projection = 4;
+ ScanLimit limit = 5;
+ Statistics statistics = 6;
+ repeated string table_partition_cols = 7;
+}
+
+message ParquetScanExecNode {
+ FileScanExecConf base_conf = 1;
Review comment:
👍
##########
File path: ballista/rust/core/src/serde/physical_plan/to_proto.rs
##########
@@ -244,90 +246,29 @@ impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
))),
})
} else if let Some(exec) = plan.downcast_ref::<CsvExec>() {
- let file_groups = exec
- .file_groups()
- .iter()
- .map(|p| p.as_slice().into())
- .collect();
Ok(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::CsvScan(
protobuf::CsvScanExecNode {
- file_groups,
- statistics: Some((&exec.statistics()).into()),
- limit: exec
- .limit()
- .map(|l| protobuf::ScanLimit { limit: l as u32 }),
- projection: exec
- .projection()
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General(
- "projection in CsvExec does not exist.".to_owned(),
- )
- })?
- .iter()
- .map(|n| *n as u32)
- .collect(),
- schema: Some(exec.file_schema().as_ref().into()),
+ base_conf: Some(exec.base_config().try_into()?),
has_header: exec.has_header(),
delimiter: byte_to_string(exec.delimiter())?,
- batch_size: exec.batch_size() as u32,
},
)),
})
} else if let Some(exec) = plan.downcast_ref::<ParquetExec>() {
- let file_groups = exec
- .file_groups()
- .iter()
- .map(|p| p.as_slice().into())
- .collect();
-
Ok(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::ParquetScan(
protobuf::ParquetScanExecNode {
- file_groups,
- statistics: Some((&exec.statistics()).into()),
- limit: exec
- .limit()
- .map(|l| protobuf::ScanLimit { limit: l as u32 }),
- schema: Some(exec.schema().as_ref().into()),
- projection: exec
- .projection()
- .as_ref()
- .iter()
- .map(|n| *n as u32)
- .collect(),
- batch_size: exec.batch_size() as u32,
+ base_conf: Some(exec.base_config().try_into()?),
+ // TODO serialize predicates
Review comment:
is this a TODO you plan for this PR? Or a follow on one?
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -51,9 +52,9 @@ pub struct ListingOptions {
/// partitioning expected should be named "a" and "b":
/// - If there is a third level of partitioning it will be ignored.
/// - Files that don't follow this partitioning will be ignored.
- /// Note that only `DataType::Utf8` is supported for the column type.
- /// TODO implement case where partitions.len() > 0
- pub partitions: Vec<String>,
+ /// Note that only `DEFAULT_PARTITION_COLUMN_DATATYPE` is currently
+ /// supported for the column type.
Review comment:
as mentioned elsewhere I think it would be fine to say "these are always dictionary coded string columns" rather than "currently" which hints at changing it in the future.
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -613,33 +614,28 @@ message ScanLimit {
uint32 limit = 1;
}
-message ParquetScanExecNode {
+message FileScanExecConf {
Review comment:
makes sense to me
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
+}
+
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
+#[derive(Debug, Clone)]
+pub struct PhysicalPlanConfig {
+ /// Store from which the `files` should be fetched
+ pub object_store: Arc<dyn ObjectStore>,
+ /// Schema before projection. It contains the columns that are expected
+ /// to be in the files without the table partition columns.
+ pub file_schema: SchemaRef,
+ /// List of files to be processed, grouped into partitions
+ pub file_groups: Vec<Vec<PartitionedFile>>,
+ /// Estimated overall statistics of the files, taking `filters` into account.
+ pub statistics: Statistics,
+ /// Columns on which to project the data. Indexes that are higher than the
+ /// number of columns of `file_schema` refer to `table_partition_cols`.
+ pub projection: Option<Vec<usize>>,
+ /// The maximum number of records per arrow column
+ pub batch_size: usize,
+ /// The minimum number of records required from this source plan
+ pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_cols: Vec<String>,
+}
+
+impl PhysicalPlanConfig {
+ /// Project the schema and the statistics on the given column indices
+ fn project(&self) -> (SchemaRef, Statistics) {
Review comment:
The one thing that seems a little less than ideal to me is that we had a nice separation between file format related options and the storage of them. Partitioning columns seems to be more related to the storage rather than the format 🤔
But on the other hand most of the rest of the fields in `PhysicalPlanConfig` are not directly related to the format either.
I don't have any concrete suggestion here
##########
File path: ballista/rust/core/src/serde/physical_plan/to_proto.rs
##########
@@ -722,3 +675,32 @@ impl From<&Statistics> for protobuf::Statistics {
}
}
}
+
+impl TryFrom<&PhysicalPlanConfig> for protobuf::FileScanExecConf {
+ type Error = BallistaError;
+ fn try_from(
+ conf: &PhysicalPlanConfig,
+ ) -> Result<protobuf::FileScanExecConf, Self::Error> {
+ let file_groups = conf
+ .file_groups
+ .iter()
+ .map(|p| p.as_slice().try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+
+ Ok(protobuf::FileScanExecConf {
+ file_groups,
+ statistics: Some((&conf.statistics).into()),
+ limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }),
+ projection: conf
+ .projection
+ .as_ref()
+ .unwrap_or(&vec![])
Review comment:
it looks like an improvement to me, but previously the code would error if `projection: None` was passed and this code will simply convert that to an empty list.
Is that intentional?
##########
File path: datafusion/tests/path_partition.rs
##########
@@ -0,0 +1,392 @@
+// 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.
+
+//! Test queries on partitioned datasets
+
+use std::{fs, io, sync::Arc};
+
+use async_trait::async_trait;
+use datafusion::{
+ assert_batches_sorted_eq,
+ datasource::{
+ file_format::{csv::CsvFormat, parquet::ParquetFormat},
+ listing::{ListingOptions, ListingTable},
+ object_store::{
+ local::LocalFileSystem, FileMeta, FileMetaStream, ListEntryStream,
+ ObjectReader, ObjectStore, SizedFile,
+ },
+ },
+ error::{DataFusionError, Result},
+ physical_plan::ColumnStatistics,
+ prelude::ExecutionContext,
+ test_util::{arrow_test_data, parquet_test_data},
+};
+use futures::{stream, StreamExt};
+
+mod common;
+
+#[tokio::test]
+async fn csv_filter_with_file_col() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+
+ register_partitioned_aggregate_csv(
+ &mut ctx,
+ &[
Review comment:
so cool
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
Review comment:
I think this could be much more concisely be expressed using `ExpressionVisitor` -- perhaps following the model in https://github.com/apache/arrow-datafusion/blob/2c8e65bcca9fe41ad80116d99ad974c86cb59654/datafusion/src/optimizer/utils.rs#L60-L98
Where you can skip all the boiler plate related to recursing through the different Expr types and make it clear what types of Expr really matter.
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
+ /// An Arrow buffer initialized to zeros that represents the key array of all partition
+ /// columns (partition columns are materialized by dictionary arrays with only one
+ /// value in the dictionary, thus all the keys are equal to zero).
+ key_buffer_cache: Option<Buffer>,
+ /// Mapping between the indexes in the list of partition columns and the target
+ /// schema. Sorted by index in the target schema so that we can iterate on it to
+ /// insert the partition columns in the target record batch.
+ projected_partition_indexes: Vec<(usize, usize)>,
+ /// The schema of the table once the projection was applied.
+ projected_schema: SchemaRef,
+}
+
+impl PartitionColumnProjector {
+ // Create a projector to insert the partitioning columns into batches read from files
+ // - projected_schema: the target schema with both file and partitioning columns
+ // - table_partition_cols: all the partitioning column names
+ fn new(projected_schema: SchemaRef, table_partition_cols: &[String]) -> Self {
+ let mut idx_map = HashMap::new();
+ for (partition_idx, partition_name) in table_partition_cols.iter().enumerate() {
+ if let Ok(schema_idx) = projected_schema.index_of(partition_name) {
+ idx_map.insert(partition_idx, schema_idx);
+ }
+ }
+
+ let mut projected_partition_indexes: Vec<_> = idx_map.into_iter().collect();
+ projected_partition_indexes.sort_by(|(_, a), (_, b)| a.cmp(b));
+
+ Self {
+ projected_partition_indexes,
+ key_buffer_cache: None,
+ projected_schema,
+ }
+ }
+
+ // Transform the batch read from the fileby inserting the partitioning columns
Review comment:
```suggestion
// Transform the batch read from the file by inserting the partitioning columns
```
##########
File path: datafusion/tests/path_partition.rs
##########
@@ -0,0 +1,235 @@
+// 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.
+
+//! Test queries on partitioned datasets
Review comment:
I think it is fine -- sql.rs as you say is too large already
##########
File path: datafusion/src/physical_plan/sort_preserving_merge.rs
##########
@@ -936,15 +936,18 @@ mod tests {
test::create_partitioned_csv("aggregate_test_100.csv", partitions).unwrap();
let csv = Arc::new(CsvExec::new(
- Arc::new(LocalFileSystem {}),
- files,
- Statistics::default(),
- Arc::clone(&schema),
+ PhysicalPlanConfig {
Review comment:
man the same code was copied over and over and over -- lol
##########
File path: datafusion/src/datasource/file_format/avro.rs
##########
@@ -61,16 +62,9 @@ impl FileFormat for AvroFormat {
async fn create_physical_plan(
&self,
conf: PhysicalPlanConfig,
+ _filters: &[Expr],
Review comment:
this is neat that the creation of the physical plan gets the filters
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
+ /// An Arrow buffer initialized to zeros that represents the key array of all partition
+ /// columns (partition columns are materialized by dictionary arrays with only one
+ /// value in the dictionary, thus all the keys are equal to zero).
+ key_buffer_cache: Option<Buffer>,
+ /// Mapping between the indexes in the list of partition columns and the target
+ /// schema. Sorted by index in the target schema so that we can iterate on it to
+ /// insert the partition columns in the target record batch.
+ projected_partition_indexes: Vec<(usize, usize)>,
+ /// The schema of the table once the projection was applied.
+ projected_schema: SchemaRef,
+}
+
+impl PartitionColumnProjector {
+ // Create a projector to insert the partitioning columns into batches read from files
+ // - projected_schema: the target schema with both file and partitioning columns
+ // - table_partition_cols: all the partitioning column names
+ fn new(projected_schema: SchemaRef, table_partition_cols: &[String]) -> Self {
+ let mut idx_map = HashMap::new();
+ for (partition_idx, partition_name) in table_partition_cols.iter().enumerate() {
+ if let Ok(schema_idx) = projected_schema.index_of(partition_name) {
+ idx_map.insert(partition_idx, schema_idx);
+ }
+ }
+
+ let mut projected_partition_indexes: Vec<_> = idx_map.into_iter().collect();
+ projected_partition_indexes.sort_by(|(_, a), (_, b)| a.cmp(b));
+
+ Self {
+ projected_partition_indexes,
+ key_buffer_cache: None,
+ projected_schema,
+ }
+ }
+
+ // Transform the batch read from the fileby inserting the partitioning columns
+ // to the right positions as deduced from `projected_schema`
+ // - file_batch: batch read from the file, with internal projection applied
+ // - partition_values: the list of partition values, one for each partition column
+ fn project(
+ &mut self,
+ file_batch: RecordBatch,
+ partition_values: &[ScalarValue],
+ ) -> ArrowResult<RecordBatch> {
+ let expected_cols =
+ self.projected_schema.fields().len() - self.projected_partition_indexes.len();
+
+ if file_batch.columns().len() != expected_cols {
+ return Err(ArrowError::SchemaError(format!(
+ "Unexpected batch schema from file, expected {} cols but got {}",
+ expected_cols,
+ file_batch.columns().len()
+ )));
+ }
+
+ let mut cols = file_batch.columns().to_vec();
+ for &(pidx, sidx) in &self.projected_partition_indexes {
+ cols.insert(
+ sidx,
+ create_dict_array(
+ &mut self.key_buffer_cache,
+ &partition_values[pidx],
+ file_batch.num_rows(),
+ ),
+ )
+ }
+ RecordBatch::try_new(Arc::clone(&self.projected_schema), cols)
+ }
+}
+
+fn create_dict_array(
+ key_buffer_cache: &mut Option<Buffer>,
+ val: &ScalarValue,
+ len: usize,
+) -> ArrayRef {
+ // build value dictionary
+ let dict_vals = val.to_array();
+
+ // build keys array
+ let sliced_key_buffer = match key_buffer_cache {
Review comment:
👍
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
+ /// An Arrow buffer initialized to zeros that represents the key array of all partition
+ /// columns (partition columns are materialized by dictionary arrays with only one
+ /// value in the dictionary, thus all the keys are equal to zero).
+ key_buffer_cache: Option<Buffer>,
+ /// Mapping between the indexes in the list of partition columns and the target
+ /// schema. Sorted by index in the target schema so that we can iterate on it to
+ /// insert the partition columns in the target record batch.
+ projected_partition_indexes: Vec<(usize, usize)>,
+ /// The schema of the table once the projection was applied.
+ projected_schema: SchemaRef,
+}
+
+impl PartitionColumnProjector {
+ // Create a projector to insert the partitioning columns into batches read from files
+ // - projected_schema: the target schema with both file and partitioning columns
+ // - table_partition_cols: all the partitioning column names
+ fn new(projected_schema: SchemaRef, table_partition_cols: &[String]) -> Self {
+ let mut idx_map = HashMap::new();
+ for (partition_idx, partition_name) in table_partition_cols.iter().enumerate() {
+ if let Ok(schema_idx) = projected_schema.index_of(partition_name) {
+ idx_map.insert(partition_idx, schema_idx);
+ }
+ }
+
+ let mut projected_partition_indexes: Vec<_> = idx_map.into_iter().collect();
+ projected_partition_indexes.sort_by(|(_, a), (_, b)| a.cmp(b));
+
+ Self {
+ projected_partition_indexes,
+ key_buffer_cache: None,
+ projected_schema,
+ }
+ }
+
+ // Transform the batch read from the fileby inserting the partitioning columns
+ // to the right positions as deduced from `projected_schema`
+ // - file_batch: batch read from the file, with internal projection applied
+ // - partition_values: the list of partition values, one for each partition column
+ fn project(
+ &mut self,
+ file_batch: RecordBatch,
+ partition_values: &[ScalarValue],
+ ) -> ArrowResult<RecordBatch> {
+ let expected_cols =
+ self.projected_schema.fields().len() - self.projected_partition_indexes.len();
+
+ if file_batch.columns().len() != expected_cols {
+ return Err(ArrowError::SchemaError(format!(
+ "Unexpected batch schema from file, expected {} cols but got {}",
+ expected_cols,
+ file_batch.columns().len()
+ )));
+ }
+
+ let mut cols = file_batch.columns().to_vec();
+ for &(pidx, sidx) in &self.projected_partition_indexes {
+ cols.insert(
+ sidx,
+ create_dict_array(
+ &mut self.key_buffer_cache,
+ &partition_values[pidx],
+ file_batch.num_rows(),
+ ),
+ )
+ }
+ RecordBatch::try_new(Arc::clone(&self.projected_schema), cols)
+ }
+}
+
+fn create_dict_array(
+ key_buffer_cache: &mut Option<Buffer>,
+ val: &ScalarValue,
+ len: usize,
+) -> ArrayRef {
+ // build value dictionary
+ let dict_vals = val.to_array();
+
+ // build keys array
+ let sliced_key_buffer = match key_buffer_cache {
+ Some(buf) if buf.len() >= len => buf.slice(buf.len() - len),
+ _ => {
+ let mut key_buffer_builder = UInt8BufferBuilder::new(len);
+ key_buffer_builder.advance(len); // keys are all 0
+ key_buffer_cache.insert(key_buffer_builder.finish()).clone()
+ }
+ };
+
+ // create data type
+ let data_type =
+ DataType::Dictionary(Box::new(DataType::UInt8), Box::new(val.get_datatype()));
+
+ debug_assert_eq!(data_type, *DEFAULT_PARTITION_COLUMN_DATATYPE);
+
+ // assemble pieces together
+ let mut builder = ArrayData::builder(data_type)
+ .len(len)
+ .add_buffer(sliced_key_buffer);
+ builder = builder.add_child_data(dict_vals.data().clone());
+ Arc::new(DictionaryArray::<UInt8Type>::from(builder.build().unwrap()))
+}
+
+#[cfg(test)]
+mod tests {
+ use crate::test::{
+ aggr_test_schema, build_table_i32, columns, object_store::TestObjectStore,
+ };
+
+ use super::*;
+
+ #[test]
+ fn physical_plan_config_no_projection() {
+ let file_schema = aggr_test_schema();
+ let conf = config_for_projection(
+ Arc::clone(&file_schema),
+ None,
+ Statistics::default(),
+ vec!["date".to_owned()],
+ );
+
+ let (proj_schema, proj_statistics) = conf.project();
+ assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
+ assert_eq!(
+ proj_schema.field(file_schema.fields().len()).name(),
+ "date",
+ "partition columns are the last columns"
+ );
+ assert_eq!(
+ proj_statistics
+ .column_statistics
+ .expect("projection creates column statistics")
+ .len(),
+ file_schema.fields().len() + 1
+ );
+ // TODO implement tests for partition column statistics once implemented
+
+ let col_names = conf.projected_file_column_names();
+ assert_eq!(col_names, None);
+
+ let col_indices = conf.file_column_projection_indices();
+ assert_eq!(col_indices, None);
+ }
+
+ #[test]
+ fn physical_plan_config_with_projection() {
+ let file_schema = aggr_test_schema();
+ let conf = config_for_projection(
+ Arc::clone(&file_schema),
+ Some(vec![file_schema.fields().len(), 0]),
+ Statistics {
+ num_rows: Some(10),
+ // assign the column index to distinct_count to help assert
+ // the source statistic after the projection
+ column_statistics: Some(
+ (0..file_schema.fields().len())
+ .map(|i| ColumnStatistics {
+ distinct_count: Some(i),
+ ..Default::default()
+ })
+ .collect(),
+ ),
+ ..Default::default()
+ },
+ vec!["date".to_owned()],
+ );
+
+ let (proj_schema, proj_statistics) = conf.project();
+ assert_eq!(
+ columns(&proj_schema),
+ vec!["date".to_owned(), "c1".to_owned()]
+ );
+ let proj_stat_cols = proj_statistics
+ .column_statistics
+ .expect("projection creates column statistics");
+ assert_eq!(proj_stat_cols.len(), 2);
+ // TODO implement tests for proj_stat_cols[0] once partition column
+ // statistics are implemented
+ assert_eq!(proj_stat_cols[1].distinct_count, Some(0));
+
+ let col_names = conf.projected_file_column_names();
+ assert_eq!(col_names, Some(vec!["c1".to_owned()]));
+
+ let col_indices = conf.file_column_projection_indices();
+ assert_eq!(col_indices, Some(vec![0]));
+ }
+
+ #[test]
+ fn partition_column_projector() {
+ let file_batch = build_table_i32(
+ ("a", &vec![0, 1, 2]),
+ ("b", &vec![-2, -1, 0]),
+ ("c", &vec![10, 11, 12]),
+ );
+ let partition_cols =
+ vec!["year".to_owned(), "month".to_owned(), "day".to_owned()];
+ // create a projected schema
+ let conf = config_for_projection(
+ file_batch.schema(),
+ // keep all cols from file and 2 from partitioning
+ Some(vec![
+ 0,
+ 1,
+ 2,
+ file_batch.schema().fields().len(),
+ file_batch.schema().fields().len() + 2,
+ ]),
+ Statistics::default(),
+ partition_cols.clone(),
+ );
+ let (proj_schema, _) = conf.project();
+ // created a projector for that projected schema
+ let mut proj = PartitionColumnProjector::new(proj_schema, &partition_cols);
+ let projected_batch = proj
+ .project(
+ // file_batch is ok here because we kept all the file cols in the projection
+ file_batch,
+ &[
+ ScalarValue::Utf8(Some("2021".to_owned())),
+ ScalarValue::Utf8(Some("10".to_owned())),
+ ScalarValue::Utf8(Some("26".to_owned())),
+ ],
+ )
+ .expect("Projection of partition columns into record batch failed");
+
+ let expected = vec![
+ "+---+----+----+------+-----+",
+ "| a | b | c | year | day |",
+ "+---+----+----+------+-----+",
+ "| 0 | -2 | 10 | 2021 | 26 |",
+ "| 1 | -1 | 11 | 2021 | 26 |",
+ "| 2 | 0 | 12 | 2021 | 26 |",
+ "+---+----+----+------+-----+",
+ ];
+ crate::assert_batches_eq!(expected, &[projected_batch]);
Review comment:
I think adding a test here that you can feed a second batch through (that is larger than the first batch size) and get the expected output would be valuable -- that is an important edge case
##########
File path: datafusion/tests/path_partition.rs
##########
@@ -0,0 +1,392 @@
+// 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.
+
+//! Test queries on partitioned datasets
+
+use std::{fs, io, sync::Arc};
+
+use async_trait::async_trait;
+use datafusion::{
+ assert_batches_sorted_eq,
+ datasource::{
+ file_format::{csv::CsvFormat, parquet::ParquetFormat},
+ listing::{ListingOptions, ListingTable},
+ object_store::{
+ local::LocalFileSystem, FileMeta, FileMetaStream, ListEntryStream,
+ ObjectReader, ObjectStore, SizedFile,
+ },
+ },
+ error::{DataFusionError, Result},
+ physical_plan::ColumnStatistics,
+ prelude::ExecutionContext,
+ test_util::{arrow_test_data, parquet_test_data},
+};
+use futures::{stream, StreamExt};
+
+mod common;
+
+#[tokio::test]
+async fn csv_filter_with_file_col() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+
+ register_partitioned_aggregate_csv(
+ &mut ctx,
+ &[
+ "mytable/date=2021-10-27/file.csv",
+ "mytable/date=2021-10-28/file.csv",
+ ],
+ &["date"],
+ "mytable",
+ );
+
+ let result = ctx
+ .sql("SELECT c1, c2 FROM t WHERE date='2021-10-27' and date!=c1 LIMIT 5")
+ .await?
+ .collect()
+ .await?;
+
+ let expected = vec![
+ "+----+----+",
+ "| c1 | c2 |",
+ "+----+----+",
+ "| a | 1 |",
+ "| b | 1 |",
+ "| b | 5 |",
+ "| c | 2 |",
+ "| d | 5 |",
+ "+----+----+",
+ ];
+ assert_batches_sorted_eq!(expected, &result);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn csv_projection_on_partition() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+
+ register_partitioned_aggregate_csv(
+ &mut ctx,
+ &[
+ "mytable/date=2021-10-27/file.csv",
+ "mytable/date=2021-10-28/file.csv",
+ ],
+ &["date"],
+ "mytable",
+ );
+
+ let result = ctx
+ .sql("SELECT c1, date FROM t WHERE date='2021-10-27' LIMIT 5")
+ .await?
+ .collect()
+ .await?;
+
+ let expected = vec![
+ "+----+------------+",
+ "| c1 | date |",
+ "+----+------------+",
+ "| a | 2021-10-27 |",
+ "| b | 2021-10-27 |",
+ "| b | 2021-10-27 |",
+ "| c | 2021-10-27 |",
+ "| d | 2021-10-27 |",
+ "+----+------------+",
+ ];
+ assert_batches_sorted_eq!(expected, &result);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn csv_grouping_by_partition() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+
+ register_partitioned_aggregate_csv(
+ &mut ctx,
+ &[
+ "mytable/date=2021-10-26/file.csv",
+ "mytable/date=2021-10-27/file.csv",
+ "mytable/date=2021-10-28/file.csv",
+ ],
+ &["date"],
+ "mytable",
+ );
+
+ let result = ctx
+ .sql("SELECT date, count(*), count(distinct(c1)) FROM t WHERE date<='2021-10-27' GROUP BY date")
+ .await?
+ .collect()
+ .await?;
+
+ let expected = vec![
+ "+------------+-----------------+----------------------+",
+ "| date | COUNT(UInt8(1)) | COUNT(DISTINCT t.c1) |",
+ "+------------+-----------------+----------------------+",
+ "| 2021-10-26 | 100 | 5 |",
+ "| 2021-10-27 | 100 | 5 |",
+ "+------------+-----------------+----------------------+",
+ ];
+ assert_batches_sorted_eq!(expected, &result);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn parquet_multiple_partitions() -> Result<()> {
+ let mut ctx = ExecutionContext::new();
+
+ register_partitioned_alltypes_parquet(
+ &mut ctx,
+ &[
+ "year=2021/month=09/day=09/file.parquet",
+ "year=2021/month=10/day=09/file.parquet",
+ "year=2021/month=10/day=28/file.parquet",
+ ],
+ &["year", "month", "day"],
+ "",
+ "alltypes_plain.parquet",
+ )
+ .await;
+
+ let result = ctx
Review comment:
this is very cool
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -162,58 +174,76 @@ impl TableProvider for ListingTable {
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
- // TODO object_store_registry should be provided as param here
- let (partitioned_file_lists, statistics) = self
- .list_files_for_scan(
- Arc::clone(&self.object_store),
- &self.path,
- filters,
- limit,
- )
- .await?;
+ let (partitioned_file_lists, statistics) =
+ self.list_files_for_scan(filters, limit).await?;
+
+ // if no files need to be read, return an `EmptyExec`
+ if partitioned_file_lists.is_empty() {
+ let schema = self.schema();
+ let projected_schema = match &projection {
+ None => schema,
+ Some(p) => Arc::new(Schema::new(
+ p.iter().map(|i| schema.field(*i).clone()).collect(),
+ )),
+ };
+ return Ok(Arc::new(EmptyExec::new(false, projected_schema)));
+ }
+
// create the execution plan
self.options
.format
- .create_physical_plan(PhysicalPlanConfig {
- object_store: Arc::clone(&self.object_store),
- schema: self.schema(),
- files: partitioned_file_lists,
- statistics,
- projection: projection.clone(),
- batch_size,
- filters: filters.to_vec(),
- limit,
- })
+ .create_physical_plan(
+ PhysicalPlanConfig {
+ object_store: Arc::clone(&self.object_store),
+ file_schema: Arc::clone(&self.file_schema),
+ file_groups: partitioned_file_lists,
+ statistics,
+ projection: projection.clone(),
+ batch_size,
+ limit,
+ table_partition_cols: self.options.table_partition_cols.clone(),
+ },
+ filters,
+ )
.await
}
fn supports_filter_pushdown(
&self,
- _filter: &Expr,
+ filter: &Expr,
) -> Result<TableProviderFilterPushDown> {
- Ok(TableProviderFilterPushDown::Inexact)
+ if expr_applicable_for_cols(&self.options.table_partition_cols, filter) {
Review comment:
❤️
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
Review comment:
Given that the partition columns come directly from a file system path typically, having them always be a `Utf8` type makes sense -- using dictionary encoding to make them more efficient is 👍
On other words I am not at all sure any other type for partitioning columns makes sense and so we could perhaps remove these comments
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+ match expr {
+ // leaf
+ Expr::Literal(_) => true,
+ // TODO how to handle qualified / unqualified names?
+ Expr::Column(logical_plan::Column { ref name, .. }) => col_names.contains(name),
+ // unary
+ Expr::Alias(child, _)
+ | Expr::Not(child)
+ | Expr::IsNotNull(child)
+ | Expr::IsNull(child)
+ | Expr::Negative(child)
+ | Expr::Cast { expr: child, .. }
+ | Expr::TryCast { expr: child, .. } => expr_applicable_for_cols(col_names, child),
+ // binary
+ Expr::BinaryExpr {
+ ref left,
+ ref right,
+ ..
+ } => {
+ expr_applicable_for_cols(col_names, left)
+ && expr_applicable_for_cols(col_names, right)
+ }
+ // ternary
+ Expr::Between {
+ expr: item,
+ low,
+ high,
+ ..
+ } => {
+ expr_applicable_for_cols(col_names, item)
+ && expr_applicable_for_cols(col_names, low)
+ && expr_applicable_for_cols(col_names, high)
+ }
+ // variadic
+ Expr::ScalarFunction { fun, args } => match fun.volatility() {
+ Volatility::Immutable => args
+ .iter()
+ .all(|arg| expr_applicable_for_cols(col_names, arg)),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable => false,
+ Volatility::Volatile => false,
+ },
+ Expr::ScalarUDF { fun, args } => match fun.signature.volatility {
+ Volatility::Immutable => args
+ .iter()
+ .all(|arg| expr_applicable_for_cols(col_names, arg)),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable => false,
+ Volatility::Volatile => false,
+ },
+ Expr::InList {
+ expr: item, list, ..
+ } => {
+ expr_applicable_for_cols(col_names, item)
+ && list.iter().all(|e| expr_applicable_for_cols(col_names, e))
+ }
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ let expr_constant = expr
+ .as_ref()
+ .map(|e| expr_applicable_for_cols(col_names, e))
+ .unwrap_or(true);
+ let else_constant = else_expr
+ .as_ref()
+ .map(|e| expr_applicable_for_cols(col_names, e))
+ .unwrap_or(true);
+ let when_then_constant = when_then_expr.iter().all(|(w, th)| {
+ expr_applicable_for_cols(col_names, w)
+ && expr_applicable_for_cols(col_names, th)
+ });
+ expr_constant && else_constant && when_then_constant
+ }
+ // TODO other expressions are not handled yet:
+ // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in some edge cases
+ // - Can `Wildcard` be considered as a `Literal`?
+ // - ScalarVariable could be `applicable`, but that would require access to the context
+ _ => false,
+ }
+}
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let chunk_size = (partitioned_files.len() + n - 1) / n;
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// that belong to irrelevant partitions using `filters` expressions.
+/// `filters` might contain expressions that can be resolved only at the
+/// file level (e.g. Parquet row group pruning).
+///
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_cols: &[String],
+) -> Result<PartitionedFileStream> {
+ // if no partition col => simply list all the files
+ if table_partition_cols.is_empty() {
+ return Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ partition_values: vec![],
+ file_meta: f?,
+ })
+ }),
+ ));
+ }
+
+ let applicable_filters: Vec<_> = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_cols, f))
+ .collect();
+ let stream_path = table_path.to_owned();
+ if applicable_filters.is_empty() {
+ // parse the partition values while listing all the files
+ // TODO we might avoid parsing the partition values if they are not used in any projection
Review comment:
I don't think it will matter -- the time to fetch the listing from object store will far dominate the parsing cost I think
##########
File path: datafusion/src/datasource/listing/mod.rs
##########
@@ -0,0 +1,24 @@
+// 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.
+
+//! A table that uses the `ObjectStore` listing capability
+//! to get the list of files to process.
+
+mod helpers;
+mod table;
Review comment:
(FWIW this is the kind of change that is an improvement but generates a large diff and makes it harder to understand what changed in the logic -- if you are looking for ways to make your PRs smaller, one way might be to make a first PR that just split up the existing code into two modules) -- then while that is being reviewed you can rebase your WIP branch on top of it and carry on
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
Review comment:
I recommend putting the context and design rationale from this comment into the code as a doc comment
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-946583530
> Thanks! You can take a quick look, but this will not be ready before 1 or 2 days... If you want, I can ping you once it's in a more reviewable state 😃
Sounds good. Thank you 🙏
--
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
[GitHub] [arrow-datafusion] rdettai commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-946491974
Thanks! You can take a quick look, but this will not be ready before 1 or 2 days... If you want, I can ping you once it's in a more reviewable state 😃
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738766844
##########
File path: ballista/rust/core/src/serde/physical_plan/to_proto.rs
##########
@@ -244,90 +246,29 @@ impl TryInto<protobuf::PhysicalPlanNode> for Arc<dyn ExecutionPlan> {
))),
})
} else if let Some(exec) = plan.downcast_ref::<CsvExec>() {
- let file_groups = exec
- .file_groups()
- .iter()
- .map(|p| p.as_slice().into())
- .collect();
Ok(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::CsvScan(
protobuf::CsvScanExecNode {
- file_groups,
- statistics: Some((&exec.statistics()).into()),
- limit: exec
- .limit()
- .map(|l| protobuf::ScanLimit { limit: l as u32 }),
- projection: exec
- .projection()
- .as_ref()
- .ok_or_else(|| {
- BallistaError::General(
- "projection in CsvExec does not exist.".to_owned(),
- )
- })?
- .iter()
- .map(|n| *n as u32)
- .collect(),
- schema: Some(exec.file_schema().as_ref().into()),
+ base_conf: Some(exec.base_config().try_into()?),
has_header: exec.has_header(),
delimiter: byte_to_string(exec.delimiter())?,
- batch_size: exec.batch_size() as u32,
},
)),
})
} else if let Some(exec) = plan.downcast_ref::<ParquetExec>() {
- let file_groups = exec
- .file_groups()
- .iter()
- .map(|p| p.as_slice().into())
- .collect();
-
Ok(protobuf::PhysicalPlanNode {
physical_plan_type: Some(PhysicalPlanType::ParquetScan(
protobuf::ParquetScanExecNode {
- file_groups,
- statistics: Some((&exec.statistics()).into()),
- limit: exec
- .limit()
- .map(|l| protobuf::ScanLimit { limit: l as u32 }),
- schema: Some(exec.schema().as_ref().into()),
- projection: exec
- .projection()
- .as_ref()
- .iter()
- .map(|n| *n as u32)
- .collect(),
- batch_size: exec.batch_size() as u32,
+ base_conf: Some(exec.base_config().try_into()?),
+ // TODO serialize predicates
Review comment:
This was already there.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735338913
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
Review comment:
I don't think we are silencing any error here. Typos in the query were already caught (or should already be caught) upstream in the SQL parser with checks like https://github.com/apache/arrow-datafusion/blob/fe1a9e2c55392b934c85098430f78a26ef71380e/datafusion/src/sql/planner.rs#L769-L775
Here `expr_applicable_for_cols` is different, it only checks whether a given expression can be resolved using the partitioning dimensions only. For instance, if the table partitioning is of type `mytable/date=2021-01-01/file.csv` with a file schema of the form `Schema([Field("a",int), Field("b",string)])`. A filter that contains `WHERE b='hello'` or `WHERE b=date` is perfectly valid, but the filter should not be kept in the list of `applicable_filters` because they cannot be resolved with the partitioning column only.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740282520
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -613,33 +614,28 @@ message ScanLimit {
uint32 limit = 1;
}
-message ParquetScanExecNode {
+message FileScanExecConf {
repeated FileGroup file_groups = 1;
Schema schema = 2;
- uint32 batch_size = 4;
- repeated uint32 projection = 6;
- ScanLimit limit = 7;
- Statistics statistics = 8;
+ uint32 batch_size = 3;
Review comment:
As long as you don't have DataFusion nodes with different versions, it should be ok!
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956566766
As this has been outstanding for a long time and is a fairly large change, the potential for conflict is large -- I am going to merge it in now and hopefully we can keep iterating in future PRs. Thanks again @rdettai @Jimexist @houqp and @Dandandan -- 🚀
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735326567
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
+/// - one column for with the original path named `_df_part_file_path_`
+/// - one column for with the last modified date named `_df_part_file_modified_`
+/// - ... one column by partition ...
+/// Note: For the last modified date, this looses precisions higher than millisecond.
+fn paths_to_batch(
+ table_partition_dims: &[String],
+ table_path: &str,
+ metas: &[FileMeta],
+) -> Result<RecordBatch> {
+ let mut key_builder = StringBuilder::new(metas.len());
+ let mut length_builder = UInt64Builder::new(metas.len());
+ let mut modified_builder = Date64Builder::new(metas.len());
+ let mut partition_builders = table_partition_dims
+ .iter()
+ .map(|_| StringBuilder::new(metas.len()))
+ .collect::<Vec<_>>();
+ for file_meta in metas {
+ if let Some(partition_values) =
+ parse_partitions_for_path(table_path, file_meta.path(), table_partition_dims)
+ {
+ key_builder.append_value(file_meta.path())?;
+ length_builder.append_value(file_meta.size())?;
+ match file_meta.last_modified {
+ Some(lm) => modified_builder.append_value(lm.timestamp_millis())?,
+ None => modified_builder.append_null()?,
+ }
+ for (i, part_val) in partition_values.iter().enumerate() {
+ partition_builders[i].append_value(part_val)?;
+ }
+ } else {
+ debug!("No partitioning for path {}", file_meta.path());
+ }
+ }
+
+ // finish all builders
+ let mut col_arrays: Vec<ArrayRef> = vec![
+ ArrayBuilder::finish(&mut key_builder),
+ ArrayBuilder::finish(&mut length_builder),
+ ArrayBuilder::finish(&mut modified_builder),
+ ];
+ for mut partition_builder in partition_builders {
+ col_arrays.push(ArrayBuilder::finish(&mut partition_builder));
+ }
+
+ // put the schema together
+ let mut fields = vec![
+ Field::new(FILE_SIZE_COLUMN_NAME, DataType::Utf8, false),
+ Field::new(FILE_PATH_COLUMN_NAME, DataType::UInt64, false),
Review comment:
yes thanks
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739011378
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
+ /// An Arrow buffer initialized to zeros that represents the key array of all partition
+ /// columns (partition columns are materialized by dictionary arrays with only one
+ /// value in the dictionary, thus all the keys are equal to zero).
+ key_buffer_cache: Option<Buffer>,
+ /// Mapping between the indexes in the list of partition columns and the target
+ /// schema. Sorted by index in the target schema so that we can iterate on it to
+ /// insert the partition columns in the target record batch.
+ projected_partition_indexes: Vec<(usize, usize)>,
+ /// The schema of the table once the projection was applied.
+ projected_schema: SchemaRef,
+}
+
+impl PartitionColumnProjector {
+ // Create a projector to insert the partitioning columns into batches read from files
+ // - projected_schema: the target schema with both file and partitioning columns
+ // - table_partition_cols: all the partitioning column names
+ fn new(projected_schema: SchemaRef, table_partition_cols: &[String]) -> Self {
+ let mut idx_map = HashMap::new();
+ for (partition_idx, partition_name) in table_partition_cols.iter().enumerate() {
+ if let Ok(schema_idx) = projected_schema.index_of(partition_name) {
+ idx_map.insert(partition_idx, schema_idx);
+ }
+ }
+
+ let mut projected_partition_indexes: Vec<_> = idx_map.into_iter().collect();
+ projected_partition_indexes.sort_by(|(_, a), (_, b)| a.cmp(b));
+
+ Self {
+ projected_partition_indexes,
+ key_buffer_cache: None,
+ projected_schema,
+ }
+ }
+
+ // Transform the batch read from the fileby inserting the partitioning columns
+ // to the right positions as deduced from `projected_schema`
+ // - file_batch: batch read from the file, with internal projection applied
+ // - partition_values: the list of partition values, one for each partition column
+ fn project(
+ &mut self,
+ file_batch: RecordBatch,
+ partition_values: &[ScalarValue],
+ ) -> ArrowResult<RecordBatch> {
+ let expected_cols =
+ self.projected_schema.fields().len() - self.projected_partition_indexes.len();
+
+ if file_batch.columns().len() != expected_cols {
+ return Err(ArrowError::SchemaError(format!(
+ "Unexpected batch schema from file, expected {} cols but got {}",
+ expected_cols,
+ file_batch.columns().len()
+ )));
+ }
+
+ let mut cols = file_batch.columns().to_vec();
+ for &(pidx, sidx) in &self.projected_partition_indexes {
+ cols.insert(
+ sidx,
+ create_dict_array(
+ &mut self.key_buffer_cache,
+ &partition_values[pidx],
+ file_batch.num_rows(),
+ ),
+ )
+ }
+ RecordBatch::try_new(Arc::clone(&self.projected_schema), cols)
+ }
+}
+
+fn create_dict_array(
+ key_buffer_cache: &mut Option<Buffer>,
+ val: &ScalarValue,
+ len: usize,
+) -> ArrayRef {
+ // build value dictionary
+ let dict_vals = val.to_array();
+
+ // build keys array
+ let sliced_key_buffer = match key_buffer_cache {
+ Some(buf) if buf.len() >= len => buf.slice(buf.len() - len),
+ _ => {
+ let mut key_buffer_builder = UInt8BufferBuilder::new(len);
+ key_buffer_builder.advance(len); // keys are all 0
+ key_buffer_cache.insert(key_buffer_builder.finish()).clone()
+ }
+ };
+
+ // create data type
+ let data_type =
+ DataType::Dictionary(Box::new(DataType::UInt8), Box::new(val.get_datatype()));
+
+ debug_assert_eq!(data_type, *DEFAULT_PARTITION_COLUMN_DATATYPE);
+
+ // assemble pieces together
+ let mut builder = ArrayData::builder(data_type)
+ .len(len)
+ .add_buffer(sliced_key_buffer);
+ builder = builder.add_child_data(dict_vals.data().clone());
+ Arc::new(DictionaryArray::<UInt8Type>::from(builder.build().unwrap()))
+}
+
+#[cfg(test)]
+mod tests {
+ use crate::test::{
+ aggr_test_schema, build_table_i32, columns, object_store::TestObjectStore,
+ };
+
+ use super::*;
+
+ #[test]
+ fn physical_plan_config_no_projection() {
+ let file_schema = aggr_test_schema();
+ let conf = config_for_projection(
+ Arc::clone(&file_schema),
+ None,
+ Statistics::default(),
+ vec!["date".to_owned()],
+ );
+
+ let (proj_schema, proj_statistics) = conf.project();
+ assert_eq!(proj_schema.fields().len(), file_schema.fields().len() + 1);
+ assert_eq!(
+ proj_schema.field(file_schema.fields().len()).name(),
+ "date",
+ "partition columns are the last columns"
+ );
+ assert_eq!(
+ proj_statistics
+ .column_statistics
+ .expect("projection creates column statistics")
+ .len(),
+ file_schema.fields().len() + 1
+ );
+ // TODO implement tests for partition column statistics once implemented
+
+ let col_names = conf.projected_file_column_names();
+ assert_eq!(col_names, None);
+
+ let col_indices = conf.file_column_projection_indices();
+ assert_eq!(col_indices, None);
+ }
+
+ #[test]
+ fn physical_plan_config_with_projection() {
+ let file_schema = aggr_test_schema();
+ let conf = config_for_projection(
+ Arc::clone(&file_schema),
+ Some(vec![file_schema.fields().len(), 0]),
+ Statistics {
+ num_rows: Some(10),
+ // assign the column index to distinct_count to help assert
+ // the source statistic after the projection
+ column_statistics: Some(
+ (0..file_schema.fields().len())
+ .map(|i| ColumnStatistics {
+ distinct_count: Some(i),
+ ..Default::default()
+ })
+ .collect(),
+ ),
+ ..Default::default()
+ },
+ vec!["date".to_owned()],
+ );
+
+ let (proj_schema, proj_statistics) = conf.project();
+ assert_eq!(
+ columns(&proj_schema),
+ vec!["date".to_owned(), "c1".to_owned()]
+ );
+ let proj_stat_cols = proj_statistics
+ .column_statistics
+ .expect("projection creates column statistics");
+ assert_eq!(proj_stat_cols.len(), 2);
+ // TODO implement tests for proj_stat_cols[0] once partition column
+ // statistics are implemented
+ assert_eq!(proj_stat_cols[1].distinct_count, Some(0));
+
+ let col_names = conf.projected_file_column_names();
+ assert_eq!(col_names, Some(vec!["c1".to_owned()]));
+
+ let col_indices = conf.file_column_projection_indices();
+ assert_eq!(col_indices, Some(vec![0]));
+ }
+
+ #[test]
+ fn partition_column_projector() {
+ let file_batch = build_table_i32(
+ ("a", &vec![0, 1, 2]),
+ ("b", &vec![-2, -1, 0]),
+ ("c", &vec![10, 11, 12]),
+ );
+ let partition_cols =
+ vec!["year".to_owned(), "month".to_owned(), "day".to_owned()];
+ // create a projected schema
+ let conf = config_for_projection(
+ file_batch.schema(),
+ // keep all cols from file and 2 from partitioning
+ Some(vec![
+ 0,
+ 1,
+ 2,
+ file_batch.schema().fields().len(),
+ file_batch.schema().fields().len() + 2,
+ ]),
+ Statistics::default(),
+ partition_cols.clone(),
+ );
+ let (proj_schema, _) = conf.project();
+ // created a projector for that projected schema
+ let mut proj = PartitionColumnProjector::new(proj_schema, &partition_cols);
+ let projected_batch = proj
+ .project(
+ // file_batch is ok here because we kept all the file cols in the projection
+ file_batch,
+ &[
+ ScalarValue::Utf8(Some("2021".to_owned())),
+ ScalarValue::Utf8(Some("10".to_owned())),
+ ScalarValue::Utf8(Some("26".to_owned())),
+ ],
+ )
+ .expect("Projection of partition columns into record batch failed");
+
+ let expected = vec![
+ "+---+----+----+------+-----+",
+ "| a | b | c | year | day |",
+ "+---+----+----+------+-----+",
+ "| 0 | -2 | 10 | 2021 | 26 |",
+ "| 1 | -1 | 11 | 2021 | 26 |",
+ "| 2 | 0 | 12 | 2021 | 26 |",
+ "+---+----+----+------+-----+",
+ ];
+ crate::assert_batches_eq!(expected, &[projected_batch]);
Review comment:
doing it for both larger and smaller batches
--
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
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
houqp commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735316543
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
Review comment:
oh, nvm then. I am fine with keeping it as is for easier review :)
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735325807
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
Review comment:
datafusion in datafusion is practical, because we have the filters as logical expressions 🙂. But you are right, it might be a bit cheaper to convert the logical expressions into physical ones manually and put together a `MemoryExec` and a `FilterExec` ourselves.
In all cases, the solution proposed here is not the final one because we would need a way to avoid materializing the entire list of files (as described in the `// TODO`). When doing that improvement, it will be interesting to re-assess if it makes sense using the high level Datafusion API or not 😉
--
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
[GitHub] [arrow-datafusion] Dandandan commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Dandandan commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735326166
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
Review comment:
Wouldn't a larger batch size also be fine (e.g. 1K or the now default batch size (8K))?
ATM we don't gain anything from having it in multiple batches as we wait on the entire results to be returned?
For a streaming `MemTable` we could revisit this and collect into smaller chunks? Maybe the page size of the backend could be retrieved from the store (e.g. S3 has 1000) or the store returns the files in chunks itself?
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-951268260
I plan to review this PR tomorrow (when I am fresher and can git it the look it deserves)
--
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
[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739140420
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,721 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr, ExpressionVisitor, Recursion},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// The `ExpressionVisitor` for `expr_applicable_for_cols`. Walks the tree to
+/// validate that the given expression is applicable with only the `col_names`
+/// set of columns.
+struct ApplicabilityVisitor<'a> {
+ col_names: &'a [String],
+ is_applicable: &'a mut bool,
+}
+
+impl ApplicabilityVisitor<'_> {
+ fn visit_volatility(self, volatility: Volatility) -> Recursion<Self> {
+ match volatility {
+ Volatility::Immutable => Recursion::Continue(self),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable | Volatility::Volatile => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ }
+ }
+}
+
+impl ExpressionVisitor for ApplicabilityVisitor<'_> {
+ fn pre_visit(self, expr: &Expr) -> Result<Recursion<Self>> {
Review comment:
👍
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
Review comment:
I personally think it is now easier to read now too as the logic is not interspersed with the recursion, but I realize that is in the eye of the beholder :)
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
Review comment:
I see -- so we are envisioning some how allowing users of this code to specify the type in the partitioning somehow (and provide their own code to determine the partition values). That makes sense
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,721 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr, ExpressionVisitor, Recursion},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// The `ExpressionVisitor` for `expr_applicable_for_cols`. Walks the tree to
+/// validate that the given expression is applicable with only the `col_names`
+/// set of columns.
+struct ApplicabilityVisitor<'a> {
+ col_names: &'a [String],
+ is_applicable: &'a mut bool,
+}
+
+impl ApplicabilityVisitor<'_> {
+ fn visit_volatility(self, volatility: Volatility) -> Recursion<Self> {
+ match volatility {
+ Volatility::Immutable => Recursion::Continue(self),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable | Volatility::Volatile => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ }
+ }
+}
+
+impl ExpressionVisitor for ApplicabilityVisitor<'_> {
+ fn pre_visit(self, expr: &Expr) -> Result<Recursion<Self>> {
+ let rec = match expr {
+ Expr::Column(logical_plan::Column { ref name, .. }) => {
+ *self.is_applicable &= self.col_names.contains(name);
+ Recursion::Stop(self) // leaf node anyway
+ }
+ Expr::Literal(_)
+ | Expr::Alias(_, _)
+ | Expr::ScalarVariable(_)
+ | Expr::Not(_)
+ | Expr::IsNotNull(_)
+ | Expr::IsNull(_)
+ | Expr::Negative(_)
+ | Expr::Cast { .. }
+ | Expr::TryCast { .. }
+ | Expr::BinaryExpr { .. }
+ | Expr::Between { .. }
+ | Expr::InList { .. }
+ | Expr::Case { .. } => Recursion::Continue(self),
+
+ Expr::ScalarFunction { fun, .. } => self.visit_volatility(fun.volatility()),
+ Expr::ScalarUDF { fun, .. } => {
+ self.visit_volatility(fun.signature.volatility)
+ }
+
+ // TODO other expressions are not handled yet:
+ // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in some edge cases
+ // - Can `Wildcard` be considered as a `Literal`?
+ // - ScalarVariable could be `applicable`, but that would require access to the context
+ Expr::AggregateUDF { .. }
+ | Expr::AggregateFunction { .. }
+ | Expr::Sort { .. }
+ | Expr::WindowFunction { .. }
+ | Expr::Wildcard => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ };
+ Ok(rec)
+ }
+}
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+ let mut is_applicable = true;
+ expr.accept(ApplicabilityVisitor {
+ col_names,
+ is_applicable: &mut is_applicable,
+ })
+ .unwrap();
+ is_applicable
+}
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let chunk_size = (partitioned_files.len() + n - 1) / n;
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// that belong to irrelevant partitions using `filters` expressions.
+/// `filters` might contain expressions that can be resolved only at the
+/// file level (e.g. Parquet row group pruning).
+///
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_cols: &[String],
+) -> Result<PartitionedFileStream> {
+ // if no partition col => simply list all the files
+ if table_partition_cols.is_empty() {
+ return Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ partition_values: vec![],
+ file_meta: f?,
+ })
+ }),
+ ));
+ }
+
+ let applicable_filters: Vec<_> = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_cols, f))
+ .collect();
+ let stream_path = table_path.to_owned();
+ if applicable_filters.is_empty() {
+ // Parse the partition values while listing all the files
+ // Note: We might avoid parsing the partition values if they are not used in any projection,
+ // but the cost of parsing will likely be far dominated by the time to fetch the listing from
+ // the object store.
+ let table_partition_cols_stream = table_partition_cols.to_vec();
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .filter_map(move |f| {
+ let stream_path = stream_path.clone();
+ let table_partition_cols_stream = table_partition_cols_stream.clone();
+ async move {
+ let file_meta = match f {
+ Ok(fm) => fm,
+ Err(err) => return Some(Err(err)),
+ };
+ let parsed_path = parse_partitions_for_path(
+ &stream_path,
+ file_meta.path(),
+ &table_partition_cols_stream,
+ )
+ .map(|p| {
+ p.iter()
+ .map(|&pn| ScalarValue::Utf8(Some(pn.to_owned())))
+ .collect()
+ });
+
+ parsed_path.map(|partition_values| {
+ Ok(PartitionedFile {
+ partition_values,
+ file_meta,
+ })
+ })
+ }
+ }),
+ ))
+ } else {
+ // parse the partition values and serde them as a RecordBatch to filter them
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ // TODO we set an arbitrary high batch size here, it does not matter as we list
+ // all the files anyway. This number will need to be adjusted according to the object
+ // store if we switch to a streaming-stlye pruning of the files. For instance S3 lists
+ // 1000 items at a time so batches of 1000 would be ideal with S3 as store.
+ .chunks(1024)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| paths_to_batch(table_partition_cols, &stream_path, &metas?))
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
Review comment:
This is a clever way to apply filtering -- convert the paths to batches and then run the evaluation on the batches, and then turn it back to paths 👍
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,721 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr, ExpressionVisitor, Recursion},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// The `ExpressionVisitor` for `expr_applicable_for_cols`. Walks the tree to
+/// validate that the given expression is applicable with only the `col_names`
+/// set of columns.
+struct ApplicabilityVisitor<'a> {
+ col_names: &'a [String],
+ is_applicable: &'a mut bool,
+}
+
+impl ApplicabilityVisitor<'_> {
+ fn visit_volatility(self, volatility: Volatility) -> Recursion<Self> {
+ match volatility {
+ Volatility::Immutable => Recursion::Continue(self),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable | Volatility::Volatile => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ }
+ }
+}
+
+impl ExpressionVisitor for ApplicabilityVisitor<'_> {
+ fn pre_visit(self, expr: &Expr) -> Result<Recursion<Self>> {
+ let rec = match expr {
+ Expr::Column(logical_plan::Column { ref name, .. }) => {
+ *self.is_applicable &= self.col_names.contains(name);
+ Recursion::Stop(self) // leaf node anyway
+ }
+ Expr::Literal(_)
+ | Expr::Alias(_, _)
+ | Expr::ScalarVariable(_)
+ | Expr::Not(_)
+ | Expr::IsNotNull(_)
+ | Expr::IsNull(_)
+ | Expr::Negative(_)
+ | Expr::Cast { .. }
+ | Expr::TryCast { .. }
+ | Expr::BinaryExpr { .. }
+ | Expr::Between { .. }
+ | Expr::InList { .. }
+ | Expr::Case { .. } => Recursion::Continue(self),
+
+ Expr::ScalarFunction { fun, .. } => self.visit_volatility(fun.volatility()),
+ Expr::ScalarUDF { fun, .. } => {
+ self.visit_volatility(fun.signature.volatility)
+ }
+
+ // TODO other expressions are not handled yet:
+ // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in some edge cases
+ // - Can `Wildcard` be considered as a `Literal`?
+ // - ScalarVariable could be `applicable`, but that would require access to the context
+ Expr::AggregateUDF { .. }
+ | Expr::AggregateFunction { .. }
+ | Expr::Sort { .. }
+ | Expr::WindowFunction { .. }
+ | Expr::Wildcard => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ };
+ Ok(rec)
+ }
+}
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+ let mut is_applicable = true;
+ expr.accept(ApplicabilityVisitor {
+ col_names,
+ is_applicable: &mut is_applicable,
+ })
+ .unwrap();
+ is_applicable
+}
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let chunk_size = (partitioned_files.len() + n - 1) / n;
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// that belong to irrelevant partitions using `filters` expressions.
+/// `filters` might contain expressions that can be resolved only at the
+/// file level (e.g. Parquet row group pruning).
+///
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_cols: &[String],
+) -> Result<PartitionedFileStream> {
+ // if no partition col => simply list all the files
+ if table_partition_cols.is_empty() {
+ return Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ partition_values: vec![],
+ file_meta: f?,
+ })
+ }),
+ ));
+ }
+
+ let applicable_filters: Vec<_> = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_cols, f))
+ .collect();
+ let stream_path = table_path.to_owned();
+ if applicable_filters.is_empty() {
+ // Parse the partition values while listing all the files
+ // Note: We might avoid parsing the partition values if they are not used in any projection,
+ // but the cost of parsing will likely be far dominated by the time to fetch the listing from
+ // the object store.
+ let table_partition_cols_stream = table_partition_cols.to_vec();
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .filter_map(move |f| {
+ let stream_path = stream_path.clone();
+ let table_partition_cols_stream = table_partition_cols_stream.clone();
+ async move {
+ let file_meta = match f {
+ Ok(fm) => fm,
+ Err(err) => return Some(Err(err)),
+ };
+ let parsed_path = parse_partitions_for_path(
+ &stream_path,
+ file_meta.path(),
+ &table_partition_cols_stream,
+ )
+ .map(|p| {
+ p.iter()
+ .map(|&pn| ScalarValue::Utf8(Some(pn.to_owned())))
+ .collect()
+ });
+
+ parsed_path.map(|partition_values| {
+ Ok(PartitionedFile {
+ partition_values,
+ file_meta,
+ })
+ })
+ }
+ }),
+ ))
+ } else {
+ // parse the partition values and serde them as a RecordBatch to filter them
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ // TODO we set an arbitrary high batch size here, it does not matter as we list
+ // all the files anyway. This number will need to be adjusted according to the object
+ // store if we switch to a streaming-stlye pruning of the files. For instance S3 lists
+ // 1000 items at a time so batches of 1000 would be ideal with S3 as store.
+ .chunks(1024)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| paths_to_batch(table_partition_cols, &stream_path, &metas?))
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
+/// - one column for with the original path named `_df_part_file_path_`
+/// - one column for with the last modified date named `_df_part_file_modified_`
+/// - ... one column by partition ...
+///
+/// Note: For the last modified date, this looses precisions higher than millisecond.
+fn paths_to_batch(
+ table_partition_cols: &[String],
+ table_path: &str,
+ metas: &[FileMeta],
+) -> Result<RecordBatch> {
+ let mut key_builder = StringBuilder::new(metas.len());
+ let mut length_builder = UInt64Builder::new(metas.len());
+ let mut modified_builder = Date64Builder::new(metas.len());
+ let mut partition_builders = table_partition_cols
+ .iter()
+ .map(|_| StringBuilder::new(metas.len()))
+ .collect::<Vec<_>>();
+ for file_meta in metas {
+ if let Some(partition_values) =
+ parse_partitions_for_path(table_path, file_meta.path(), table_partition_cols)
+ {
+ key_builder.append_value(file_meta.path())?;
+ length_builder.append_value(file_meta.size())?;
+ match file_meta.last_modified {
+ Some(lm) => modified_builder.append_value(lm.timestamp_millis())?,
+ None => modified_builder.append_null()?,
+ }
+ for (i, part_val) in partition_values.iter().enumerate() {
+ partition_builders[i].append_value(part_val)?;
+ }
+ } else {
+ debug!("No partitioning for path {}", file_meta.path());
+ }
+ }
+
+ // finish all builders
+ let mut col_arrays: Vec<ArrayRef> = vec![
+ ArrayBuilder::finish(&mut key_builder),
+ ArrayBuilder::finish(&mut length_builder),
+ ArrayBuilder::finish(&mut modified_builder),
+ ];
+ for mut partition_builder in partition_builders {
+ col_arrays.push(ArrayBuilder::finish(&mut partition_builder));
+ }
+
+ // put the schema together
+ let mut fields = vec![
+ Field::new(FILE_PATH_COLUMN_NAME, DataType::Utf8, false),
+ Field::new(FILE_SIZE_COLUMN_NAME, DataType::UInt64, false),
+ Field::new(FILE_MODIFIED_COLUMN_NAME, DataType::Date64, false),
+ ];
+ for pn in table_partition_cols {
+ fields.push(Field::new(pn, DataType::Utf8, false));
+ }
+
+ let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), col_arrays)?;
+ Ok(batch)
+}
+
+/// convert a set of record batches created by `paths_to_batch()` back to partitioned files.
+fn batches_to_paths(batches: &[RecordBatch]) -> Vec<PartitionedFile> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ let key_array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let length_array = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap();
+ let modified_array = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Date64Array>()
+ .unwrap();
+
+ (0..batch.num_rows()).map(move |row| PartitionedFile {
+ file_meta: FileMeta {
+ last_modified: match modified_array.is_null(row) {
+ false => Some(Utc.timestamp_millis(modified_array.value(row))),
+ true => None,
+ },
+ sized_file: SizedFile {
+ path: key_array.value(row).to_owned(),
+ size: length_array.value(row),
+ },
+ },
+ partition_values: (3..batch.columns().len())
+ .map(|col| {
+ ScalarValue::try_from_array(batch.column(col), row).unwrap()
+ })
+ .collect(),
+ })
+ })
+ .collect()
+}
+
+/// Extract the partition values for the given `file_path` (in the given `table_path`)
+/// associated to the partitions defined by `table_partition_cols`
+fn parse_partitions_for_path<'a>(
+ table_path: &str,
+ file_path: &'a str,
+ table_partition_cols: &[String],
+) -> Option<Vec<&'a str>> {
+ let subpath = file_path.strip_prefix(table_path)?;
+
+ // ignore whether table_path ended with "/" or not
+ let subpath = match subpath.strip_prefix('/') {
+ Some(subpath) => subpath,
+ None => subpath,
+ };
+
+ let mut part_values = vec![];
+ for (path, pn) in subpath.split('/').zip(table_partition_cols) {
+ match path.split_once('=') {
+ Some((name, val)) if name == pn => part_values.push(val),
+ _ => return None,
+ }
+ }
+ Some(part_values)
+}
+
+#[cfg(test)]
+mod tests {
+ use crate::{
+ logical_plan::{case, col, lit},
+ test::object_store::TestObjectStore,
+ };
+
+ use super::*;
+
+ #[test]
+ fn test_split_files() {
+ let new_partitioned_file = |path: &str| PartitionedFile::new(path.to_owned(), 10);
+ let files = vec![
+ new_partitioned_file("a"),
+ new_partitioned_file("b"),
+ new_partitioned_file("c"),
+ new_partitioned_file("d"),
+ new_partitioned_file("e"),
+ ];
+
+ let chunks = split_files(files.clone(), 1);
+ assert_eq!(1, chunks.len());
+ assert_eq!(5, chunks[0].len());
+
+ let chunks = split_files(files.clone(), 2);
+ assert_eq!(2, chunks.len());
+ assert_eq!(3, chunks[0].len());
+ assert_eq!(2, chunks[1].len());
+
+ let chunks = split_files(files.clone(), 5);
+ assert_eq!(5, chunks.len());
+ assert_eq!(1, chunks[0].len());
+ assert_eq!(1, chunks[1].len());
+ assert_eq!(1, chunks[2].len());
+ assert_eq!(1, chunks[3].len());
+ assert_eq!(1, chunks[4].len());
+
+ let chunks = split_files(files, 123);
+ assert_eq!(5, chunks.len());
+ assert_eq!(1, chunks[0].len());
+ assert_eq!(1, chunks[1].len());
+ assert_eq!(1, chunks[2].len());
+ assert_eq!(1, chunks[3].len());
+ assert_eq!(1, chunks[4].len());
+
+ let chunks = split_files(vec![], 2);
+ assert_eq!(0, chunks.len());
+ }
+
+ #[tokio::test]
+ async fn test_pruned_partition_list_empty() {
+ let store = TestObjectStore::new_arc(&[
+ ("tablepath/mypartition=val1/notparquetfile", 100),
+ ("tablepath/file.parquet", 100),
+ ]);
+ let filter = Expr::eq(col("mypartition"), lit("val1"));
+ let pruned = pruned_partition_list(
+ store.as_ref(),
+ "tablepath/",
+ &[filter],
+ ".parquet",
+ &[String::from("mypartition")],
+ )
+ .await
+ .expect("partition pruning failed")
+ .collect::<Vec<_>>()
+ .await;
+
+ assert_eq!(pruned.len(), 0);
+ }
+
+ #[tokio::test]
+ async fn test_pruned_partition_list() {
+ let store = TestObjectStore::new_arc(&[
+ ("tablepath/mypartition=val1/file.parquet", 100),
+ ("tablepath/mypartition=val2/file.parquet", 100),
+ ("tablepath/mypartition=val1/other=val3/file.parquet", 100),
+ ]);
+ let filter = Expr::eq(col("mypartition"), lit("val1"));
+ let pruned = pruned_partition_list(
+ store.as_ref(),
+ "tablepath/",
+ &[filter],
+ ".parquet",
+ &[String::from("mypartition")],
+ )
+ .await
+ .expect("partition pruning failed")
+ .collect::<Vec<_>>()
+ .await;
+
+ assert_eq!(pruned.len(), 2);
+ let f1 = pruned[0].as_ref().expect("first item not an error");
+ assert_eq!(
+ &f1.file_meta.sized_file.path,
+ "tablepath/mypartition=val1/file.parquet"
+ );
+ assert_eq!(
+ &f1.partition_values,
+ &[ScalarValue::Utf8(Some(String::from("val1"))),]
+ );
+ let f2 = pruned[1].as_ref().expect("second item not an error");
+ assert_eq!(
+ &f2.file_meta.sized_file.path,
+ "tablepath/mypartition=val1/other=val3/file.parquet"
+ );
+ assert_eq!(
+ &f2.partition_values,
+ &[ScalarValue::Utf8(Some(String::from("val1"))),]
+ );
+ }
+
+ #[tokio::test]
+ async fn test_pruned_partition_list_multi() {
+ let store = TestObjectStore::new_arc(&[
+ ("tablepath/part1=p1v1/file.parquet", 100),
+ ("tablepath/part1=p1v2/part2=p2v1/file1.parquet", 100),
+ ("tablepath/part1=p1v2/part2=p2v1/file2.parquet", 100),
+ ("tablepath/part1=p1v3/part2=p2v1/file2.parquet", 100),
+ ("tablepath/part1=p1v2/part2=p2v2/file2.parquet", 100),
+ ]);
+ let filter1 = Expr::eq(col("part1"), lit("p1v2"));
+ let filter2 = Expr::eq(col("part2"), lit("p2v1"));
+ // filter3 cannot be resolved at partition pruning
+ let filter3 = Expr::eq(col("part2"), col("other"));
+ let pruned = pruned_partition_list(
+ store.as_ref(),
+ "tablepath/",
+ &[filter1, filter2, filter3],
+ ".parquet",
+ &[String::from("part1"), String::from("part2")],
+ )
+ .await
+ .expect("partition pruning failed")
+ .collect::<Vec<_>>()
+ .await;
+
+ assert_eq!(pruned.len(), 2);
+ let f1 = pruned[0].as_ref().expect("first item not an error");
+ assert_eq!(
+ &f1.file_meta.sized_file.path,
+ "tablepath/part1=p1v2/part2=p2v1/file1.parquet"
+ );
+ assert_eq!(
+ &f1.partition_values,
+ &[
+ ScalarValue::Utf8(Some(String::from("p1v2"))),
+ ScalarValue::Utf8(Some(String::from("p2v1")))
+ ]
+ );
+ let f2 = pruned[1].as_ref().expect("second item not an error");
+ assert_eq!(
+ &f2.file_meta.sized_file.path,
+ "tablepath/part1=p1v2/part2=p2v1/file2.parquet"
+ );
+ assert_eq!(
+ &f2.partition_values,
+ &[
+ ScalarValue::Utf8(Some(String::from("p1v2"))),
+ ScalarValue::Utf8(Some(String::from("p2v1")))
+ ]
+ );
+ }
+
+ #[test]
+ fn test_parse_partitions_for_path() {
+ assert_eq!(
+ Some(vec![]),
+ parse_partitions_for_path("bucket/mytable", "bucket/mytable/file.csv", &[])
+ );
+ assert_eq!(
+ None,
+ parse_partitions_for_path(
+ "bucket/othertable",
+ "bucket/mytable/file.csv",
+ &[]
+ )
+ );
+ assert_eq!(
+ None,
+ parse_partitions_for_path(
+ "bucket/mytable",
+ "bucket/mytable/file.csv",
+ &[String::from("mypartition")]
+ )
+ );
+ assert_eq!(
+ Some(vec!["v1"]),
+ parse_partitions_for_path(
+ "bucket/mytable",
+ "bucket/mytable/mypartition=v1/file.csv",
+ &[String::from("mypartition")]
+ )
+ );
+ assert_eq!(
+ Some(vec!["v1"]),
+ parse_partitions_for_path(
+ "bucket/mytable/",
+ "bucket/mytable/mypartition=v1/file.csv",
+ &[String::from("mypartition")]
+ )
+ );
+ // Only hive style partitioning supported for now:
+ assert_eq!(
+ None,
+ parse_partitions_for_path(
+ "bucket/mytable",
+ "bucket/mytable/v1/file.csv",
+ &[String::from("mypartition")]
+ )
+ );
+ assert_eq!(
+ Some(vec!["v1", "v2"]),
+ parse_partitions_for_path(
+ "bucket/mytable",
+ "bucket/mytable/mypartition=v1/otherpartition=v2/file.csv",
+ &[String::from("mypartition"), String::from("otherpartition")]
+ )
+ );
+ assert_eq!(
+ Some(vec!["v1"]),
+ parse_partitions_for_path(
+ "bucket/mytable",
+ "bucket/mytable/mypartition=v1/otherpartition=v2/file.csv",
+ &[String::from("mypartition")]
+ )
+ );
+ }
+
+ #[test]
+ fn test_path_batch_roundtrip_no_partiton() {
+ let files = vec![
+ FileMeta {
+ sized_file: SizedFile {
+ path: String::from("mybucket/tablepath/part1=val1/file.parquet"),
+ size: 100,
+ },
+ last_modified: Some(Utc.timestamp_millis(1634722979123)),
+ },
+ FileMeta {
+ sized_file: SizedFile {
+ path: String::from("mybucket/tablepath/part1=val2/file.parquet"),
+ size: 100,
+ },
+ last_modified: None,
+ },
+ ];
+
+ let batches = paths_to_batch(&[], "mybucket/tablepath", &files)
+ .expect("Serialization of file list to batch failed");
+
+ let parsed_files = batches_to_paths(&[batches]);
+ assert_eq!(parsed_files.len(), 2);
+ assert_eq!(&parsed_files[0].partition_values, &[]);
+ assert_eq!(&parsed_files[1].partition_values, &[]);
+
+ let parsed_metas = parsed_files
Review comment:
👍
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-946150606
I plan to check this out carefully tomorrow
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735318173
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
Review comment:
It is kind of a balance to avoid having to much overhead when wrapping it with RecordBatch while keeping it as small as possible to accommodate all the possible sizes returned by the paging of the object store. I might as well have chosen 42... 👽 Do you have any idea how to provide a better default here?
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735309514
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -109,7 +108,8 @@ impl ListingOptions {
/// or file system listing capability to get the list of files.
pub struct ListingTable {
object_store: Arc<dyn ObjectStore>,
- path: String,
+ table_path: String,
Review comment:
no need because it is injected with the right object store. URI is required when it is injected with the registry.
--
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
[GitHub] [arrow-datafusion] houqp commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
houqp commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-950592101
Epic work @rdettai !
> when applying the filtering expression the dictionary needs to be expanded because many kernel ops are not supported on Dictionaries for now
I think we won't need to worry too much about this for this particular use-case because it won't be a problem anymore after we move to use scalar columnar value to store the partition values.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738254353
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -613,33 +614,28 @@ message ScanLimit {
uint32 limit = 1;
}
-message ParquetScanExecNode {
+message FileScanExecConf {
Review comment:
Promoted the `PhysicalPlanConfig` entity to the `physical_plan/file_format`, which enables us to factorize a lot of configuration that is the same for all file formats. Here for instance we serialize the configuration in one common entity to avoid having the same parameters over and over in all `XxxScanExecConf`.
The rational for doing this here is that clippy was complaining more and more often that methods had too many arguments 😉
Suggestions on the naming of the config entity (`PhysicalPlanConfig`) or the serialized version (`FileScanExecConf`) are welcome 😄
##########
File path: datafusion/src/datasource/mod.rs
##########
@@ -128,8 +130,8 @@ pub async fn get_statistics_with_limit(
pub struct PartitionedFile {
/// Path for the file (e.g. URL, filesystem path, etc)
pub file_meta: FileMeta,
- // Values of partition columns to be appended to each row
- // pub partition_value: Option<Vec<ScalarValue>>,
+ /// Values of partition columns to be appended to each row
+ pub partition_values: Vec<ScalarValue>,
Review comment:
Scalar value is always `Utf8` for now in the `ListingTable` implementation. Even though I usually prefer to apply YAGNI in these cases, this is a core abstraction that might be used by custom table providers, and those might want to allow stronger typing for their partition columns.
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -294,7 +295,7 @@ message ListingTableScanNode {
ProjectionColumns projection = 4;
Schema schema = 5;
repeated LogicalExprNode filters = 6;
- repeated string partitions = 7;
+ repeated string table_partition_cols = 7;
Review comment:
Renamed the `partitions` field to `table_partition_cols` to make it more explicit (cf https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735280489)
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
Review comment:
For now the file format executor only support string partition columns encoded as dictionaries to optimize space useage (partition columns have constant values for every row of the files, which is perfect for dictionary encoding).
##########
File path: datafusion/src/datasource/file_format/mod.rs
##########
@@ -29,33 +29,12 @@ use std::sync::Arc;
use crate::arrow::datatypes::SchemaRef;
use crate::error::Result;
use crate::logical_plan::Expr;
+use crate::physical_plan::file_format::PhysicalPlanConfig;
use crate::physical_plan::{ExecutionPlan, Statistics};
use async_trait::async_trait;
-use super::object_store::{ObjectReader, ObjectReaderStream, ObjectStore};
-use super::PartitionedFile;
-
-/// The configurations to be passed when creating a physical plan for
-/// a given file format.
-pub struct PhysicalPlanConfig {
Review comment:
moved this to physical_plan/file_format/mod.rs as it is a means of factorizing file format execution plan inputs
##########
File path: datafusion/src/datasource/listing/mod.rs
##########
@@ -0,0 +1,24 @@
+// 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.
+
+//! A table that uses the `ObjectStore` listing capability
+//! to get the list of files to process.
+
+mod helpers;
+mod table;
Review comment:
The `listing.rs` file was becoming huge, so it was split into two files. This also helps maintaining tests closer to the code. An unfortunate side effect is that we loose the git history for some of the code that was moved to `helpers.rs`.
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
Review comment:
`PartitionColumnProjector` helps injecting partition columns into file record batches both in `FileStream` and the `ParquetExec`. One interesting trick is the usage of a cache for the key part of the partition column dictionaries. Indeed, these columns are constant, so the dictionaries that encode them have all their keys equal to 0, so we can re-use the same "all-zero" buffer across batches. This makes the space consumption of the partition columns O(batch_size) instead of O(record_count).
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
+}
+
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
+#[derive(Debug, Clone)]
+pub struct PhysicalPlanConfig {
+ /// Store from which the `files` should be fetched
+ pub object_store: Arc<dyn ObjectStore>,
+ /// Schema before projection. It contains the columns that are expected
+ /// to be in the files without the table partition columns.
+ pub file_schema: SchemaRef,
+ /// List of files to be processed, grouped into partitions
+ pub file_groups: Vec<Vec<PartitionedFile>>,
+ /// Estimated overall statistics of the files, taking `filters` into account.
+ pub statistics: Statistics,
+ /// Columns on which to project the data. Indexes that are higher than the
+ /// number of columns of `file_schema` refer to `table_partition_cols`.
+ pub projection: Option<Vec<usize>>,
+ /// The maximum number of records per arrow column
+ pub batch_size: usize,
+ /// The minimum number of records required from this source plan
+ pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_cols: Vec<String>,
+}
+
+impl PhysicalPlanConfig {
+ /// Project the schema and the statistics on the given column indices
+ fn project(&self) -> (SchemaRef, Statistics) {
Review comment:
`PhysicalPlanConfig` was originally intended to reduce the number of arguments passed around from the table provider to the execution plans, but it ended up as a nice receiver for various schema and column name projection/extraction operations that where present in multiple places.
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
Review comment:
I tried to keep the methods well documented, even though these are private methods. I hope this will help reviewers in their task. Note here that `pruned_partition_list` and `split_files` appear as new but were actually ported from `listing.rs`.
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
+}
+
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
+#[derive(Debug, Clone)]
+pub struct PhysicalPlanConfig {
+ /// Store from which the `files` should be fetched
+ pub object_store: Arc<dyn ObjectStore>,
+ /// Schema before projection. It contains the columns that are expected
+ /// to be in the files without the table partition columns.
+ pub file_schema: SchemaRef,
Review comment:
we specify systematically what the schema represents:
- `file_schema` means that only the columns that can be found within the file are described
- `table_schema` means that on top of the file columns, we have partition columns
- `projected_schema` means that the column projection was applied
##########
File path: datafusion/src/physical_plan/file_format/avro.rs
##########
@@ -185,69 +141,134 @@ impl ExecutionPlan for AvroExec {
write!(
f,
"AvroExec: files={}, batch_size={}, limit={:?}",
- super::FileGroupsDisplay(&self.file_groups),
- self.batch_size,
- self.limit,
+ super::FileGroupsDisplay(&self.base_config.file_groups),
+ self.base_config.batch_size,
+ self.base_config.limit,
)
}
}
}
fn statistics(&self) -> Statistics {
- self.statistics.clone()
+ self.projected_statistics.clone()
}
}
#[cfg(test)]
#[cfg(feature = "avro")]
mod tests {
+ use crate::datasource::file_format::{avro::AvroFormat, FileFormat};
use crate::datasource::object_store::local::{
- local_file_meta, local_object_reader_stream, LocalFileSystem,
+ local_object_reader_stream, local_unpartitioned_file, LocalFileSystem,
};
+ use crate::scalar::ScalarValue;
+ use futures::StreamExt;
use super::*;
#[tokio::test]
- async fn test() -> Result<()> {
- use futures::StreamExt;
-
- use crate::datasource::file_format::{avro::AvroFormat, FileFormat};
-
+ async fn avro_exec_without_partition() -> Result<()> {
let testdata = crate::test_util::arrow_test_data();
let filename = format!("{}/avro/alltypes_plain.avro", testdata);
- let avro_exec = AvroExec::new(
- Arc::new(LocalFileSystem {}),
- vec![vec![PartitionedFile {
- file_meta: local_file_meta(filename.clone()),
- }]],
- Statistics::default(),
- AvroFormat {}
+ let avro_exec = AvroExec::new(PhysicalPlanConfig {
+ object_store: Arc::new(LocalFileSystem {}),
+ file_groups: vec![vec![local_unpartitioned_file(filename.clone())]],
+ file_schema: AvroFormat {}
.infer_schema(local_object_reader_stream(vec![filename]))
.await?,
- Some(vec![0, 1, 2]),
- 1024,
- None,
- );
+ statistics: Statistics::default(),
+ projection: Some(vec![0, 1, 2]),
+ batch_size: 1024,
+ limit: None,
+ table_partition_cols: vec![],
+ });
assert_eq!(avro_exec.output_partitioning().partition_count(), 1);
- let mut results = avro_exec.execute(0).await?;
- let batch = results.next().await.unwrap()?;
+ let mut results = avro_exec.execute(0).await.expect("plan execution failed");
+ let batch = results
+ .next()
+ .await
+ .expect("plan iterator empty")
+ .expect("plan iterator returned an error");
- assert_eq!(8, batch.num_rows());
- assert_eq!(3, batch.num_columns());
+ let expected = vec![
+ "+----+----------+-------------+",
+ "| id | bool_col | tinyint_col |",
+ "+----+----------+-------------+",
+ "| 4 | true | 0 |",
+ "| 5 | false | 1 |",
+ "| 6 | true | 0 |",
+ "| 7 | false | 1 |",
+ "| 2 | true | 0 |",
+ "| 3 | false | 1 |",
+ "| 0 | true | 0 |",
+ "| 1 | false | 1 |",
+ "+----+----------+-------------+",
+ ];
- let schema = batch.schema();
- let field_names: Vec<&str> =
- schema.fields().iter().map(|f| f.name().as_str()).collect();
- assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);
+ crate::assert_batches_eq!(expected, &[batch]);
Review comment:
moved to `assert_batches_eq` in both existing tests as well as new tests (focused on execution plans with partitioning) because this helps visualizing the differences in a glance.
##########
File path: datafusion/src/physical_plan/file_format/file_stream.rs
##########
@@ -63,52 +64,73 @@ impl<T> FormatReaderOpener for T where
pub struct FileStream<F: FormatReaderOpener> {
/// An iterator over record batches of the last file returned by file_iter
batch_iter: BatchIter,
- /// An iterator over input files
+ /// Partitioning column values for the current batch_iter
+ partition_values: Vec<ScalarValue>,
+ /// An iterator over input files.
file_iter: FileIter,
- /// The stream schema (file schema after projection)
- schema: SchemaRef,
+ /// The stream schema (file schema including partition columns and after
+ /// projection).
+ projected_schema: SchemaRef,
/// The remaining number of records to parse, None if no limit
remain: Option<usize>,
/// A closure that takes a reader and an optional remaining number of lines
/// (before reaching the limit) and returns a batch iterator. If the file reader
/// is not capable of limiting the number of records in the last batch, the file
/// stream will take care of truncating it.
file_reader: F,
+ /// The partition column projector
+ pc_projector: PartitionColumnProjector,
+ /// the store from which to source the files.
+ object_store: Arc<dyn ObjectStore>,
}
Review comment:
Reworked the `FileStream` helper to allow the injection of partition columns into the file record batches when necessary.
##########
File path: datafusion/tests/path_partition.rs
##########
@@ -0,0 +1,235 @@
+// 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.
+
+//! Test queries on partitioned datasets
Review comment:
We introduced a whole new file for tests that focus on the partitioning because the `sql.rs` test file was already saturated. This is maybe not the best split for integration tests.
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -96,37 +97,48 @@ impl ListingOptions {
.await?
.map(move |file_meta| object_store.file_reader(file_meta?.sized_file));
let file_schema = self.format.infer_schema(Box::pin(file_stream)).await?;
- // Add the partition columns to the file schema
- let mut fields = file_schema.fields().clone();
- for part in &self.partitions {
- fields.push(Field::new(part, DataType::Utf8, false));
- }
- Ok(Arc::new(Schema::new(fields)))
+ Ok(file_schema)
}
}
/// An implementation of `TableProvider` that uses the object store
/// or file system listing capability to get the list of files.
pub struct ListingTable {
object_store: Arc<dyn ObjectStore>,
- path: String,
- schema: SchemaRef,
+ table_path: String,
+ /// File fields only
+ file_schema: SchemaRef,
+ /// File fields + partition columns
+ table_schema: SchemaRef,
options: ListingOptions,
}
impl ListingTable {
/// Create new table that lists the FS to get the files to scan.
+ /// The provided `schema` must be resolved before creating the table
+ /// and should contain the fields of the file without the table
+ /// partitioning columns.
pub fn new(
object_store: Arc<dyn ObjectStore>,
- path: String,
- // the schema must be resolved before creating the table
- schema: SchemaRef,
+ table_path: String,
Review comment:
Just as with schemas, we try to be more explicit, indicating whether we are dealing with `table_path` or `file_path`
--
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
[GitHub] [arrow-datafusion] Dandandan commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Dandandan commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956567388
Thank you @rdettai ! Really nice work
--
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
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
houqp commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735323298
##########
File path: datafusion/src/datasource/file_format/mod.rs
##########
@@ -36,25 +36,25 @@ use async_trait::async_trait;
use super::object_store::{ObjectReader, ObjectReaderStream, ObjectStore};
use super::PartitionedFile;
-/// The configurations to be passed when creating a physical plan for
-/// a given file format.
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
pub struct PhysicalPlanConfig {
/// Store from which the `files` should be fetched
pub object_store: Arc<dyn ObjectStore>,
/// Schema before projection
- pub schema: SchemaRef,
+ pub file_schema: SchemaRef,
/// List of files to be processed, grouped into partitions
- pub files: Vec<Vec<PartitionedFile>>,
+ pub file_groups: Vec<Vec<PartitionedFile>>,
/// Estimated overall statistics of the plan, taking `filters` into account
pub statistics: Statistics,
/// Columns on which to project the data
pub projection: Option<Vec<usize>>,
/// The maximum number of records per arrow column
pub batch_size: usize,
- /// The filters that were pushed down to this execution plan
- pub filters: Vec<Expr>,
/// The minimum number of records required from this source plan
pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_dims: Vec<String>,
Review comment:
Conceptually they are handled as "virtual columns" during compute right? for example, when a user is writing a SQL query to filter against a partition, they will apply the filter to that partition just like other regular columns. I am suggesting partition column here because it's the term used in hive and spark, so readers would be more familiar with it. Are there systems that use partition dimensions as the naming convention?
--
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
[GitHub] [arrow-datafusion] rdettai commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-949804675
This is a bit harder than I thought it would be 😅. I will have to keep working on this next week. Feel free to give some feedback already, most of the important elements are already there.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740288428
##########
File path: datafusion/src/datasource/object_store/mod.rs
##########
@@ -70,7 +70,7 @@ pub enum ListEntry {
}
/// The path and size of the file.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
Review comment:
Not sure if it is really useful to add a trait that is not used 😊
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956248575
FYI fixed a logical conflict in https://github.com/apache/arrow-datafusion/pull/1141/commits/5d34be6df60fb2c36129d51b41ffce80c749ee67
--
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
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740239794
##########
File path: datafusion/tests/common.rs
##########
@@ -0,0 +1,40 @@
+// 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.
+
+//! methods that are common to multiple integration test setups
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+
+pub fn aggr_test_schema() -> SchemaRef {
Review comment:
```suggestion
pub(crate) fn aggr_test_schema() -> SchemaRef {
```
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739011698
##########
File path: datafusion/src/physical_plan/sort_preserving_merge.rs
##########
@@ -936,15 +936,18 @@ mod tests {
test::create_partitioned_csv("aggregate_test_100.csv", partitions).unwrap();
let csv = Arc::new(CsvExec::new(
- Arc::new(LocalFileSystem {}),
- files,
- Statistics::default(),
- Arc::clone(&schema),
+ PhysicalPlanConfig {
Review comment:
this is hell on earth 😄
--
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
[GitHub] [arrow-datafusion] alamb edited a comment on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb edited a comment on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956220581
I plan to merge from master and if all the tests pass put this one in. FYI @houqp / @Dandandan @Jimexist
Please let me know if you want more time to review
--
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
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740226457
##########
File path: ballista/rust/core/proto/ballista.proto
##########
@@ -613,33 +614,28 @@ message ScanLimit {
uint32 limit = 1;
}
-message ParquetScanExecNode {
+message FileScanExecConf {
repeated FileGroup file_groups = 1;
Schema schema = 2;
- uint32 batch_size = 4;
- repeated uint32 projection = 6;
- ScanLimit limit = 7;
- Statistics statistics = 8;
+ uint32 batch_size = 3;
Review comment:
would this be back-compatible?
--
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
[GitHub] [arrow-datafusion] Dandandan commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Dandandan commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956567388
Thank you @rdettai ! Really nice work
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956566766
As this has been outstanding for a long time and is a fairly large change, the potential for conflict is large -- I am going to merge it in now and hopefully we can keep iterating in future PRs. Thanks again @rdettai @Jimexist @houqp and @Dandandan -- 🚀
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735312508
##########
File path: datafusion/src/datasource/file_format/mod.rs
##########
@@ -36,25 +36,25 @@ use async_trait::async_trait;
use super::object_store::{ObjectReader, ObjectReaderStream, ObjectStore};
use super::PartitionedFile;
-/// The configurations to be passed when creating a physical plan for
-/// a given file format.
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
pub struct PhysicalPlanConfig {
/// Store from which the `files` should be fetched
pub object_store: Arc<dyn ObjectStore>,
/// Schema before projection
- pub schema: SchemaRef,
+ pub file_schema: SchemaRef,
/// List of files to be processed, grouped into partitions
- pub files: Vec<Vec<PartitionedFile>>,
+ pub file_groups: Vec<Vec<PartitionedFile>>,
/// Estimated overall statistics of the plan, taking `filters` into account
pub statistics: Statistics,
/// Columns on which to project the data
pub projection: Option<Vec<usize>>,
/// The maximum number of records per arrow column
pub batch_size: usize,
- /// The filters that were pushed down to this execution plan
- pub filters: Vec<Expr>,
/// The minimum number of records required from this source plan
pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_dims: Vec<String>,
Review comment:
Haha, I have to admit that I am in a huge hesitation regarding naming 😅. I am wondering if it's not the comment that should be changed. This partitions are originally encoded in the file path, that we then parse and project into a column if necessary. So they end up as columns, but they are not columns per say.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735788958
##########
File path: datafusion/src/datasource/file_format/mod.rs
##########
@@ -36,25 +36,25 @@ use async_trait::async_trait;
use super::object_store::{ObjectReader, ObjectReaderStream, ObjectStore};
use super::PartitionedFile;
-/// The configurations to be passed when creating a physical plan for
-/// a given file format.
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
pub struct PhysicalPlanConfig {
/// Store from which the `files` should be fetched
pub object_store: Arc<dyn ObjectStore>,
/// Schema before projection
- pub schema: SchemaRef,
+ pub file_schema: SchemaRef,
/// List of files to be processed, grouped into partitions
- pub files: Vec<Vec<PartitionedFile>>,
+ pub file_groups: Vec<Vec<PartitionedFile>>,
/// Estimated overall statistics of the plan, taking `filters` into account
pub statistics: Statistics,
/// Columns on which to project the data
pub projection: Option<Vec<usize>>,
/// The maximum number of records per arrow column
pub batch_size: usize,
- /// The filters that were pushed down to this execution plan
- pub filters: Vec<Expr>,
/// The minimum number of records required from this source plan
pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_dims: Vec<String>,
Review comment:
not that I am aware of, I'll change this to cols 😉
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739046851
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
Review comment:
I got 10 lines of code less with `ExpressionVisitor` than before, but more importantly, it is **much** fancier 😄
--
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
[GitHub] [arrow-datafusion] rdettai commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-953870803
I just noticed that there seems to be a small irregularity in terms of columns statistics vector length. I'm adding one last test about that with the appropriate patch if required 😉
--
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
[GitHub] [arrow-datafusion] alamb merged pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb merged pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740285300
##########
File path: datafusion/tests/common.rs
##########
@@ -0,0 +1,40 @@
+// 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.
+
+//! methods that are common to multiple integration test setups
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+
+pub fn aggr_test_schema() -> SchemaRef {
Review comment:
The `tests` module is private anyway, so the `(crate)` modifier does not have any effects here
--
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
[GitHub] [arrow-datafusion] Dandandan commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Dandandan commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956567388
Thank you @rdettai ! Really nice work
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739057046
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -50,3 +165,260 @@ impl<'a> Display for FileGroupsDisplay<'a> {
write!(f, "[{}]", parts.join(", "))
}
}
+
+/// A helper that projects partition columns into the file record batches
+struct PartitionColumnProjector {
Review comment:
done
--
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
[GitHub] [arrow-datafusion] Jimexist commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
Jimexist commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740232427
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,722 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr, ExpressionVisitor, Recursion},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// The `ExpressionVisitor` for `expr_applicable_for_cols`. Walks the tree to
+/// validate that the given expression is applicable with only the `col_names`
+/// set of columns.
+struct ApplicabilityVisitor<'a> {
+ col_names: &'a [String],
+ is_applicable: &'a mut bool,
+}
+
+impl ApplicabilityVisitor<'_> {
+ fn visit_volatility(self, volatility: Volatility) -> Recursion<Self> {
+ match volatility {
+ Volatility::Immutable => Recursion::Continue(self),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable | Volatility::Volatile => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ }
+ }
+}
+
+impl ExpressionVisitor for ApplicabilityVisitor<'_> {
+ fn pre_visit(self, expr: &Expr) -> Result<Recursion<Self>> {
+ let rec = match expr {
+ Expr::Column(logical_plan::Column { ref name, .. }) => {
+ *self.is_applicable &= self.col_names.contains(name);
+ Recursion::Stop(self) // leaf node anyway
+ }
+ Expr::Literal(_)
+ | Expr::Alias(_, _)
+ | Expr::ScalarVariable(_)
+ | Expr::Not(_)
+ | Expr::IsNotNull(_)
+ | Expr::IsNull(_)
+ | Expr::Negative(_)
+ | Expr::Cast { .. }
+ | Expr::TryCast { .. }
+ | Expr::BinaryExpr { .. }
+ | Expr::Between { .. }
+ | Expr::InList { .. }
+ | Expr::GetIndexedField { .. }
+ | Expr::Case { .. } => Recursion::Continue(self),
+
+ Expr::ScalarFunction { fun, .. } => self.visit_volatility(fun.volatility()),
+ Expr::ScalarUDF { fun, .. } => {
+ self.visit_volatility(fun.signature.volatility)
+ }
+
+ // TODO other expressions are not handled yet:
+ // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in some edge cases
+ // - Can `Wildcard` be considered as a `Literal`?
+ // - ScalarVariable could be `applicable`, but that would require access to the context
+ Expr::AggregateUDF { .. }
+ | Expr::AggregateFunction { .. }
+ | Expr::Sort { .. }
+ | Expr::WindowFunction { .. }
+ | Expr::Wildcard => {
+ *self.is_applicable = false;
+ Recursion::Stop(self)
+ }
+ };
+ Ok(rec)
+ }
+}
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+ let mut is_applicable = true;
+ expr.accept(ApplicabilityVisitor {
+ col_names,
+ is_applicable: &mut is_applicable,
+ })
+ .unwrap();
+ is_applicable
+}
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let chunk_size = (partitioned_files.len() + n - 1) / n;
Review comment:
```suggestion
// effectively this is div with rounding up instead of truncating
let chunk_size = (partitioned_files.len() + n - 1) / n;
```
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956566766
As this has been outstanding for a long time and is a fairly large change, the potential for conflict is large -- I am going to merge it in now and hopefully we can keep iterating in future PRs. Thanks again @rdettai @Jimexist @houqp and @Dandandan -- 🚀
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-956220581
I am going to merge from master and if all the tests pass put this one in. FYI @houqp / @Dandandan @Jimexist
--
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
[GitHub] [arrow-datafusion] alamb merged pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb merged pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735806636
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
+/// - one column for with the original path named `_df_part_file_path_`
+/// - one column for with the last modified date named `_df_part_file_modified_`
+/// - ... one column by partition ...
+/// Note: For the last modified date, this looses precisions higher than millisecond.
+fn paths_to_batch(
+ table_partition_dims: &[String],
+ table_path: &str,
+ metas: &[FileMeta],
+) -> Result<RecordBatch> {
+ let mut key_builder = StringBuilder::new(metas.len());
+ let mut length_builder = UInt64Builder::new(metas.len());
+ let mut modified_builder = Date64Builder::new(metas.len());
+ let mut partition_builders = table_partition_dims
+ .iter()
+ .map(|_| StringBuilder::new(metas.len()))
+ .collect::<Vec<_>>();
+ for file_meta in metas {
+ if let Some(partition_values) =
+ parse_partitions_for_path(table_path, file_meta.path(), table_partition_dims)
+ {
+ key_builder.append_value(file_meta.path())?;
+ length_builder.append_value(file_meta.size())?;
+ match file_meta.last_modified {
+ Some(lm) => modified_builder.append_value(lm.timestamp_millis())?,
+ None => modified_builder.append_null()?,
+ }
+ for (i, part_val) in partition_values.iter().enumerate() {
+ partition_builders[i].append_value(part_val)?;
+ }
+ } else {
+ debug!("No partitioning for path {}", file_meta.path());
+ }
+ }
+
+ // finish all builders
+ let mut col_arrays: Vec<ArrayRef> = vec![
+ ArrayBuilder::finish(&mut key_builder),
+ ArrayBuilder::finish(&mut length_builder),
+ ArrayBuilder::finish(&mut modified_builder),
+ ];
+ for mut partition_builder in partition_builders {
+ col_arrays.push(ArrayBuilder::finish(&mut partition_builder));
+ }
+
+ // put the schema together
+ let mut fields = vec![
+ Field::new(FILE_SIZE_COLUMN_NAME, DataType::Utf8, false),
+ Field::new(FILE_PATH_COLUMN_NAME, DataType::UInt64, false),
+ Field::new(FILE_MODIFIED_COLUMN_NAME, DataType::Date64, false),
+ ];
+ for pn in table_partition_dims {
+ fields.push(Field::new(pn, DataType::Utf8, false));
+ }
+
+ let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), col_arrays)?;
+ Ok(batch)
+}
+
+/// convert a set of record batches created by `paths_to_batch()` back to partitioned files.
+fn batches_to_paths(batches: &[RecordBatch]) -> Vec<PartitionedFile> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ let key_array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let length_array = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap();
+ let modified_array = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Date64Array>()
+ .unwrap();
+
+ (0..batch.num_rows()).map(move |row| PartitionedFile {
+ file_meta: FileMeta {
+ last_modified: match modified_array.is_null(row) {
+ false => Some(Utc.timestamp_millis(modified_array.value(row))),
+ true => None,
+ },
+ sized_file: SizedFile {
+ path: key_array.value(row).to_owned(),
+ size: length_array.value(row),
+ },
+ },
+ partition_values: (3..batch.columns().len())
+ .map(|col| {
+ ScalarValue::try_from_array(batch.column(col), row).unwrap()
+ })
+ .collect(),
+ })
+ })
+ .collect()
+}
+
+/// Extract the partition values for the given `file_path` (in the given `table_path`)
+/// associated to the partitions defined by `table_partition_dims`
+fn parse_partitions_for_path<'a>(
+ table_path: &str,
+ file_path: &'a str,
+ table_partition_dims: &[String],
+) -> Option<Vec<&'a str>> {
+ let subpath = file_path.strip_prefix(table_path)?;
+
+ // ignore whether table_path ended with "/" or not
+ let subpath = match subpath.strip_prefix('/') {
+ Some(subpath) => subpath,
+ None => subpath,
+ };
+
+ let mut part_values = vec![];
+ for (path, pn) in subpath.split('/').zip(table_partition_dims) {
+ if let Some(val) = path.strip_prefix(&format!("{}=", pn)) {
Review comment:
agreed, let's do something a bit nicer here 😃
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738768467
##########
File path: ballista/rust/core/src/serde/physical_plan/to_proto.rs
##########
@@ -722,3 +675,32 @@ impl From<&Statistics> for protobuf::Statistics {
}
}
}
+
+impl TryFrom<&PhysicalPlanConfig> for protobuf::FileScanExecConf {
+ type Error = BallistaError;
+ fn try_from(
+ conf: &PhysicalPlanConfig,
+ ) -> Result<protobuf::FileScanExecConf, Self::Error> {
+ let file_groups = conf
+ .file_groups
+ .iter()
+ .map(|p| p.as_slice().try_into())
+ .collect::<Result<Vec<_>, _>>()?;
+
+ Ok(protobuf::FileScanExecConf {
+ file_groups,
+ statistics: Some((&conf.statistics).into()),
+ limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }),
+ projection: conf
+ .projection
+ .as_ref()
+ .unwrap_or(&vec![])
Review comment:
`None` is encoded with empty vec. Not the cleanest but works here as projection with no column could not have another meaning.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735343229
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
Review comment:
> ATM we don't gain anything from having it in multiple batches as we wait on the entire results to be returned?
You are 100% right, my logic only holds with a streaming datasource, and I gave up on that for this PR. I'll rise this to the default batch size for now and add a comment.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735315517
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
Review comment:
This is not new code but appears as such because I refactored the existing one into different files. I don't mind adding your optimization though 😃 (it is hard to read, but so is the current version 😄)
--
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
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
houqp commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735280489
##########
File path: datafusion/src/datasource/file_format/mod.rs
##########
@@ -36,25 +36,25 @@ use async_trait::async_trait;
use super::object_store::{ObjectReader, ObjectReaderStream, ObjectStore};
use super::PartitionedFile;
-/// The configurations to be passed when creating a physical plan for
-/// a given file format.
+/// The base configurations to provide when creating a physical plan for
+/// any given file format.
pub struct PhysicalPlanConfig {
/// Store from which the `files` should be fetched
pub object_store: Arc<dyn ObjectStore>,
/// Schema before projection
- pub schema: SchemaRef,
+ pub file_schema: SchemaRef,
/// List of files to be processed, grouped into partitions
- pub files: Vec<Vec<PartitionedFile>>,
+ pub file_groups: Vec<Vec<PartitionedFile>>,
/// Estimated overall statistics of the plan, taking `filters` into account
pub statistics: Statistics,
/// Columns on which to project the data
pub projection: Option<Vec<usize>>,
/// The maximum number of records per arrow column
pub batch_size: usize,
- /// The filters that were pushed down to this execution plan
- pub filters: Vec<Expr>,
/// The minimum number of records required from this source plan
pub limit: Option<usize>,
+ /// The partitioning column names
+ pub table_partition_dims: Vec<String>,
Review comment:
curious why not name it as `table_partition_cols` to better align with the comment? the type should make it clear that it's storing column names.
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -109,7 +108,8 @@ impl ListingOptions {
/// or file system listing capability to get the list of files.
pub struct ListingTable {
object_store: Arc<dyn ObjectStore>,
- path: String,
+ table_path: String,
Review comment:
I think the table_path should contain contain uri scheme as well? if so, `table_uri` might be a better name here.
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
Review comment:
any particular reason for picking 64 here?
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
Review comment:
here we go datafusion in datafusion :D I do think in the future we could optimize this to only use the physical execution code to directly evaluate the record batches.
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
Review comment:
minor optimizaton:
```suggestion
let chunk_size = (partitioned_files.len()+n-1) / n;
```
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
Review comment:
I think we can also avoid the filter when no filter is provided?
```suggestion
if table_partition_dims.is_empty() || filters.empty() {
```
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
+/// - one column for with the original path named `_df_part_file_path_`
+/// - one column for with the last modified date named `_df_part_file_modified_`
+/// - ... one column by partition ...
+/// Note: For the last modified date, this looses precisions higher than millisecond.
+fn paths_to_batch(
+ table_partition_dims: &[String],
+ table_path: &str,
+ metas: &[FileMeta],
+) -> Result<RecordBatch> {
+ let mut key_builder = StringBuilder::new(metas.len());
+ let mut length_builder = UInt64Builder::new(metas.len());
+ let mut modified_builder = Date64Builder::new(metas.len());
+ let mut partition_builders = table_partition_dims
+ .iter()
+ .map(|_| StringBuilder::new(metas.len()))
+ .collect::<Vec<_>>();
+ for file_meta in metas {
+ if let Some(partition_values) =
+ parse_partitions_for_path(table_path, file_meta.path(), table_partition_dims)
+ {
+ key_builder.append_value(file_meta.path())?;
+ length_builder.append_value(file_meta.size())?;
+ match file_meta.last_modified {
+ Some(lm) => modified_builder.append_value(lm.timestamp_millis())?,
+ None => modified_builder.append_null()?,
+ }
+ for (i, part_val) in partition_values.iter().enumerate() {
+ partition_builders[i].append_value(part_val)?;
+ }
+ } else {
+ debug!("No partitioning for path {}", file_meta.path());
+ }
+ }
+
+ // finish all builders
+ let mut col_arrays: Vec<ArrayRef> = vec![
+ ArrayBuilder::finish(&mut key_builder),
+ ArrayBuilder::finish(&mut length_builder),
+ ArrayBuilder::finish(&mut modified_builder),
+ ];
+ for mut partition_builder in partition_builders {
+ col_arrays.push(ArrayBuilder::finish(&mut partition_builder));
+ }
+
+ // put the schema together
+ let mut fields = vec![
+ Field::new(FILE_SIZE_COLUMN_NAME, DataType::Utf8, false),
+ Field::new(FILE_PATH_COLUMN_NAME, DataType::UInt64, false),
Review comment:
typo?
```suggestion
Field::new(FILE_SIZE_COLUMN_NAME, DataType::Uint64, false),
Field::new(FILE_PATH_COLUMN_NAME, DataType::Utf8, false),
```
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
+ .map(|v| v.into_iter().collect::<Result<Vec<_>>>())
+ .map(move |metas| {
+ paths_to_batch(&table_partition_dims, &stream_path, &metas?)
+ })
+ .try_collect()
+ .await?;
+
+ let mem_table = MemTable::try_new(batches[0].schema(), vec![batches])?;
+
+ // Filter the partitions using a local datafusion context
+ // TODO having the external context would allow us to resolve `Volatility::Stable`
+ // scalar functions (`ScalarFunction` & `ScalarUDF`) and `ScalarVariable`s
+ let mut ctx = ExecutionContext::new();
+ let mut df = ctx.read_table(Arc::new(mem_table))?;
+ for filter in applicable_filters {
+ df = df.filter(filter.clone())?;
+ }
+ let filtered_batches = df.collect().await?;
+
+ Ok(Box::pin(stream::iter(
+ batches_to_paths(&filtered_batches).into_iter().map(Ok),
+ )))
+ }
+}
+
+/// convert the paths of the files to a record batch with the following columns:
+/// - one column for the file size named `_df_part_file_size_`
+/// - one column for with the original path named `_df_part_file_path_`
+/// - one column for with the last modified date named `_df_part_file_modified_`
+/// - ... one column by partition ...
+/// Note: For the last modified date, this looses precisions higher than millisecond.
+fn paths_to_batch(
+ table_partition_dims: &[String],
+ table_path: &str,
+ metas: &[FileMeta],
+) -> Result<RecordBatch> {
+ let mut key_builder = StringBuilder::new(metas.len());
+ let mut length_builder = UInt64Builder::new(metas.len());
+ let mut modified_builder = Date64Builder::new(metas.len());
+ let mut partition_builders = table_partition_dims
+ .iter()
+ .map(|_| StringBuilder::new(metas.len()))
+ .collect::<Vec<_>>();
+ for file_meta in metas {
+ if let Some(partition_values) =
+ parse_partitions_for_path(table_path, file_meta.path(), table_partition_dims)
+ {
+ key_builder.append_value(file_meta.path())?;
+ length_builder.append_value(file_meta.size())?;
+ match file_meta.last_modified {
+ Some(lm) => modified_builder.append_value(lm.timestamp_millis())?,
+ None => modified_builder.append_null()?,
+ }
+ for (i, part_val) in partition_values.iter().enumerate() {
+ partition_builders[i].append_value(part_val)?;
+ }
+ } else {
+ debug!("No partitioning for path {}", file_meta.path());
+ }
+ }
+
+ // finish all builders
+ let mut col_arrays: Vec<ArrayRef> = vec![
+ ArrayBuilder::finish(&mut key_builder),
+ ArrayBuilder::finish(&mut length_builder),
+ ArrayBuilder::finish(&mut modified_builder),
+ ];
+ for mut partition_builder in partition_builders {
+ col_arrays.push(ArrayBuilder::finish(&mut partition_builder));
+ }
+
+ // put the schema together
+ let mut fields = vec![
+ Field::new(FILE_SIZE_COLUMN_NAME, DataType::Utf8, false),
+ Field::new(FILE_PATH_COLUMN_NAME, DataType::UInt64, false),
+ Field::new(FILE_MODIFIED_COLUMN_NAME, DataType::Date64, false),
+ ];
+ for pn in table_partition_dims {
+ fields.push(Field::new(pn, DataType::Utf8, false));
+ }
+
+ let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), col_arrays)?;
+ Ok(batch)
+}
+
+/// convert a set of record batches created by `paths_to_batch()` back to partitioned files.
+fn batches_to_paths(batches: &[RecordBatch]) -> Vec<PartitionedFile> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ let key_array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let length_array = batch
+ .column(1)
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap();
+ let modified_array = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<Date64Array>()
+ .unwrap();
+
+ (0..batch.num_rows()).map(move |row| PartitionedFile {
+ file_meta: FileMeta {
+ last_modified: match modified_array.is_null(row) {
+ false => Some(Utc.timestamp_millis(modified_array.value(row))),
+ true => None,
+ },
+ sized_file: SizedFile {
+ path: key_array.value(row).to_owned(),
+ size: length_array.value(row),
+ },
+ },
+ partition_values: (3..batch.columns().len())
+ .map(|col| {
+ ScalarValue::try_from_array(batch.column(col), row).unwrap()
+ })
+ .collect(),
+ })
+ })
+ .collect()
+}
+
+/// Extract the partition values for the given `file_path` (in the given `table_path`)
+/// associated to the partitions defined by `table_partition_dims`
+fn parse_partitions_for_path<'a>(
+ table_path: &str,
+ file_path: &'a str,
+ table_partition_dims: &[String],
+) -> Option<Vec<&'a str>> {
+ let subpath = file_path.strip_prefix(table_path)?;
+
+ // ignore whether table_path ended with "/" or not
+ let subpath = match subpath.strip_prefix('/') {
+ Some(subpath) => subpath,
+ None => subpath,
+ };
+
+ let mut part_values = vec![];
+ for (path, pn) in subpath.split('/').zip(table_partition_dims) {
+ if let Some(val) = path.strip_prefix(&format!("{}=", pn)) {
Review comment:
I think we could avoid the `format!` string creation here by applying https://doc.rust-lang.org/stable/std/primitive.str.html#method.split_once with `=` on path and compare `pn` with the first component.
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
Review comment:
to avoid the complexity of `expr_applicable_for_cols`, perhaps we could just throw a runtime error to the user if invalid filters are provided? this also makes sure that if there is a typo in the filter, we won't silent the error.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738991233
##########
File path: datafusion/src/physical_plan/file_format/mod.rs
##########
@@ -24,19 +24,134 @@ mod json;
mod parquet;
pub use self::parquet::ParquetExec;
+use arrow::{
+ array::{ArrayData, ArrayRef, DictionaryArray, UInt8BufferBuilder},
+ buffer::Buffer,
+ datatypes::{DataType, Field, Schema, SchemaRef, UInt8Type},
+ error::{ArrowError, Result as ArrowResult},
+ record_batch::RecordBatch,
+};
pub use avro::AvroExec;
pub use csv::CsvExec;
pub use json::NdJsonExec;
-use crate::datasource::PartitionedFile;
-use std::fmt::{Display, Formatter, Result};
+use crate::{
+ datasource::{object_store::ObjectStore, PartitionedFile},
+ scalar::ScalarValue,
+};
+use std::{
+ collections::HashMap,
+ fmt::{Display, Formatter, Result as FmtResult},
+ sync::Arc,
+ vec,
+};
+
+use super::{ColumnStatistics, Statistics};
+
+lazy_static! {
+ /// The datatype used for all partitioning columns for now
+ pub static ref DEFAULT_PARTITION_COLUMN_DATATYPE: DataType = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8));
Review comment:
My tendency would also be to force the listing provider to use `Utf8` only and let the user cast explicitly in its query if necessary. But after various discussions with @Dandandan @yjshen and @houqp, the following points emerged in favor of letting `ListingTable` perform the cast:
- "Better query user experience so users won't need to manually add the casting in their queries if the type is already known"
- "Help with spark migration where they even support automatic partition column type inference (can be turned off)"
- Avoid the cost of going through an intermediate Utf8 batch materialization (though the cast should be pretty inexpensive as partition columns are represented by dictionaries with only one dict value)
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735319602
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
Review comment:
correct!
--
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
[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
houqp commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r735327286
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,682 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let mut chunk_size = partitioned_files.len() / n;
+ if partitioned_files.len() % n > 0 {
+ chunk_size += 1;
+ }
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// relative to irrelevant partitions using `filters` expressions
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_dims: &[String],
+) -> Result<PartitionedFileStream> {
+ if table_partition_dims.is_empty() {
+ Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ file_meta: f?,
+ partition_values: vec![],
+ })
+ }),
+ ))
+ } else {
+ let applicable_filters = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_dims, f));
+
+ let table_partition_dims = table_partition_dims.to_vec();
+ let stream_path = table_path.to_owned();
+ // TODO avoid collecting but have a streaming memory table instead
+ let batches: Vec<RecordBatch> = store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .chunks(64)
Review comment:
got it, I don't have good suggestion on this either. But I think it would be good to comment inline that this number is not set in stone and might need further tuning.
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739051018
##########
File path: datafusion/src/datasource/listing/helpers.rs
##########
@@ -0,0 +1,729 @@
+// 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.
+
+//! Helper functions for the table implementation
+
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayBuilder, ArrayRef, Date64Array, Date64Builder, StringArray,
+ StringBuilder, UInt64Array, UInt64Builder,
+ },
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
+use chrono::{TimeZone, Utc};
+use futures::{
+ stream::{self},
+ StreamExt, TryStreamExt,
+};
+use log::debug;
+
+use crate::{
+ error::Result,
+ execution::context::ExecutionContext,
+ logical_plan::{self, Expr},
+ physical_plan::functions::Volatility,
+ scalar::ScalarValue,
+};
+
+use crate::datasource::{
+ object_store::{FileMeta, ObjectStore, SizedFile},
+ MemTable, PartitionedFile, PartitionedFileStream,
+};
+
+const FILE_SIZE_COLUMN_NAME: &str = "_df_part_file_size_";
+const FILE_PATH_COLUMN_NAME: &str = "_df_part_file_path_";
+const FILE_MODIFIED_COLUMN_NAME: &str = "_df_part_file_modified_";
+
+/// Check whether the given expression can be resolved using only the columns `col_names`.
+/// This means that if this function returns true:
+/// - the table provider can filter the table partition values with this expression
+/// - the expression can be marked as `TableProviderFilterPushDown::Exact` once this filtering
+/// was performed
+pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
+ match expr {
+ // leaf
+ Expr::Literal(_) => true,
+ // TODO how to handle qualified / unqualified names?
+ Expr::Column(logical_plan::Column { ref name, .. }) => col_names.contains(name),
+ // unary
+ Expr::Alias(child, _)
+ | Expr::Not(child)
+ | Expr::IsNotNull(child)
+ | Expr::IsNull(child)
+ | Expr::Negative(child)
+ | Expr::Cast { expr: child, .. }
+ | Expr::TryCast { expr: child, .. } => expr_applicable_for_cols(col_names, child),
+ // binary
+ Expr::BinaryExpr {
+ ref left,
+ ref right,
+ ..
+ } => {
+ expr_applicable_for_cols(col_names, left)
+ && expr_applicable_for_cols(col_names, right)
+ }
+ // ternary
+ Expr::Between {
+ expr: item,
+ low,
+ high,
+ ..
+ } => {
+ expr_applicable_for_cols(col_names, item)
+ && expr_applicable_for_cols(col_names, low)
+ && expr_applicable_for_cols(col_names, high)
+ }
+ // variadic
+ Expr::ScalarFunction { fun, args } => match fun.volatility() {
+ Volatility::Immutable => args
+ .iter()
+ .all(|arg| expr_applicable_for_cols(col_names, arg)),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable => false,
+ Volatility::Volatile => false,
+ },
+ Expr::ScalarUDF { fun, args } => match fun.signature.volatility {
+ Volatility::Immutable => args
+ .iter()
+ .all(|arg| expr_applicable_for_cols(col_names, arg)),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ Volatility::Stable => false,
+ Volatility::Volatile => false,
+ },
+ Expr::InList {
+ expr: item, list, ..
+ } => {
+ expr_applicable_for_cols(col_names, item)
+ && list.iter().all(|e| expr_applicable_for_cols(col_names, e))
+ }
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ let expr_constant = expr
+ .as_ref()
+ .map(|e| expr_applicable_for_cols(col_names, e))
+ .unwrap_or(true);
+ let else_constant = else_expr
+ .as_ref()
+ .map(|e| expr_applicable_for_cols(col_names, e))
+ .unwrap_or(true);
+ let when_then_constant = when_then_expr.iter().all(|(w, th)| {
+ expr_applicable_for_cols(col_names, w)
+ && expr_applicable_for_cols(col_names, th)
+ });
+ expr_constant && else_constant && when_then_constant
+ }
+ // TODO other expressions are not handled yet:
+ // - AGGREGATE, WINDOW and SORT should not end up in filter conditions, except maybe in some edge cases
+ // - Can `Wildcard` be considered as a `Literal`?
+ // - ScalarVariable could be `applicable`, but that would require access to the context
+ _ => false,
+ }
+}
+
+/// Partition the list of files into `n` groups
+pub fn split_files(
+ partitioned_files: Vec<PartitionedFile>,
+ n: usize,
+) -> Vec<Vec<PartitionedFile>> {
+ if partitioned_files.is_empty() {
+ return vec![];
+ }
+ let chunk_size = (partitioned_files.len() + n - 1) / n;
+ partitioned_files
+ .chunks(chunk_size)
+ .map(|c| c.to_vec())
+ .collect()
+}
+
+/// Discover the partitions on the given path and prune out files
+/// that belong to irrelevant partitions using `filters` expressions.
+/// `filters` might contain expressions that can be resolved only at the
+/// file level (e.g. Parquet row group pruning).
+///
+/// TODO for tables with many files (10k+), it will usually more efficient
+/// to first list the folders relative to the first partition dimension,
+/// prune those, then list only the contain of the remaining folders.
+pub async fn pruned_partition_list(
+ store: &dyn ObjectStore,
+ table_path: &str,
+ filters: &[Expr],
+ file_extension: &str,
+ table_partition_cols: &[String],
+) -> Result<PartitionedFileStream> {
+ // if no partition col => simply list all the files
+ if table_partition_cols.is_empty() {
+ return Ok(Box::pin(
+ store
+ .list_file_with_suffix(table_path, file_extension)
+ .await?
+ .map(|f| {
+ Ok(PartitionedFile {
+ partition_values: vec![],
+ file_meta: f?,
+ })
+ }),
+ ));
+ }
+
+ let applicable_filters: Vec<_> = filters
+ .iter()
+ .filter(|f| expr_applicable_for_cols(table_partition_cols, f))
+ .collect();
+ let stream_path = table_path.to_owned();
+ if applicable_filters.is_empty() {
+ // parse the partition values while listing all the files
+ // TODO we might avoid parsing the partition values if they are not used in any projection
Review comment:
Agreed, I'll update the comment and remove the `// TODO`
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r739048453
##########
File path: datafusion/src/datasource/listing/table.rs
##########
@@ -51,9 +52,9 @@ pub struct ListingOptions {
/// partitioning expected should be named "a" and "b":
/// - If there is a third level of partitioning it will be ignored.
/// - Files that don't follow this partitioning will be ignored.
- /// Note that only `DataType::Utf8` is supported for the column type.
- /// TODO implement case where partitions.len() > 0
- pub partitions: Vec<String>,
+ /// Note that only `DEFAULT_PARTITION_COLUMN_DATATYPE` is currently
+ /// supported for the column type.
Review comment:
Let's discuss this in the other https://github.com/apache/arrow-datafusion/pull/1141#discussion_r738738023
--
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
[GitHub] [arrow-datafusion] alamb commented on pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb commented on pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#issuecomment-951813525
> I plan to review this PR tomorrow (when I am fresher and can git it the look it deserves)
@rdettai mentioned he has some more work planned for this PR so I will postpone additional review until that is complete
--
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
[GitHub] [arrow-datafusion] rdettai commented on a change in pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
rdettai commented on a change in pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141#discussion_r740285300
##########
File path: datafusion/tests/common.rs
##########
@@ -0,0 +1,40 @@
+// 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.
+
+//! methods that are common to multiple integration test setups
+
+use std::sync::Arc;
+
+use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+
+pub fn aggr_test_schema() -> SchemaRef {
Review comment:
The `tests` module is enabled for tests only anyway, so the `(crate)` modifier does not have much effects here (you would need to build DataFusion in test mode to use it).
--
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
[GitHub] [arrow-datafusion] alamb merged pull request #1141: File partitioning for ListingTable
Posted by GitBox <gi...@apache.org>.
alamb merged pull request #1141:
URL: https://github.com/apache/arrow-datafusion/pull/1141
--
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