ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all \u003cpre\u003e\u003ccode\u003e blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks"); } } catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); } })(); (function(){ try { var __m = "github.com"; var __re = new RegExp('^' + "github\\.com" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length \u003e 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions.#8155

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

Closed
jorgecarleitao wants to merge 2 commits into from
Closed

ARROW-9954: [Rust] [DataFusion] Made aggregates support the same signatures as functions. #8155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4348f10
Made aggregates support same signatures as functions.
jorgecarleitao Aug 31, 2020
1ea0094
Fixed error in computing signature types.
jorgecarleitao Sep 12, 2020
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
4 changes: 2 additions & 2 deletions rust/datafusion/src/execution/context.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -514,7 +514,7 @@ mod tests {

use super::*;
use crate::datasource::MemTable;
use crate::logical_plan::{aggregate_expr, col, create_udf};
use crate::logical_plan::{col, create_udf, sum};
use crate::physical_plan::functions::ScalarFunctionImplementation;
use crate::test;
use crate::variable::VarType;
Expand DownExpand Up@@ -941,7 +941,7 @@ mod tests {
]));

let plan = LogicalPlanBuilder::scan("default", "test", schema.as_ref(), None)?
.aggregate(vec![col("c1")], vec![aggregate_expr("SUM", col("c2"))])?
.aggregate(vec![col("c1")], vec![sum(col("c2"))])?
.project(vec![col("c1"), col("SUM(c2)").alias("total_salary")])?
.build()?;

