Skip to content
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.

Sign up for GitHub

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

Jump to bottom

feat(rust,python): support horizontal concatenation of LazyFrames #13139

Merged
merged 17 commits into from
Dec 28, 2023
Merged

feat(rust,python): support horizontal concatenation of LazyFrames #13139

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
a175dbe
Initial outline of lazy horizontal concat
adamreeve Dec 14, 2023
9c85968
Implement parallel evaluation of inputs for hconcat
adamreeve Dec 18, 2023
d01b649
Support slice pushdown with horizontal concatenation
adamreeve Dec 19, 2023
4190cce
Update test for error
adamreeve Dec 20, 2023
d91bc5e
Add test for predicate handling with hconcat
adamreeve Dec 20, 2023
248f938
Add test with streaming and hconcat
adamreeve Dec 20, 2023
394c8e6
Make total_read_count check a debug assertion
adamreeve Dec 21, 2023
93a7d51
Fix handling of non-streamable nodes with multiple inputs
adamreeve Dec 21, 2023
db8863d
Merge branch 'main' into lazy-hconcat
adamreeve Dec 21, 2023
e74f9b8
Implement projection pushdown for horizontal concatenation
adamreeve Dec 19, 2023
a082896
Merge branch 'main' into lazy-hconcat
adamreeve Dec 22, 2023
bc7c1f3
Add comment about why we don't use par_iter for hconcat executor
adamreeve Dec 27, 2023
b320ef3
Use debug_assert_ne
adamreeve Dec 27, 2023
4eb1efe
Move test from Rust to Python
adamreeve Dec 27, 2023
e9bfb50
Introduce new StackFrame struct for constructing streaming pipelines
adamreeve Dec 27, 2023
b119d1e
Make shared merge_schemas function
adamreeve Dec 27, 2023
10b5a28
Merge branch 'main' into lazy-hconcat
adamreeve Dec 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 4 additions & 16 deletions crates/polars-lazy/src/dsl/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,31 +185,19 @@ pub fn concat_lf_horizontal<L: AsRef<[LazyFrame]>>(
opt_state.file_caching |= lf.opt_state.file_caching;
}

let schema_size = lfs
.iter()
.map(|lf| lf.schema().map(|schema| schema.len()))
.sum::<PolarsResult<_>>()?;
let mut column_names = PlHashSet::with_capacity(schema_size);
let mut combined_schema = Schema::with_capacity(schema_size);

let mut lps = Vec::with_capacity(lfs.len());
let mut schemas = Vec::with_capacity(lfs.len());

for lf in lfs.iter() {
let mut lf = lf.clone();
let schema = lf.schema()?;
schema.iter().try_for_each(|(name, dtype)| {
if !column_names.contains(name) {
column_names.insert(name.clone());
combined_schema.with_column(name.clone(), dtype.clone());
Ok(())
} else {
Err(polars_err!(Duplicate: "Column with name '{}' has more than one occurrence", name))
}
})?;
schemas.push(schema);
let lp = std::mem::take(&mut lf.logical_plan);
lps.push(lp);
}

let combined_schema = merge_schemas(&schemas)?;

let options = HConcatOptions {
parallel: args.parallel,
};
Expand Down
14 changes: 4 additions & 10 deletions crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/hconcat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ pub(super) fn process_hconcat(

for input in inputs.iter() {
let mut input_pushdown = Vec::new();
let input_schema = lp_arena.get(*input).schema(lp_arena);

for proj in remaining_projections.iter() {
let input_schema = lp_arena.get(*input).schema(lp_arena);
if check_input_node(*proj, input_schema.as_ref(), expr_arena) {
input_pushdown.push(*proj);
}
Expand All @@ -47,18 +47,12 @@ pub(super) fn process_hconcat(
)?;
}

let schema_size = inputs
.iter()
.map(|input| lp_arena.get(*input).schema(lp_arena).len())
.sum();
let mut new_schema = Schema::with_capacity(schema_size);
let mut schemas = Vec::with_capacity(inputs.len());
for input in inputs.iter() {
let schema = lp_arena.get(*input).schema(lp_arena);
schema.as_ref().iter().for_each(|(name, dtype)| {
new_schema.with_column(name.clone(), dtype.clone());
});
schemas.push(schema.as_ref().clone());
}

let new_schema = merge_schemas(&schemas)?;
Arc::new(new_schema)
};

Expand Down
18 changes: 18 additions & 0 deletions crates/polars-plan/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -370,6 +370,24 @@ pub(crate) fn aexprs_to_schema(
.collect()
}

/// Concatenate multiple schemas into one, disallowing duplicate field names
pub fn merge_schemas(schemas: &[SchemaRef]) -> PolarsResult<Schema> {
let schema_size = schemas.iter().map(|schema| schema.len()).sum();
let mut merged_schema = Schema::with_capacity(schema_size);

for schema in schemas {
schema.iter().try_for_each(|(name, dtype)| {
if merged_schema.with_column(name.clone(), dtype.clone()).is_none() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a blocker for this PR, but this can be done elegantly with ok_or_else(|| polars_err!(..).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FWIW, this is actually doing the opposite of ok_or_else, a Some result should be mapped to an Err and a None to Ok(()). I couldn't find a more concise way to express this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh right!

Ok(())
} else {
Err(polars_err!(Duplicate: "Column with name '{}' has more than one occurrence", name))
}
})?;
}

Ok(merged_schema)
}

pub fn combine_predicates_expr<I>(iter: I) -> Expr
where
I: Iterator<Item = Expr>,
Expand Down