This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 48416e57f4 Fix projection name with DataFrame::with_column and window
functions (#12000)
48416e57f4 is described below
commit 48416e57f4671286a292c565682feae53af88a69
Author: WeblWabl <[email protected]>
AuthorDate: Sat Aug 17 05:45:01 2024 -0500
Fix projection name with DataFrame::with_column and window functions
(#12000)
* fix/11982: resolves projection issue found in with_column window fn usage
Signed-off-by: Devan <[email protected]>
* fix/11982: resolves projection issue found in with_column window fn usage
Signed-off-by: Devan <[email protected]>
* fmt
Signed-off-by: Devan <[email protected]>
* fmt
Signed-off-by: Devan <[email protected]>
* refactor to get tests working
Signed-off-by: Devan <[email protected]>
* change test to use test harness
Signed-off-by: Devan <[email protected]>
* use row_number method and add comment about test
Signed-off-by: Devan <[email protected]>
* add back import
Signed-off-by: Devan <[email protected]>
---------
Signed-off-by: Devan <[email protected]>
---
datafusion/core/src/dataframe/mod.rs | 44 ++++++++++++++++++++++++++++++++----
1 file changed, 40 insertions(+), 4 deletions(-)
diff --git a/datafusion/core/src/dataframe/mod.rs
b/datafusion/core/src/dataframe/mod.rs
index 3705873ce3..760ebd7392 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -1441,14 +1441,18 @@ impl DataFrame {
/// ```
pub fn with_column(self, name: &str, expr: Expr) -> Result<DataFrame> {
let window_func_exprs = find_window_exprs(&[expr.clone()]);
- let plan = if window_func_exprs.is_empty() {
- self.plan
+
+ let (plan, mut col_exists, window_func) = if
window_func_exprs.is_empty() {
+ (self.plan, false, false)
} else {
- LogicalPlanBuilder::window_plan(self.plan, window_func_exprs)?
+ (
+ LogicalPlanBuilder::window_plan(self.plan, window_func_exprs)?,
+ true,
+ true,
+ )
};
let new_column = expr.alias(name);
- let mut col_exists = false;
let mut fields: Vec<Expr> = plan
.schema()
.iter()
@@ -1456,6 +1460,8 @@ impl DataFrame {
if field.name() == name {
col_exists = true;
new_column.clone()
+ } else if window_func && qualifier.is_none() {
+ col(Column::from((qualifier, field))).alias(name)
} else {
col(Column::from((qualifier, field)))
}
@@ -1704,6 +1710,7 @@ mod tests {
use datafusion_common::{Constraint, Constraints, ScalarValue};
use datafusion_common_runtime::SpawnedTask;
use datafusion_expr::expr::WindowFunction;
+ use datafusion_expr::window_function::row_number;
use datafusion_expr::{
cast, create_udf, expr, lit, BuiltInWindowFunction, ExprFunctionExt,
ScalarFunctionImplementation, Volatility, WindowFrame,
WindowFrameBound,
@@ -2956,6 +2963,35 @@ mod tests {
Ok(())
}
+ // Test issue: https://github.com/apache/datafusion/issues/11982
+ // Window function was creating unwanted projection when using
with_column() method.
+ #[tokio::test]
+ async fn test_window_function_with_column() -> Result<()> {
+ let df = test_table().await?.select_columns(&["c1", "c2", "c3"])?;
+ let ctx = SessionContext::new();
+ let df_impl = DataFrame::new(ctx.state(), df.plan.clone());
+ let func = row_number().alias("row_num");
+
+ // Should create an additional column with alias 'r' that has window
func results
+ let df = df_impl.with_column("r", func)?.limit(0, Some(2))?;
+ assert_eq!(4, df.schema().fields().len());
+
+ let df_results = df.clone().collect().await?;
+ assert_batches_sorted_eq!(
+ [
+ "+----+----+-----+---+",
+ "| c1 | c2 | c3 | r |",
+ "+----+----+-----+---+",
+ "| c | 2 | 1 | 1 |",
+ "| d | 5 | -40 | 2 |",
+ "+----+----+-----+---+",
+ ],
+ &df_results
+ );
+
+ Ok(())
+ }
+
// Test issue: https://github.com/apache/datafusion/issues/7790
// The join operation outputs two identical column names, but they belong
to different relations.
#[tokio::test]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