Expand Down
94 changes: 33 additions & 61 deletions rust/datafusion/src/logical_plan/mod.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -32,7 +32,7 @@ use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
use crate::{
physical_plan::{
expressions::binary_operator_data_type, functions,
aggregates, expressions::binary_operator_data_type, functions,
type_coercion::can_coerce_from, udf::ScalarUDF,
},
sql::parser::FileType,
Expand DownExpand Up@@ -210,12 +210,12 @@ fn create_name(e: &Expr, input_schema: &Schema) -> Result<String> {
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::AggregateFunction { name, args, .. } => {
Expr::AggregateFunction { fun, args, .. } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", name, names.join(",")))
Ok(format!("{}({})", fun, names.join(",")))
}
other => Err(ExecutionError::NotImplemented(format!(
"Physical plan does not support logical expression {:?}",
Expand DownExpand Up@@ -290,7 +290,7 @@ pub enum Expr {
/// aggregate function
AggregateFunction {
/// Name of the function
name: String,
fun: aggregates::AggregateFunction,
/// List of expressions to feed to the functions as arguments
args: Vec<Expr>,
},
Expand DownExpand Up@@ -321,47 +321,12 @@ impl Expr {
.collect::<Result<Vec<_>>>()?;
functions::return_type(fun, &data_types)
}
Expr::AggregateFunction { name, args, .. } => {
match name.to_uppercase().as_str() {
"MIN" | "MAX" => args[0].get_type(schema),
"SUM" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64 => Ok(DataType::Int64),
DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64 => Ok(DataType::UInt64),
DataType::Float32 => Ok(DataType::Float32),
DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"SUM does not support {:?}",
other
))),
},
"AVG" => match args[0].get_type(schema)? {
DataType::Int8
| DataType::Int16
| DataType::Int32
| DataType::Int64
| DataType::UInt8
| DataType::UInt16
| DataType::UInt32
| DataType::UInt64
| DataType::Float32
| DataType::Float64 => Ok(DataType::Float64),
other => Err(ExecutionError::General(format!(
"AVG does not support {:?}",
other
))),
},
"COUNT" => Ok(DataType::UInt64),
other => Err(ExecutionError::General(format!(
"Invalid aggregate function '{:?}'",
other
))),
}
Expr::AggregateFunction { fun, args, .. } => {
let data_types = args
.iter()
.map(|e| e.get_type(schema))
.collect::<Result<Vec<_>>>()?;
aggregates::return_type(fun, &data_types)
}
Expr::Not(_) => Ok(DataType::Boolean),
Expr::IsNull(_) => Ok(DataType::Boolean),
Expand DownExpand Up@@ -573,27 +538,42 @@ pub fn col(name: &str) -> Expr {

/// Create an expression to represent the min() aggregate function
pub fn min(expr: Expr) -> Expr {
aggregate_expr("MIN", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Min,
args: vec![expr],
}
}

/// Create an expression to represent the max() aggregate function
pub fn max(expr: Expr) -> Expr {
aggregate_expr("MAX", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Max,
args: vec![expr],
}
}

/// Create an expression to represent the sum() aggregate function
pub fn sum(expr: Expr) -> Expr {
aggregate_expr("SUM", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Sum,
args: vec![expr],
}
}

/// Create an expression to represent the avg() aggregate function
pub fn avg(expr: Expr) -> Expr {
aggregate_expr("AVG", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Avg,
args: vec![expr],
}
}

/// Create an expression to represent the count() aggregate function
pub fn count(expr: Expr) -> Expr {
aggregate_expr("COUNT", expr)
Expr::AggregateFunction {
fun: aggregates::AggregateFunction::Count,
args: vec![expr],
}
}

/// Whether it can be represented as a literal expression
Expand DownExpand Up@@ -690,14 +670,6 @@ pub fn concat(args: Vec<Expr>) -> Expr {
}
}

/// Create an aggregate expression
pub fn aggregate_expr(name: &str, expr: Expr) -> Expr {
Expr::AggregateFunction {
name: name.to_owned(),
args: vec![expr],
}
}

/// Creates a new UDF with a specific signature and specific return type.
/// This is a helper function to create a new UDF.
/// The function `create_udf` returns a subset of all possible `ScalarFunction`:
Expand DownExpand Up@@ -767,8 +739,8 @@ impl fmt::Debug for Expr {

write!(f, ")")
}
Expr::AggregateFunction { name, ref args, .. } => {
write!(f, "{}(", name)?;
Expr::AggregateFunction { fun, ref args, .. } => {
write!(f, "{}(", fun)?;
for i in 0..args.len() {
if i > 0 {
write!(f, ", ")?;
Expand DownExpand Up@@ -1429,7 +1401,7 @@ mod tests {
)?
.aggregate(
vec![col("state")],
vec![aggregate_expr("SUM", col("salary")).alias("total_salary")],
vec![sum(col("salary")).alias("total_salary")],
)?
.project(vec![col("state"), col("total_salary")])?
.build()?;
Expand Down
14 changes: 4 additions & 10 deletions rust/datafusion/src/optimizer/filter_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -303,7 +303,7 @@ fn rewrite(expr: &Expr, projection: &HashMap<String, Expr>) -> Result<Expr> {
mod tests {
use super::*;
use crate::logical_plan::col;
use crate::logical_plan::{aggregate_expr, lit, Expr, LogicalPlanBuilder, Operator};
use crate::logical_plan::{lit, sum, Expr, LogicalPlanBuilder, Operator};
use crate::test::*;

fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
Expand DownExpand Up@@ -369,10 +369,7 @@ mod tests {
fn filter_move_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("total_salary")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("total_salary")])?
.filter(col("a").gt(lit(10i64)))?
.build()?;
// filter of key aggregation is commutative
Expand All@@ -388,10 +385,7 @@ mod tests {
fn filter_keep_agg() -> Result<()> {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.aggregate(
vec![col("a")],
vec![aggregate_expr("SUM", col("b")).alias("b")],
)?
.aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])?
.filter(col("b").gt(lit(10i64)))?
.build()?;
// filter of aggregate is after aggregation since they are non-commutative
Expand DownExpand Up@@ -508,7 +502,7 @@ mod tests {
let table_scan = test_table_scan()?;
let plan = LogicalPlanBuilder::from(&table_scan)
.project(vec![col("a").alias("b"), col("c")])?
.aggregate(vec![col("b")], vec![aggregate_expr("SUM", col("c"))])?
.aggregate(vec![col("b")], vec![sum(col("c"))])?
.filter(col("b").gt(lit(10i64)))?
.filter(col("SUM(c)").gt(lit(10i64)))?
.build()?;
Expand Down
2 changes: 1 addition & 1 deletion rust/datafusion/src/optimizer/projection_push_down.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -338,7 +338,7 @@ mod tests {

use super::*;
use crate::logical_plan::{col, lit};
use crate::logical_plan::{Expr, LogicalPlanBuilder};
use crate::logical_plan::{max, min, Expr, LogicalPlanBuilder};
use crate::test::*;
use arrow::datatypes::DataType;

Expand Down
4 changes: 2 additions & 2 deletions rust/datafusion/src/optimizer/utils.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -239,8 +239,8 @@ pub fn rewrite_expression(expr: &Expr, expressions: &Vec<Expr>) -> Result<Expr>
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::AggregateFunction { name, .. } => Ok(Expr::AggregateFunction {
name: name.clone(),
Expr::AggregateFunction { fun, .. } => Ok(Expr::AggregateFunction {
fun: fun.clone(),
args: expressions.clone(),
}),
Expr::Cast { data_type, .. } => Ok(Expr::Cast {
Expand Down
Loading