-
Notifications
You must be signed in to change notification settings - Fork 1.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Simplify file struct abstractions #1120
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -23,6 +23,7 @@ use std::{any::Any, convert::TryInto}; | |
|
||
use crate::datasource::file_format::parquet::ChunkObjectReader; | ||
use crate::datasource::object_store::ObjectStore; | ||
use crate::datasource::PartitionedFile; | ||
use crate::{ | ||
error::{DataFusionError, Result}, | ||
logical_plan::{Column, Expr}, | ||
|
@@ -59,14 +60,12 @@ use tokio::{ | |
|
||
use async_trait::async_trait; | ||
|
||
use crate::datasource::{FilePartition, PartitionedFile}; | ||
|
||
/// Execution plan for scanning one or more Parquet partitions | ||
#[derive(Debug, Clone)] | ||
pub struct ParquetExec { | ||
object_store: Arc<dyn ObjectStore>, | ||
/// Parquet partitions to read | ||
partitions: Vec<ParquetPartition>, | ||
/// List of parquet files, grouped by output partition | ||
file_groups: Vec<Vec<PartitionedFile>>, | ||
/// Schema after projection is applied | ||
schema: SchemaRef, | ||
/// Projection for which columns to load | ||
|
@@ -83,23 +82,6 @@ pub struct ParquetExec { | |
limit: Option<usize>, | ||
} | ||
|
||
/// Represents one partition of a Parquet data set and this currently means one Parquet file. | ||
/// | ||
/// In the future it would be good to support subsets of files based on ranges of row groups | ||
/// so that we can better parallelize reads of large files across available cores (see | ||
/// [ARROW-10995](https://issues.apache.org/jira/browse/ARROW-10995)). | ||
/// | ||
/// We may also want to support reading Parquet files that are partitioned based on a key and | ||
/// in this case we would want this partition struct to represent multiple files for a given | ||
/// partition key (see [ARROW-11019](https://issues.apache.org/jira/browse/ARROW-11019)). | ||
#[derive(Debug, Clone)] | ||
pub struct ParquetPartition { | ||
Comment on lines
-86
to
-96
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These comments were mostly outdated and the other features mentioned are now planned in the |
||
/// The Parquet filename for this partition | ||
pub file_partition: FilePartition, | ||
/// Execution metrics | ||
metrics: ExecutionPlanMetricsSet, | ||
} | ||
|
||
/// Stores metrics about the parquet execution for a particular parquet file | ||
#[derive(Debug, Clone)] | ||
struct ParquetFileMetrics { | ||
|
@@ -115,24 +97,16 @@ impl ParquetExec { | |
#[allow(clippy::too_many_arguments)] | ||
pub fn new( | ||
object_store: Arc<dyn ObjectStore>, | ||
files: Vec<Vec<PartitionedFile>>, | ||
file_groups: Vec<Vec<PartitionedFile>>, | ||
statistics: Statistics, | ||
schema: SchemaRef, | ||
projection: Option<Vec<usize>>, | ||
predicate: Option<Expr>, | ||
batch_size: usize, | ||
limit: Option<usize>, | ||
) -> Self { | ||
debug!("Creating ParquetExec, desc: {:?}, projection {:?}, predicate: {:?}, limit: {:?}", | ||
files, projection, predicate, limit); | ||
|
||
let metrics = ExecutionPlanMetricsSet::new(); | ||
|
||
let partitions = files | ||
.into_iter() | ||
.enumerate() | ||
.map(|(i, f)| ParquetPartition::new(f, i, metrics.clone())) | ||
.collect::<Vec<_>>(); | ||
debug!("Creating ParquetExec, files: {:?}, projection {:?}, predicate: {:?}, limit: {:?}", | ||
file_groups, projection, predicate, limit); | ||
|
||
let metrics = ExecutionPlanMetricsSet::new(); | ||
let predicate_creation_errors = | ||
|
@@ -162,7 +136,7 @@ impl ParquetExec { | |
|
||
Self { | ||
object_store, | ||
partitions, | ||
file_groups, | ||
schema: projected_schema, | ||
projection, | ||
metrics, | ||
|
@@ -204,11 +178,8 @@ impl ParquetExec { | |
} | ||
|
||
/// List of data files | ||
pub fn partitions(&self) -> Vec<&[PartitionedFile]> { | ||
self.partitions | ||
.iter() | ||
.map(|fp| fp.file_partition.files.as_slice()) | ||
.collect() | ||
pub fn file_groups(&self) -> &[Vec<PartitionedFile>] { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We replace
|
||
&self.file_groups | ||
} | ||
/// Optional projection for which columns to load | ||
pub fn projection(&self) -> &[usize] { | ||
|
@@ -225,20 +196,6 @@ impl ParquetExec { | |
} | ||
} | ||
|
||
impl ParquetPartition { | ||
/// Create a new parquet partition | ||
pub fn new( | ||
files: Vec<PartitionedFile>, | ||
index: usize, | ||
metrics: ExecutionPlanMetricsSet, | ||
) -> Self { | ||
Self { | ||
file_partition: FilePartition { index, files }, | ||
metrics, | ||
} | ||
} | ||
} | ||
|
||
impl ParquetFileMetrics { | ||
/// Create new metrics | ||
pub fn new( | ||
|
@@ -279,7 +236,7 @@ impl ExecutionPlan for ParquetExec { | |
|
||
/// Get the output partitioning of this plan | ||
fn output_partitioning(&self) -> Partitioning { | ||
Partitioning::UnknownPartitioning(self.partitions.len()) | ||
Partitioning::UnknownPartitioning(self.file_groups.len()) | ||
} | ||
|
||
fn with_new_children( | ||
|
@@ -304,7 +261,7 @@ impl ExecutionPlan for ParquetExec { | |
Receiver<ArrowResult<RecordBatch>>, | ||
) = channel(2); | ||
|
||
let partition = self.partitions[partition_index].clone(); | ||
let partition = self.file_groups[partition_index].clone(); | ||
let metrics = self.metrics.clone(); | ||
let projection = self.projection.clone(); | ||
let predicate_builder = self.predicate_builder.clone(); | ||
|
@@ -338,18 +295,12 @@ impl ExecutionPlan for ParquetExec { | |
) -> std::fmt::Result { | ||
match t { | ||
DisplayFormatType::Default => { | ||
let files: Vec<_> = self | ||
.partitions | ||
.iter() | ||
.map(|pp| format!("{}", pp.file_partition)) | ||
.collect(); | ||
|
||
write!( | ||
f, | ||
"ParquetExec: batch_size={}, limit={:?}, partitions=[{}]", | ||
"ParquetExec: batch_size={}, limit={:?}, partitions={}", | ||
self.batch_size, | ||
self.limit, | ||
files.join(", ") | ||
super::FileGroupsDisplay(&self.file_groups) | ||
) | ||
} | ||
} | ||
|
@@ -493,7 +444,7 @@ fn build_row_group_predicate( | |
fn read_partition( | ||
object_store: &dyn ObjectStore, | ||
partition_index: usize, | ||
partition: ParquetPartition, | ||
partition: Vec<PartitionedFile>, | ||
metrics: ExecutionPlanMetricsSet, | ||
projection: &[usize], | ||
predicate_builder: &Option<PruningPredicate>, | ||
|
@@ -502,8 +453,7 @@ fn read_partition( | |
limit: Option<usize>, | ||
) -> Result<()> { | ||
let mut total_rows = 0; | ||
let all_files = partition.file_partition.files; | ||
'outer: for partitioned_file in all_files { | ||
'outer: for partitioned_file in partition { | ||
let file_metrics = ParquetFileMetrics::new( | ||
partition_index, | ||
&*partitioned_file.file_meta.path(), | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
"output partition" is vague here.
file_group
, i.e.Vec<PartitionedFile>
, is the unit of parallelism and will be processed by one single executor/thread.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
are you suggesting us coming up with a name that's semantically closer to concurrency?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was referring to the
ExecutionPlan.output_partitioning()
. Let me change this for something slightly more explicit 🙂 (I'll try to update this later today)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes. I think we can rephrase this line to avoid the ambiguity
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Lgtm