From 2cdc453a16dc2aa561dffaf5e7ab8ce5c4d917d0 Mon Sep 17 00:00:00 2001 From: silezhou Date: Wed, 12 Mar 2025 23:23:14 +0800 Subject: [PATCH 01/10] chore(deps): Update sqlparser to 0.55.0 --- Cargo.toml | 2 +- datafusion/expr/src/window_frame.rs | 27 ++- datafusion/sql/src/expr/function.rs | 6 +- datafusion/sql/src/expr/identifier.rs | 36 ++-- datafusion/sql/src/expr/mod.rs | 25 +-- datafusion/sql/src/expr/order_by.rs | 12 +- datafusion/sql/src/expr/unary_op.rs | 9 +- datafusion/sql/src/expr/value.rs | 14 +- datafusion/sql/src/parser.rs | 43 +++-- datafusion/sql/src/planner.rs | 181 +++++++++---------- datafusion/sql/src/query.rs | 27 ++- datafusion/sql/src/relation/join.rs | 8 +- datafusion/sql/src/relation/mod.rs | 3 +- datafusion/sql/src/resolve.rs | 4 +- datafusion/sql/src/select.rs | 20 ++- datafusion/sql/src/set_expr.rs | 2 +- datafusion/sql/src/statement.rs | 90 ++++++++-- datafusion/sql/src/unparser/ast.rs | 36 ++-- datafusion/sql/src/unparser/dialect.rs | 6 +- datafusion/sql/src/unparser/expr.rs | 203 ++++++++++++---------- datafusion/sql/src/unparser/plan.rs | 20 ++- datafusion/sql/src/unparser/utils.rs | 4 +- datafusion/sql/tests/cases/plan_to_sql.rs | 22 +-- 23 files changed, 465 insertions(+), 335 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index efdcb1082b1d..40064d045382 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -151,7 +151,7 @@ recursive = "0.1.1" regex = "1.8" rstest = "0.24.0" serde_json = "1" -sqlparser = { version = "0.54.0", features = ["visitor"] } +sqlparser = { version = "0.55.0", features = ["visitor"] } tempfile = "3" tokio = { version = "1.43", features = ["macros", "rt", "sync"] } url = "2.5.4" diff --git a/datafusion/expr/src/window_frame.rs b/datafusion/expr/src/window_frame.rs index 82b33650523b..8771b25137cf 100644 --- a/datafusion/expr/src/window_frame.rs +++ b/datafusion/expr/src/window_frame.rs @@ -29,7 +29,7 @@ use std::fmt::{self, Formatter}; use std::hash::Hash; use datafusion_common::{plan_err, sql_err, DataFusionError, Result, ScalarValue}; -use sqlparser::ast; +use sqlparser::ast::{self, ValueWithSpan}; use sqlparser::parser::ParserError::ParserError; /// The frame specification determines which output rows are read by an aggregate @@ -368,7 +368,7 @@ fn convert_frame_bound_to_scalar_value( match units { // For ROWS and GROUPS we are sure that the ScalarValue must be a non-negative integer ... ast::WindowFrameUnits::Rows | ast::WindowFrameUnits::Groups => match v { - ast::Expr::Value(ast::Value::Number(value, false)) => { + ast::Expr::Value(ValueWithSpan{value: ast::Value::Number(value, false), span: _}) => { Ok(ScalarValue::try_from_string(value, &DataType::UInt64)?) }, ast::Expr::Interval(ast::Interval { @@ -379,7 +379,7 @@ fn convert_frame_bound_to_scalar_value( fractional_seconds_precision: None, }) => { let value = match *value { - ast::Expr::Value(ast::Value::SingleQuotedString(item)) => item, + ast::Expr::Value(ValueWithSpan{value: ast::Value::SingleQuotedString(item), span: _}) => item, e => { return sql_err!(ParserError(format!( "INTERVAL expression cannot be {e:?}" @@ -395,14 +395,14 @@ fn convert_frame_bound_to_scalar_value( // ... instead for RANGE it could be anything depending on the type of the ORDER BY clause, // so we use a ScalarValue::Utf8. ast::WindowFrameUnits::Range => Ok(ScalarValue::Utf8(Some(match v { - ast::Expr::Value(ast::Value::Number(value, false)) => value, + ast::Expr::Value(ValueWithSpan{value: ast::Value::Number(value, false), span: _}) => value, ast::Expr::Interval(ast::Interval { value, leading_field, .. }) => { let result = match *value { - ast::Expr::Value(ast::Value::SingleQuotedString(item)) => item, + ast::Expr::Value(ValueWithSpan{value: ast::Value::SingleQuotedString(item), span: _}) => item, e => { return sql_err!(ParserError(format!( "INTERVAL expression cannot be {e:?}" @@ -514,10 +514,10 @@ mod tests { let window_frame = ast::WindowFrame { units: ast::WindowFrameUnits::Rows, start_bound: ast::WindowFrameBound::Preceding(Some(Box::new( - ast::Expr::Value(ast::Value::Number("2".to_string(), false)), + ast::Expr::value(ast::Value::Number("2".to_string(), false)), ))), end_bound: Some(ast::WindowFrameBound::Preceding(Some(Box::new( - ast::Expr::Value(ast::Value::Number("1".to_string(), false)), + ast::Expr::value(ast::Value::Number("1".to_string(), false)), )))), }; @@ -575,10 +575,9 @@ mod tests { test_bound!(Range, None, ScalarValue::Null); // Number - let number = Some(Box::new(ast::Expr::Value(ast::Value::Number( - "42".to_string(), - false, - )))); + let number = Some(Box::new(ast::Expr::Value( + ast::Value::Number("42".to_string(), false).into(), + ))); test_bound!(Rows, number.clone(), ScalarValue::UInt64(Some(42))); test_bound!(Groups, number.clone(), ScalarValue::UInt64(Some(42))); test_bound!( @@ -589,9 +588,9 @@ mod tests { // Interval let number = Some(Box::new(ast::Expr::Interval(ast::Interval { - value: Box::new(ast::Expr::Value(ast::Value::SingleQuotedString( - "1".to_string(), - ))), + value: Box::new(ast::Expr::Value( + ast::Value::SingleQuotedString("1".to_string()).into(), + )), leading_field: Some(ast::DateTimeField::Day), fractional_seconds_precision: None, last_field: None, diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index cdf61183eb3d..9c74a0371449 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -31,7 +31,7 @@ use datafusion_expr::{ use sqlparser::ast::{ DuplicateTreatment, Expr as SQLExpr, Function as SQLFunction, FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments, - NullTreatment, ObjectName, OrderByExpr, WindowType, + NullTreatment, ObjectName, OrderByExpr, Spanned, WindowType, }; /// Suggest a valid function based on an invalid input function name @@ -217,13 +217,13 @@ impl SqlToRel<'_, S> { // it shouldn't have ordering requirement as function argument // required ordering should be defined in OVER clause. let is_function_window = over.is_some(); - let sql_parser_span = name.0[0].span; + let sql_parser_span = name.0[0].span(); let name = if name.0.len() > 1 { // DF doesn't handle compound identifiers // (e.g. "foo.bar") for function names yet name.to_string() } else { - crate::utils::normalize_ident(name.0[0].clone()) + crate::utils::normalize_ident(name.0[0].as_ident().unwrap().clone()) }; if name.eq("make_map") { diff --git a/datafusion/sql/src/expr/identifier.rs b/datafusion/sql/src/expr/identifier.rs index 7d358d0b6624..7c276ce53e35 100644 --- a/datafusion/sql/src/expr/identifier.rs +++ b/datafusion/sql/src/expr/identifier.rs @@ -22,7 +22,7 @@ use datafusion_common::{ }; use datafusion_expr::planner::PlannerResult; use datafusion_expr::{Case, Expr}; -use sqlparser::ast::{Expr as SQLExpr, Ident}; +use sqlparser::ast::{CaseWhen, Expr as SQLExpr, Ident}; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use datafusion_expr::UNNAMED_TABLE; @@ -216,8 +216,7 @@ impl SqlToRel<'_, S> { pub(super) fn sql_case_identifier_to_expr( &self, operand: Option>, - conditions: Vec, - results: Vec, + conditions: Vec, else_result: Option>, schema: &DFSchema, planner_context: &mut PlannerContext, @@ -231,13 +230,22 @@ impl SqlToRel<'_, S> { } else { None }; - let when_expr = conditions + let when_then_expr = conditions .into_iter() - .map(|e| self.sql_expr_to_logical_expr(e, schema, planner_context)) - .collect::>>()?; - let then_expr = results - .into_iter() - .map(|e| self.sql_expr_to_logical_expr(e, schema, planner_context)) + .map(|e| { + Ok(( + Box::new(self.sql_expr_to_logical_expr( + e.condition, + schema, + planner_context, + )?), + Box::new(self.sql_expr_to_logical_expr( + e.result, + schema, + planner_context, + )?), + )) + }) .collect::>>()?; let else_expr = if let Some(e) = else_result { Some(Box::new(self.sql_expr_to_logical_expr( @@ -249,15 +257,7 @@ impl SqlToRel<'_, S> { None }; - Ok(Expr::Case(Case::new( - expr, - when_expr - .iter() - .zip(then_expr.iter()) - .map(|(w, t)| (Box::new(w.to_owned()), Box::new(t.to_owned()))) - .collect(), - else_expr, - ))) + Ok(Expr::Case(Case::new(expr, when_then_expr, else_expr))) } } diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index fa2619111e7e..596be3527b22 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -22,7 +22,7 @@ use datafusion_expr::planner::{ use sqlparser::ast::{ AccessExpr, BinaryOperator, CastFormat, CastKind, DataType as SQLDataType, DictionaryField, Expr as SQLExpr, ExprWithAlias as SQLExprWithAlias, MapEntry, - StructField, Subscript, TrimWhereField, Value, + StructField, Subscript, TrimWhereField, Value, ValueWithSpan, }; use datafusion_common::{ @@ -211,7 +211,7 @@ impl SqlToRel<'_, S> { // more context. match sql { SQLExpr::Value(value) => { - self.parse_value(value, planner_context.prepare_param_data_types()) + self.parse_value(value.into(), planner_context.prepare_param_data_types()) } SQLExpr::Extract { field, expr, .. } => { let mut extract_args = vec![ @@ -253,12 +253,10 @@ impl SqlToRel<'_, S> { SQLExpr::Case { operand, conditions, - results, else_result, } => self.sql_case_identifier_to_expr( operand, conditions, - results, else_result, schema, planner_context, @@ -292,7 +290,7 @@ impl SqlToRel<'_, S> { } SQLExpr::TypedString { data_type, value } => Ok(Expr::Cast(Cast::new( - Box::new(lit(value)), + Box::new(lit(value.into_string().unwrap())), self.convert_data_type(&data_type)?, ))), @@ -544,9 +542,10 @@ impl SqlToRel<'_, S> { planner_context, )?), match *time_zone { - SQLExpr::Value(Value::SingleQuotedString(s)) => { - DataType::Timestamp(TimeUnit::Nanosecond, Some(s.into())) - } + SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(s), + span: _, + }) => DataType::Timestamp(TimeUnit::Nanosecond, Some(s.into())), _ => { return not_impl_err!( "Unsupported ast node in sqltorel: {time_zone:?}" @@ -997,10 +996,12 @@ impl SqlToRel<'_, S> { Subscript::Index { index } => { // index can be a name, in which case it is a named field access match index { - SQLExpr::Value( - Value::SingleQuotedString(s) - | Value::DoubleQuotedString(s), - ) => Ok(Some(GetFieldAccess::NamedStructField { + SQLExpr::Value(ValueWithSpan { + value: + Value::SingleQuotedString(s) + | Value::DoubleQuotedString(s), + span: _, + }) => Ok(Some(GetFieldAccess::NamedStructField { name: ScalarValue::from(s), })), SQLExpr::JsonAccess { .. } => { diff --git a/datafusion/sql/src/expr/order_by.rs b/datafusion/sql/src/expr/order_by.rs index b7ed04326f40..cce3f3004809 100644 --- a/datafusion/sql/src/expr/order_by.rs +++ b/datafusion/sql/src/expr/order_by.rs @@ -21,7 +21,9 @@ use datafusion_common::{ }; use datafusion_expr::expr::Sort; use datafusion_expr::{Expr, SortExpr}; -use sqlparser::ast::{Expr as SQLExpr, OrderByExpr, Value}; +use sqlparser::ast::{ + Expr as SQLExpr, OrderByExpr, OrderByOptions, Value, ValueWithSpan, +}; impl SqlToRel<'_, S> { /// Convert sql [OrderByExpr] to `Vec`. @@ -62,9 +64,8 @@ impl SqlToRel<'_, S> { let mut expr_vec = vec![]; for e in exprs { let OrderByExpr { - asc, expr, - nulls_first, + options: OrderByOptions { asc, nulls_first }, with_fill, } = e; @@ -73,7 +74,10 @@ impl SqlToRel<'_, S> { } let expr = match expr { - SQLExpr::Value(Value::Number(v, _)) if literal_to_column => { + SQLExpr::Value(ValueWithSpan { + value: Value::Number(v, _), + span: _, + }) if literal_to_column => { let field_index = v .parse::() .map_err(|err| plan_datafusion_err!("{}", err))?; diff --git a/datafusion/sql/src/expr/unary_op.rs b/datafusion/sql/src/expr/unary_op.rs index a4096ec2355b..b8379b1a518b 100644 --- a/datafusion/sql/src/expr/unary_op.rs +++ b/datafusion/sql/src/expr/unary_op.rs @@ -21,7 +21,7 @@ use datafusion_expr::{ type_coercion::{is_interval, is_timestamp}, Expr, ExprSchemable, }; -use sqlparser::ast::{Expr as SQLExpr, UnaryOperator, Value}; +use sqlparser::ast::{Expr as SQLExpr, UnaryOperator, Value, ValueWithSpan}; impl SqlToRel<'_, S> { pub(crate) fn parse_sql_unary_op( @@ -52,9 +52,10 @@ impl SqlToRel<'_, S> { match expr { // Optimization: if it's a number literal, we apply the negative operator // here directly to calculate the new literal. - SQLExpr::Value(Value::Number(n, _)) => { - self.parse_sql_number(&n, true) - } + SQLExpr::Value(ValueWithSpan { + value: Value::Number(n, _), + span: _, + }) => self.parse_sql_number(&n, true), SQLExpr::Interval(interval) => { self.sql_interval_to_expr(true, interval) } diff --git a/datafusion/sql/src/expr/value.rs b/datafusion/sql/src/expr/value.rs index 168348aee222..d53691ef05d1 100644 --- a/datafusion/sql/src/expr/value.rs +++ b/datafusion/sql/src/expr/value.rs @@ -32,7 +32,9 @@ use datafusion_expr::expr::{BinaryExpr, Placeholder}; use datafusion_expr::planner::PlannerResult; use datafusion_expr::{lit, Expr, Operator}; use log::debug; -use sqlparser::ast::{BinaryOperator, Expr as SQLExpr, Interval, UnaryOperator, Value}; +use sqlparser::ast::{ + BinaryOperator, Expr as SQLExpr, Interval, UnaryOperator, Value, ValueWithSpan, +}; use sqlparser::parser::ParserError::ParserError; use std::borrow::Cow; use std::cmp::Ordering; @@ -254,8 +256,14 @@ impl SqlToRel<'_, S> { fn interval_literal(interval_value: SQLExpr, negative: bool) -> Result { let s = match interval_value { - SQLExpr::Value(Value::SingleQuotedString(s) | Value::DoubleQuotedString(s)) => s, - SQLExpr::Value(Value::Number(ref v, long)) => { + SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(s) | Value::DoubleQuotedString(s), + span: _, + }) => s, + SQLExpr::Value(ValueWithSpan { + value: Value::Number(ref v, long), + span: _, + }) => { if long { return not_impl_err!( "Unsupported interval argument. Long number not supported: {interval_value:?}" diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs index 9725166b8ae0..1e6e5621b8fe 100644 --- a/datafusion/sql/src/parser.rs +++ b/datafusion/sql/src/parser.rs @@ -20,7 +20,7 @@ use std::collections::VecDeque; use std::fmt; -use sqlparser::ast::ExprWithAlias; +use sqlparser::ast::{ExprWithAlias, OrderByOptions}; use sqlparser::tokenizer::TokenWithSpan; use sqlparser::{ ast::{ @@ -628,8 +628,7 @@ impl<'a> DFParser<'a> { Ok(OrderByExpr { expr, - asc, - nulls_first, + options: OrderByOptions { asc, nulls_first }, with_fill: None, }) } @@ -676,11 +675,6 @@ impl<'a> DFParser<'a> { fn parse_column_def(&mut self) -> Result { let name = self.parser.parse_identifier()?; let data_type = self.parser.parse_data_type()?; - let collation = if self.parser.parse_keyword(Keyword::COLLATE) { - Some(self.parser.parse_object_name(false)?) - } else { - None - }; let mut options = vec![]; loop { if self.parser.parse_keyword(Keyword::CONSTRAINT) { @@ -702,7 +696,6 @@ impl<'a> DFParser<'a> { Ok(ColumnDef { name, data_type, - collation, options, }) } @@ -925,7 +918,6 @@ mod tests { span: Span::empty(), }, data_type, - collation: None, options: vec![], } } @@ -935,7 +927,7 @@ mod tests { // positive case let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv'"; let display = None; - let name = ObjectName(vec![Ident::from("t")]); + let name = ObjectName::from(vec![Ident::from("t")]); let expected = Statement::CreateExternalTable(CreateExternalTable { name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(display))], @@ -1233,8 +1225,7 @@ mod tests { quote_style: None, span: Span::empty(), }), - asc, - nulls_first, + options: OrderByOptions { asc, nulls_first }, with_fill: None, }]], if_not_exists: false, @@ -1265,8 +1256,10 @@ mod tests { quote_style: None, span: Span::empty(), }), - asc: Some(true), - nulls_first: None, + options: OrderByOptions { + asc: Some(true), + nulls_first: None, + }, with_fill: None, }, OrderByExpr { @@ -1275,8 +1268,10 @@ mod tests { quote_style: None, span: Span::empty(), }), - asc: Some(false), - nulls_first: Some(true), + options: OrderByOptions { + asc: Some(false), + nulls_first: Some(true), + }, with_fill: None, }, ]], @@ -1314,8 +1309,10 @@ mod tests { span: Span::empty(), })), }, - asc: Some(true), - nulls_first: None, + options: OrderByOptions { + asc: Some(true), + nulls_first: None, + }, with_fill: None, }]], if_not_exists: false, @@ -1361,8 +1358,10 @@ mod tests { span: Span::empty(), })), }, - asc: Some(true), - nulls_first: None, + options: OrderByOptions { + asc: Some(true), + nulls_first: None, + }, with_fill: None, }]], if_not_exists: true, @@ -1575,7 +1574,7 @@ mod tests { // For error cases, see: `copy.slt` fn object_name(name: &str) -> CopyToSource { - CopyToSource::Relation(ObjectName(vec![Ident::new(name)])) + CopyToSource::Relation(ObjectName::from(vec![Ident::new(name)])) } // Based on sqlparser-rs diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index bc7c2b7f4377..7af754ed5625 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -23,20 +23,18 @@ use std::vec; use arrow::datatypes::*; use datafusion_common::config::SqlParserOptions; use datafusion_common::error::add_possible_columns_to_diag; +use datafusion_common::TableReference; use datafusion_common::{ field_not_found, internal_err, plan_datafusion_err, DFSchemaRef, Diagnostic, SchemaError, }; -use sqlparser::ast::TimezoneInfo; -use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo}; -use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption}; -use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias}; - -use datafusion_common::TableReference; use datafusion_common::{not_impl_err, plan_err, DFSchema, DataFusionError, Result}; use datafusion_expr::logical_plan::{LogicalPlan, LogicalPlanBuilder}; use datafusion_expr::utils::find_column_exprs; use datafusion_expr::{col, Expr}; +use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo, TimezoneInfo}; +use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption}; +use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias}; use crate::utils::make_decimal_type; pub use datafusion_expr::planner::ContextProvider; @@ -550,97 +548,94 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16), SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => Ok(DataType::Int32), SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64), - SQLDataType::UnsignedTinyInt(_) => Ok(DataType::UInt8), - SQLDataType::UnsignedSmallInt(_) | SQLDataType::UnsignedInt2(_) => Ok(DataType::UInt16), - SQLDataType::UnsignedInt(_) | SQLDataType::UnsignedInteger(_) | SQLDataType::UnsignedInt4(_) => { - Ok(DataType::UInt32) - } + SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8), + SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => Ok(DataType::UInt16), + SQLDataType::IntUnsigned(_) | SQLDataType::IntegerUnsigned(_) | SQLDataType::Int4Unsigned(_) => { + Ok(DataType::UInt32) + } SQLDataType::Varchar(length) => { - match (length, self.options.support_varchar_with_length) { - (Some(_), false) => plan_err!("does not support Varchar with length, please set `support_varchar_with_length` to be true"), - _ => Ok(DataType::Utf8), - } - } - SQLDataType::UnsignedBigInt(_) | SQLDataType::UnsignedInt8(_) => Ok(DataType::UInt64), + match (length, self.options.support_varchar_with_length) { + (Some(_), false) => plan_err!("does not support Varchar with length, please set `support_varchar_with_length` to be true"), + _ => Ok(DataType::Utf8), + } + } + SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), SQLDataType::Float(_) => Ok(DataType::Float32), SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), SQLDataType::Double(ExactNumberInfo::None) | SQLDataType::DoublePrecision | SQLDataType::Float8 => Ok(DataType::Float64), SQLDataType::Double(ExactNumberInfo::Precision(_)|ExactNumberInfo::PrecisionAndScale(_, _)) => { - not_impl_err!("Unsupported SQL type (precision/scale not supported) {sql_type}") - } + not_impl_err!("Unsupported SQL type (precision/scale not supported) {sql_type}") + } SQLDataType::Char(_) - | SQLDataType::Text - | SQLDataType::String(_) => Ok(DataType::Utf8), + | SQLDataType::Text + | SQLDataType::String(_) => Ok(DataType::Utf8), SQLDataType::Timestamp(precision, tz_info) - if precision.is_none() || [0, 3, 6, 9].contains(&precision.unwrap()) => { - let tz = if matches!(tz_info, TimezoneInfo::Tz) - || matches!(tz_info, TimezoneInfo::WithTimeZone) - { - // Timestamp With Time Zone - // INPUT : [SQLDataType] TimestampTz + [Config] Time Zone - // OUTPUT: [ArrowDataType] Timestamp - self.context_provider.options().execution.time_zone.clone() - } else { - // Timestamp Without Time zone - None - }; - let precision = match precision { - Some(0) => TimeUnit::Second, - Some(3) => TimeUnit::Millisecond, - Some(6) => TimeUnit::Microsecond, - None | Some(9) => TimeUnit::Nanosecond, - _ => unreachable!(), - }; - Ok(DataType::Timestamp(precision, tz.map(Into::into))) - } + if precision.is_none() || [0, 3, 6, 9].contains(&precision.unwrap()) => { + let tz = if matches!(tz_info, TimezoneInfo::Tz) + || matches!(tz_info, TimezoneInfo::WithTimeZone) + { + // Timestamp With Time Zone + // INPUT : [SQLDataType] TimestampTz + [Config] Time Zone + // OUTPUT: [ArrowDataType] Timestamp + self.context_provider.options().execution.time_zone.clone() + } else { + // Timestamp Without Time zone + None + }; + let precision = match precision { + Some(0) => TimeUnit::Second, + Some(3) => TimeUnit::Millisecond, + Some(6) => TimeUnit::Microsecond, + None | Some(9) => TimeUnit::Nanosecond, + _ => unreachable!(), + }; + Ok(DataType::Timestamp(precision, tz.map(Into::into))) + } SQLDataType::Date => Ok(DataType::Date32), SQLDataType::Time(None, tz_info) => { - if matches!(tz_info, TimezoneInfo::None) - || matches!(tz_info, TimezoneInfo::WithoutTimeZone) - { - Ok(DataType::Time64(TimeUnit::Nanosecond)) - } else { - // We don't support TIMETZ and TIME WITH TIME ZONE for now - not_impl_err!( - "Unsupported SQL type {sql_type:?}" - ) - } - } + if matches!(tz_info, TimezoneInfo::None) + || matches!(tz_info, TimezoneInfo::WithoutTimeZone) + { + Ok(DataType::Time64(TimeUnit::Nanosecond)) + } else { + // We don't support TIMETZ and TIME WITH TIME ZONE for now + not_impl_err!( + "Unsupported SQL type {sql_type:?}" + ) + } + } SQLDataType::Numeric(exact_number_info) - | SQLDataType::Decimal(exact_number_info) => { - let (precision, scale) = match *exact_number_info { - ExactNumberInfo::None => (None, None), - ExactNumberInfo::Precision(precision) => (Some(precision), None), - ExactNumberInfo::PrecisionAndScale(precision, scale) => { - (Some(precision), Some(scale)) + | SQLDataType::Decimal(exact_number_info) => { + let (precision, scale) = match *exact_number_info { + ExactNumberInfo::None => (None, None), + ExactNumberInfo::Precision(precision) => (Some(precision), None), + ExactNumberInfo::PrecisionAndScale(precision, scale) => { + (Some(precision), Some(scale)) + } + }; + make_decimal_type(precision, scale) } - }; - make_decimal_type(precision, scale) - } SQLDataType::Bytea => Ok(DataType::Binary), SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), SQLDataType::Struct(fields, _) => { - let fields = fields - .iter() - .enumerate() - .map(|(idx, field)| { - let data_type = self.convert_data_type(&field.field_type)?; - let field_name = match &field.field_name{ - Some(ident) => ident.clone(), - None => Ident::new(format!("c{idx}")) - }; - Ok(Arc::new(Field::new( - self.ident_normalizer.normalize(field_name), - data_type, - true, - ))) - }) - .collect::>>()?; - Ok(DataType::Struct(Fields::from(fields))) - } - // Explicitly list all other types so that if sqlparser - // adds/changes the `SQLDataType` the compiler will tell us on upgrade - // and avoid bugs like https://github.com/apache/datafusion/issues/3059 + let fields = fields + .iter() + .enumerate() + .map(|(idx, field)| { + let data_type = self.convert_data_type(&field.field_type)?; + let field_name = match &field.field_name{ + Some(ident) => ident.clone(), + None => Ident::new(format!("c{idx}")) + }; + Ok(Arc::new(Field::new( + self.ident_normalizer.normalize(field_name), + data_type, + true, + ))) + }) + .collect::>>()?; + Ok(DataType::Struct(Fields::from(fields))) + } SQLDataType::Nvarchar(_) | SQLDataType::JSON | SQLDataType::Uuid @@ -654,7 +649,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::Enum(_, _) | SQLDataType::Set(_) | SQLDataType::MediumInt(_) - | SQLDataType::UnsignedMediumInt(_) + | SQLDataType::MediumIntUnsigned(_) | SQLDataType::Character(_) | SQLDataType::CharacterVarying(_) | SQLDataType::CharVarying(_) @@ -704,8 +699,16 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::LongText | SQLDataType::Bit(_) | SQLDataType::BitVarying(_) + | SQLDataType::Signed + | SQLDataType::SignedInteger + | SQLDataType::Unsigned + | SQLDataType::UnsignedInteger // BigQuery UDFs | SQLDataType::AnyType + // Postgres datatypes + | SQLDataType::Table(_) + | SQLDataType::VarBit(_) + | SQLDataType::GeometricType(_) => not_impl_err!( "Unsupported SQL type {sql_type:?}" ), @@ -738,7 +741,11 @@ pub fn object_name_to_table_reference( enable_normalization: bool, ) -> Result { // Use destructure to make it clear no fields on ObjectName are ignored - let ObjectName(idents) = object_name; + let ObjectName(object_name_parts) = object_name; + let idents = object_name_parts + .into_iter() + .map(|object_name_part| object_name_part.as_ident().unwrap().clone()) + .collect(); idents_to_table_reference(idents, enable_normalization) } @@ -828,11 +835,11 @@ pub fn object_name_to_qualifier( .iter() .rev() .zip(columns) - .map(|(ident, column_name)| { + .map(|(object_name_part, column_name)| { format!( r#"{} = '{}'"#, column_name, - normalizer.normalize(ident.clone()) + normalizer.normalize(object_name_part.as_ident().unwrap().clone()) ) }) .collect::>() diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 9d5a54d90b2c..cbbf83ff056e 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -23,11 +23,11 @@ use crate::stack::StackGuard; use datafusion_common::{not_impl_err, Constraints, DFSchema, Result}; use datafusion_expr::expr::Sort; use datafusion_expr::{ - CreateMemoryTable, DdlStatement, Distinct, LogicalPlan, LogicalPlanBuilder, + CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder, }; use sqlparser::ast::{ - Expr as SQLExpr, Offset as SQLOffset, OrderBy, OrderByExpr, Query, SelectInto, - SetExpr, + Expr as SQLExpr, Offset as SQLOffset, OrderBy, OrderByExpr, OrderByKind, Query, + SelectInto, SetExpr, }; impl SqlToRel<'_, S> { @@ -50,10 +50,8 @@ impl SqlToRel<'_, S> { match set_expr { SetExpr::Select(mut select) => { let select_into = select.into.take(); - // Order-by expressions may refer to columns in the `FROM` clause, - // so we need to process `SELECT` and `ORDER BY` together. - let oby_exprs = to_order_by_exprs(query.order_by)?; - let plan = self.select_to_plan(*select, oby_exprs, planner_context)?; + let plan = + self.select_to_plan(*select, query.order_by, planner_context)?; let plan = self.limit(plan, query.offset, query.limit, planner_context)?; // Process the `SELECT INTO` after `LIMIT`. @@ -153,12 +151,23 @@ impl SqlToRel<'_, S> { /// Returns the order by expressions from the query. fn to_order_by_exprs(order_by: Option) -> Result> { - let Some(OrderBy { exprs, interpolate }) = order_by else { + to_order_by_exprs_with_select(order_by, None) +} + +/// Returns the order by expressions from the query with the select expressions. +pub(crate) fn to_order_by_exprs_with_select( + order_by: Option, + _select_exprs: Option>, // TODO: ORDER BY ALL +) -> Result> { + let Some(OrderBy { kind, interpolate }) = order_by else { // If no order by, return an empty array. return Ok(vec![]); }; if let Some(_interpolate) = interpolate { return not_impl_err!("ORDER BY INTERPOLATE is not supported"); } - Ok(exprs) + match kind { + OrderByKind::All(_) => not_impl_err!("ORDER BY ALL is not supported"), + OrderByKind::Expressions(order_by_exprs) => Ok(order_by_exprs), + } } diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 88665401dc31..5c818fffa2e0 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -55,13 +55,13 @@ impl SqlToRel<'_, S> { self.create_relation(join.relation, planner_context)? }; match join.join_operator { - JoinOperator::LeftOuter(constraint) => { + JoinOperator::LeftOuter(constraint) | JoinOperator::Left(constraint) => { self.parse_join(left, right, constraint, JoinType::Left, planner_context) } - JoinOperator::RightOuter(constraint) => { + JoinOperator::RightOuter(constraint) | JoinOperator::Right(constraint) => { self.parse_join(left, right, constraint, JoinType::Right, planner_context) } - JoinOperator::Inner(constraint) => { + JoinOperator::Inner(constraint) | JoinOperator::Join(constraint) => { self.parse_join(left, right, constraint, JoinType::Inner, planner_context) } JoinOperator::LeftSemi(constraint) => self.parse_join( @@ -136,7 +136,7 @@ impl SqlToRel<'_, S> { ) } else { let id = object_names.swap_remove(0); - Ok(self.ident_normalizer.normalize(id)) + Ok(self.ident_normalizer.normalize(id.as_ident().unwrap().clone())) } }) .collect::>>()?; diff --git a/datafusion/sql/src/relation/mod.rs b/datafusion/sql/src/relation/mod.rs index 800dd151a124..43a8a02c6c05 100644 --- a/datafusion/sql/src/relation/mod.rs +++ b/datafusion/sql/src/relation/mod.rs @@ -43,7 +43,8 @@ impl SqlToRel<'_, S> { name, alias, args, .. } => { if let Some(func_args) = args { - let tbl_func_name = name.0.first().unwrap().value.to_string(); + let tbl_func_name = + name.0.first().unwrap().as_ident().unwrap().to_string(); let args = func_args .args .into_iter() diff --git a/datafusion/sql/src/resolve.rs b/datafusion/sql/src/resolve.rs index 88416dfe0324..96012a92c09a 100644 --- a/datafusion/sql/src/resolve.rs +++ b/datafusion/sql/src/resolve.rs @@ -81,7 +81,7 @@ impl Visitor for RelationVisitor { cte.visit(self); } self.ctes_in_scope - .push(ObjectName(vec![cte.alias.name.clone()])); + .push(ObjectName::from(vec![cte.alias.name.clone()])); } } ControlFlow::Continue(()) @@ -120,7 +120,7 @@ impl Visitor for RelationVisitor { ); if requires_information_schema { for s in INFORMATION_SCHEMA_TABLES { - self.relations.insert(ObjectName(vec![ + self.relations.insert(ObjectName::from(vec![ Ident::new(INFORMATION_SCHEMA), Ident::new(*s), ])); diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index e21def4c3941..8385682d42bc 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -19,6 +19,7 @@ use std::collections::HashSet; use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; +use crate::query::to_order_by_exprs_with_select; use crate::utils::{ check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs, rewrite_recursive_unnests_bottom_up, @@ -44,8 +45,8 @@ use datafusion_expr::{ use indexmap::IndexMap; use sqlparser::ast::{ - Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderByExpr, - WildcardAdditionalOptions, WindowType, + Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderBy, + SelectItemQualifiedWildcardKind, WildcardAdditionalOptions, WindowType, }; use sqlparser::ast::{NamedWindowDefinition, Select, SelectItem, TableWithJoins}; @@ -54,7 +55,7 @@ impl SqlToRel<'_, S> { pub(super) fn select_to_plan( &self, mut select: Select, - order_by: Vec, + query_order_by: Option, planner_context: &mut PlannerContext, ) -> Result { // Check for unsupported syntax first @@ -92,6 +93,9 @@ impl SqlToRel<'_, S> { planner_context, )?; + let order_by = + to_order_by_exprs_with_select(query_order_by, Some(select_exprs.clone()))?; + // Having and group by clause may reference aliases defined in select projection let projected_plan = self.project(base_plan.clone(), select_exprs.clone())?; @@ -639,6 +643,16 @@ impl SqlToRel<'_, S> { } SelectItem::QualifiedWildcard(object_name, options) => { Self::check_wildcard_options(&options)?; + let object_name = match object_name { + SelectItemQualifiedWildcardKind::ObjectName(object_name) => { + object_name + } + SelectItemQualifiedWildcardKind::Expr(_) => { + return plan_err!( + "Qualified wildcard with expression not supported" + ) + } + }; let qualifier = self.object_name_to_table_reference(object_name)?; let planned_options = self.plan_wildcard_options( plan, diff --git a/datafusion/sql/src/set_expr.rs b/datafusion/sql/src/set_expr.rs index a55b3b039087..272d6f874b4d 100644 --- a/datafusion/sql/src/set_expr.rs +++ b/datafusion/sql/src/set_expr.rs @@ -31,7 +31,7 @@ impl SqlToRel<'_, S> { ) -> Result { let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); match set_expr { - SetExpr::Select(s) => self.select_to_plan(*s, vec![], planner_context), + SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), SetExpr::SetOperation { op, diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index fbe6d6501c86..415c16bce3fc 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -57,6 +57,7 @@ use datafusion_expr::{ use sqlparser::ast::{ self, BeginTransactionKind, NullsDistinctOption, ShowStatementIn, ShowStatementOptions, SqliteOnConflict, TableObject, UpdateTableFromKind, + ValueWithSpan, }; use sqlparser::ast::{ Assignment, AssignmentTarget, ColumnDef, CreateIndex, CreateTable, @@ -75,7 +76,11 @@ fn object_name_to_string(object_name: &ObjectName) -> String { object_name .0 .iter() - .map(ident_to_string) + .map(|object_name_part| { + object_name_part + .as_ident() + .map_or_else(String::new, ident_to_string) + }) .collect::>() .join(".") } @@ -160,7 +165,8 @@ fn calc_inline_constraints_from_columns(columns: &[ColumnDef]) -> Vec {} + | ast::ColumnOption::Alias(_) + | ast::ColumnOption::Collation(_) => {} } } } @@ -273,6 +279,12 @@ impl SqlToRel<'_, S> { with_aggregation_policy, with_row_access_policy, with_tags, + iceberg, + external_volume, + base_location, + catalog, + catalog_sync, + storage_serialization_policy, }) if table_properties.is_empty() && with_options.is_empty() => { if temporary { return not_impl_err!("Temporary tables not supported")?; @@ -393,6 +405,24 @@ impl SqlToRel<'_, S> { if with_tags.is_some() { return not_impl_err!("With tags not supported")?; } + if iceberg { + return not_impl_err!("Iceberg not supported")?; + } + if external_volume.is_some() { + return not_impl_err!("External volume not supported")?; + } + if base_location.is_some() { + return not_impl_err!("Base location not supported")?; + } + if catalog.is_some() { + return not_impl_err!("Catalog not supported")?; + } + if catalog_sync.is_some() { + return not_impl_err!("Catalog sync not supported")?; + } + if storage_serialization_policy.is_some() { + return not_impl_err!("Storage serialization policy not supported")?; + } // Merge inline constraints and existing constraints let mut all_constraints = constraints; @@ -687,6 +717,8 @@ impl SqlToRel<'_, S> { // has_parentheses specifies the syntax, but the plan is the // same no matter the syntax used, so ignore it has_parentheses: _, + immediate, + into, } => { // `USING` is a MySQL-specific syntax and currently not supported. if !using.is_empty() { @@ -694,7 +726,14 @@ impl SqlToRel<'_, S> { "Execute statement with USING is not supported" ); } - + if immediate { + return not_impl_err!( + "Execute statement with IMMEDIATE is not supported" + ); + } + if !into.is_empty() { + return not_impl_err!("Execute statement with INTO is not supported"); + } let empty_schema = DFSchema::empty(); let parameters = parameters .into_iter() @@ -702,7 +741,7 @@ impl SqlToRel<'_, S> { .collect::>>()?; Ok(LogicalPlan::Statement(PlanStatement::Execute(Execute { - name: object_name_to_string(&name), + name: object_name_to_string(&name.unwrap()), parameters, }))) } @@ -905,8 +944,8 @@ impl SqlToRel<'_, S> { } => { let from = from.map(|update_table_from_kind| match update_table_from_kind { - UpdateTableFromKind::BeforeSet(from) => from, - UpdateTableFromKind::AfterSet(from) => from, + UpdateTableFromKind::BeforeSet(from) => from[0].clone(), + UpdateTableFromKind::AfterSet(from) => from[0].clone(), }); if returning.is_some() { plan_err!("Update-returning clause not yet supported")?; @@ -955,12 +994,28 @@ impl SqlToRel<'_, S> { begin: false, modifier, transaction, + statements, + exception_statements, + has_end_keyword, } => { if let Some(modifier) = modifier { return not_impl_err!( "Transaction modifier not supported: {modifier}" ); } + if !statements.is_empty() { + return not_impl_err!( + "Transaction with multiple statements not supported" + ); + } + if exception_statements.is_some() { + return not_impl_err!( + "Transaction with exception statements not supported" + ); + } + if has_end_keyword { + return not_impl_err!("Transaction with END keyword not supported"); + } self.validate_transaction_kind(transaction)?; let isolation_level: ast::TransactionIsolationLevel = modes .iter() @@ -1085,7 +1140,7 @@ impl SqlToRel<'_, S> { // At the moment functions can't be qualified `schema.name` let name = match &name.0[..] { [] => exec_err!("Function should have name")?, - [n] => n.value.clone(), + [n] => n.as_ident().unwrap().value.clone(), [..] => not_impl_err!("Qualified functions are not supported")?, }; // @@ -1143,7 +1198,7 @@ impl SqlToRel<'_, S> { // At the moment functions can't be qualified `schema.name` let name = match &desc.name.0[..] { [] => exec_err!("Function should have name")?, - [n] => n.value.clone(), + [n] => n.as_ident().unwrap().value.clone(), [..] => not_impl_err!("Qualified functions are not supported")?, }; let statement = DdlStatement::DropFunction(DropFunction { @@ -1344,8 +1399,9 @@ impl SqlToRel<'_, S> { planner_context, ) .unwrap(); - let asc = order_by_expr.asc.unwrap_or(true); - let nulls_first = order_by_expr.nulls_first.unwrap_or(!asc); + let asc = order_by_expr.options.asc.unwrap_or(true); + let nulls_first = + order_by_expr.options.nulls_first.unwrap_or(!asc); SortExpr::new(ordered_expr, asc, nulls_first) }) @@ -1612,7 +1668,7 @@ impl SqlToRel<'_, S> { variable_vec = variable_vec.split_at(variable_vec.len() - 1).0.to_vec(); } - let variable = object_name_to_string(&ObjectName(variable_vec)); + let variable = object_name_to_string(&ObjectName::from(variable_vec)); let base_query = format!("SELECT {columns} FROM information_schema.df_settings"); let query = if variable == "all" { // Add an ORDER BY so the output comes out in a consistent order @@ -1679,7 +1735,7 @@ impl SqlToRel<'_, S> { // Parse value string from Expr let value_string = match &value[0] { SQLExpr::Identifier(i) => ident_to_string(i), - SQLExpr::Value(v) => match crate::utils::value_to_string(v) { + SQLExpr::Value(v) => match crate::utils::value_to_string(&v.value) { None => { return plan_err!("Unsupported Value {}", value[0]); } @@ -1779,7 +1835,9 @@ impl SqlToRel<'_, S> { .0 .iter() .last() - .ok_or_else(|| plan_datafusion_err!("Empty column id"))?; + .ok_or_else(|| plan_datafusion_err!("Empty column id"))? + .as_ident() + .unwrap(); // Validate that the assignment target column exists table_schema.field_with_unqualified_name(&col_name.value)?; Ok((col_name.value.clone(), assign.value.clone())) @@ -1917,7 +1975,11 @@ impl SqlToRel<'_, S> { if let SetExpr::Values(ast::Values { rows, .. }) = (*source.body).clone() { for row in rows.iter() { for (idx, val) in row.iter().enumerate() { - if let SQLExpr::Value(Value::Placeholder(name)) = val { + if let SQLExpr::Value(ValueWithSpan { + value: Value::Placeholder(name), + span: _, + }) = val + { let name = name.replace('$', "").parse::().map_err(|_| { plan_datafusion_err!("Can't parse placeholder: {name}") diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index 6d77c01ea888..0a7f433b846f 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -17,14 +17,14 @@ use core::fmt; -use sqlparser::ast; use sqlparser::ast::helpers::attached_token::AttachedToken; +use sqlparser::ast::{self, OrderByKind, SelectFlavor}; #[derive(Clone)] pub struct QueryBuilder { with: Option, body: Option>, - order_by: Vec, + order_by_kind: Option, limit: Option, limit_by: Vec, offset: Option, @@ -46,8 +46,8 @@ impl QueryBuilder { pub fn take_body(&mut self) -> Option> { self.body.take() } - pub fn order_by(&mut self, value: Vec) -> &mut Self { - self.order_by = value; + pub fn order_by(&mut self, value: OrderByKind) -> &mut Self { + self.order_by_kind = Some(value); self } pub fn limit(&mut self, value: Option) -> &mut Self { @@ -75,14 +75,13 @@ impl QueryBuilder { self } pub fn build(&self) -> Result { - let order_by = if self.order_by.is_empty() { - None - } else { - Some(ast::OrderBy { - exprs: self.order_by.clone(), + let order_by = self + .order_by_kind + .as_ref() + .map(|order_by_kind| ast::OrderBy { + kind: order_by_kind.clone(), interpolate: None, - }) - }; + }); Ok(ast::Query { with: self.with.clone(), @@ -105,7 +104,7 @@ impl QueryBuilder { Self { with: Default::default(), body: Default::default(), - order_by: Default::default(), + order_by_kind: Default::default(), limit: Default::default(), limit_by: Default::default(), offset: Default::default(), @@ -138,6 +137,7 @@ pub struct SelectBuilder { named_window: Vec, qualify: Option, value_table_mode: Option, + flavor: Option, } #[allow(dead_code)] @@ -264,6 +264,10 @@ impl SelectBuilder { window_before_qualify: false, prewhere: None, select_token: AttachedToken::empty(), + flavor: match self.flavor { + Some(ref value) => value.clone(), + None => return Err(Into::into(UninitializedFieldError::from("flavor"))), + }, }) } fn create_empty() -> Self { @@ -283,6 +287,7 @@ impl SelectBuilder { named_window: Default::default(), qualify: Default::default(), value_table_mode: Default::default(), + flavor: Some(SelectFlavor::Standard), } } } @@ -422,6 +427,7 @@ pub struct TableRelationBuilder { with_hints: Vec, version: Option, partitions: Vec, + index_hints: Vec, } #[allow(dead_code)] @@ -450,6 +456,10 @@ impl TableRelationBuilder { self.partitions = value; self } + pub fn index_hints(&mut self, value: Vec) -> &mut Self { + self.index_hints = value; + self + } pub fn build(&self) -> Result { Ok(ast::TableFactor::Table { name: match self.name { @@ -467,6 +477,7 @@ impl TableRelationBuilder { with_ordinality: false, json_path: None, sample: None, + index_hints: self.index_hints.clone(), }) } fn create_empty() -> Self { @@ -477,6 +488,7 @@ impl TableRelationBuilder { with_hints: Default::default(), version: Default::default(), partitions: Default::default(), + index_hints: Default::default(), } } } diff --git a/datafusion/sql/src/unparser/dialect.rs b/datafusion/sql/src/unparser/dialect.rs index 399f0df0a699..0966c6e421ff 100644 --- a/datafusion/sql/src/unparser/dialect.rs +++ b/datafusion/sql/src/unparser/dialect.rs @@ -313,7 +313,7 @@ impl PostgreSqlDialect { } Ok(ast::Expr::Function(Function { - name: ObjectName(vec![Ident { + name: ObjectName::from(vec![Ident { value: func_name.to_string(), quote_style: None, span: Span::empty(), @@ -421,11 +421,11 @@ impl Dialect for MySqlDialect { } fn int64_cast_dtype(&self) -> ast::DataType { - ast::DataType::Custom(ObjectName(vec![Ident::new("SIGNED")]), vec![]) + ast::DataType::Custom(ObjectName::from(vec![Ident::new("SIGNED")]), vec![]) } fn int32_cast_dtype(&self) -> ast::DataType { - ast::DataType::Custom(ObjectName(vec![Ident::new("SIGNED")]), vec![]) + ast::DataType::Custom(ObjectName::from(vec![Ident::new("SIGNED")]), vec![]) } fn timestamp_cast_dtype( diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index d051cb78a8d5..21a78b5a569d 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -18,8 +18,8 @@ use datafusion_expr::expr::{AggregateFunctionParams, Unnest, WindowFunctionParams}; use sqlparser::ast::Value::SingleQuotedString; use sqlparser::ast::{ - self, Array, BinaryOperator, Expr as AstExpr, Function, Ident, Interval, ObjectName, - Subscript, TimezoneInfo, UnaryOperator, + self, Array, BinaryOperator, CaseWhen, Expr as AstExpr, Function, Ident, Interval, + ObjectName, OrderByOptions, Subscript, TimezoneInfo, UnaryOperator, ValueWithSpan, }; use std::sync::Arc; use std::vec; @@ -154,12 +154,14 @@ impl Unparser<'_> { }) => { let conditions = when_then_expr .iter() - .map(|(w, _)| self.expr_to_sql_inner(w)) - .collect::>>()?; - let results = when_then_expr - .iter() - .map(|(_, t)| self.expr_to_sql_inner(t)) - .collect::>>()?; + .map(|(cond, result)| { + Ok(CaseWhen { + condition: self.expr_to_sql_inner(cond)?, + result: self.expr_to_sql_inner(result)?, + }) + }) + .collect::>>()?; + let operand = match expr.as_ref() { Some(e) => match self.expr_to_sql_inner(e) { Ok(sql_expr) => Some(Box::new(sql_expr)), @@ -178,7 +180,6 @@ impl Unparser<'_> { Ok(ast::Expr::Case { operand, conditions, - results, else_result, }) } @@ -247,7 +248,7 @@ impl Unparser<'_> { })); Ok(ast::Expr::Function(Function { - name: ObjectName(vec![Ident { + name: ObjectName::from(vec![Ident { value: func_name.to_string(), quote_style: None, span: Span::empty(), @@ -300,7 +301,7 @@ impl Unparser<'_> { None => None, }; Ok(ast::Expr::Function(Function { - name: ObjectName(vec![Ident { + name: ObjectName::from(vec![Ident { value: func_name.to_string(), quote_style: None, span: Span::empty(), @@ -435,7 +436,7 @@ impl Unparser<'_> { let idents: Vec = qualifier.to_vec().into_iter().map(Ident::new).collect(); Ok(ast::Expr::QualifiedWildcard( - ObjectName(idents), + ObjectName::from(idents), attached_token, )) } else { @@ -477,7 +478,7 @@ impl Unparser<'_> { } }, Expr::Placeholder(p) => { - Ok(ast::Expr::Value(ast::Value::Placeholder(p.id.to_string()))) + Ok(ast::Expr::value(ast::Value::Placeholder(p.id.to_string()))) } Expr::OuterReferenceColumn(_, col) => self.col_to_sql(col), Expr::Unnest(unnest) => self.unnest_to_sql(unnest), @@ -507,7 +508,7 @@ impl Unparser<'_> { ) -> Result { let args = self.function_args_to_sql(args)?; Ok(ast::Expr::Function(Function { - name: ObjectName(vec![Ident { + name: ObjectName::from(vec![Ident { value: func_name.to_string(), quote_style: None, span: Span::empty(), @@ -659,8 +660,10 @@ impl Unparser<'_> { Ok(ast::OrderByExpr { expr: sql_parser_expr, - asc: Some(*asc), - nulls_first, + options: OrderByOptions { + asc: Some(*asc), + nulls_first, + }, with_fill: None, }) } @@ -700,7 +703,11 @@ impl Unparser<'_> { datafusion_expr::window_frame::WindowFrameBound::Preceding(val) => { Ok(ast::WindowFrameBound::Preceding({ let val = self.scalar_to_sql(val)?; - if let ast::Expr::Value(ast::Value::Null) = &val { + if let ast::Expr::Value(ValueWithSpan { + value: ast::Value::Null, + span: _, + }) = &val + { None } else { Some(Box::new(val)) @@ -710,7 +717,11 @@ impl Unparser<'_> { datafusion_expr::window_frame::WindowFrameBound::Following(val) => { Ok(ast::WindowFrameBound::Following({ let val = self.scalar_to_sql(val)?; - if let ast::Expr::Value(ast::Value::Null) = &val { + if let ast::Expr::Value(ValueWithSpan { + value: ast::Value::Null, + span: _, + }) = &val + { None } else { Some(Box::new(val)) @@ -997,7 +1008,7 @@ impl Unparser<'_> { Ok(ast::Expr::Cast { kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(SingleQuotedString(ts))), + expr: Box::new(ast::Expr::value(SingleQuotedString(ts))), data_type: self.dialect.timestamp_cast_dtype(&time_unit, &None), format: None, }) @@ -1019,7 +1030,7 @@ impl Unparser<'_> { .to_string(); Ok(ast::Expr::Cast { kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(SingleQuotedString(time))), + expr: Box::new(ast::Expr::value(SingleQuotedString(time))), data_type: ast::DataType::Time(None, TimezoneInfo::None), format: None, }) @@ -1054,102 +1065,102 @@ impl Unparser<'_> { /// For example ScalarValue::Date32(d) corresponds to the ast::Expr CAST('datestr' as DATE) fn scalar_to_sql(&self, v: &ScalarValue) -> Result { match v { - ScalarValue::Null => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Null => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Boolean(Some(b)) => { - Ok(ast::Expr::Value(ast::Value::Boolean(b.to_owned()))) + Ok(ast::Expr::value(ast::Value::Boolean(b.to_owned()))) } - ScalarValue::Boolean(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Boolean(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Float16(Some(f)) => { - Ok(ast::Expr::Value(ast::Value::Number(f.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(f.to_string(), false))) } - ScalarValue::Float16(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Float16(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Float32(Some(f)) => { let f_val = match f.fract() { 0.0 => format!("{:.1}", f), _ => format!("{}", f), }; - Ok(ast::Expr::Value(ast::Value::Number(f_val, false))) + Ok(ast::Expr::value(ast::Value::Number(f_val, false))) } - ScalarValue::Float32(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Float32(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Float64(Some(f)) => { let f_val = match f.fract() { 0.0 => format!("{:.1}", f), _ => format!("{}", f), }; - Ok(ast::Expr::Value(ast::Value::Number(f_val, false))) + Ok(ast::Expr::value(ast::Value::Number(f_val, false))) } - ScalarValue::Float64(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Float64(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Decimal128(Some(value), precision, scale) => { - Ok(ast::Expr::Value(ast::Value::Number( + Ok(ast::Expr::value(ast::Value::Number( Decimal128Type::format_decimal(*value, *precision, *scale), false, ))) } - ScalarValue::Decimal128(None, ..) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Decimal128(None, ..) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Decimal256(Some(value), precision, scale) => { - Ok(ast::Expr::Value(ast::Value::Number( + Ok(ast::Expr::value(ast::Value::Number( Decimal256Type::format_decimal(*value, *precision, *scale), false, ))) } - ScalarValue::Decimal256(None, ..) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Decimal256(None, ..) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Int8(Some(i)) => { - Ok(ast::Expr::Value(ast::Value::Number(i.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(i.to_string(), false))) } - ScalarValue::Int8(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Int8(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Int16(Some(i)) => { - Ok(ast::Expr::Value(ast::Value::Number(i.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(i.to_string(), false))) } - ScalarValue::Int16(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Int16(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Int32(Some(i)) => { - Ok(ast::Expr::Value(ast::Value::Number(i.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(i.to_string(), false))) } - ScalarValue::Int32(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Int32(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Int64(Some(i)) => { - Ok(ast::Expr::Value(ast::Value::Number(i.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(i.to_string(), false))) } - ScalarValue::Int64(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Int64(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::UInt8(Some(ui)) => { - Ok(ast::Expr::Value(ast::Value::Number(ui.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(ui.to_string(), false))) } - ScalarValue::UInt8(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::UInt8(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::UInt16(Some(ui)) => { - Ok(ast::Expr::Value(ast::Value::Number(ui.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(ui.to_string(), false))) } - ScalarValue::UInt16(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::UInt16(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::UInt32(Some(ui)) => { - Ok(ast::Expr::Value(ast::Value::Number(ui.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(ui.to_string(), false))) } - ScalarValue::UInt32(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::UInt32(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::UInt64(Some(ui)) => { - Ok(ast::Expr::Value(ast::Value::Number(ui.to_string(), false))) + Ok(ast::Expr::value(ast::Value::Number(ui.to_string(), false))) } - ScalarValue::UInt64(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::UInt64(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Utf8(Some(str)) => { - Ok(ast::Expr::Value(SingleQuotedString(str.to_string()))) + Ok(ast::Expr::value(SingleQuotedString(str.to_string()))) } - ScalarValue::Utf8(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Utf8(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Utf8View(Some(str)) => { - Ok(ast::Expr::Value(SingleQuotedString(str.to_string()))) + Ok(ast::Expr::value(SingleQuotedString(str.to_string()))) } - ScalarValue::Utf8View(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Utf8View(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::LargeUtf8(Some(str)) => { - Ok(ast::Expr::Value(SingleQuotedString(str.to_string()))) + Ok(ast::Expr::value(SingleQuotedString(str.to_string()))) } - ScalarValue::LargeUtf8(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::LargeUtf8(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Binary(Some(_)) => not_impl_err!("Unsupported scalar: {v:?}"), - ScalarValue::Binary(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Binary(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::BinaryView(Some(_)) => { not_impl_err!("Unsupported scalar: {v:?}") } - ScalarValue::BinaryView(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::BinaryView(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::FixedSizeBinary(..) => { not_impl_err!("Unsupported scalar: {v:?}") } ScalarValue::LargeBinary(Some(_)) => { not_impl_err!("Unsupported scalar: {v:?}") } - ScalarValue::LargeBinary(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::LargeBinary(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::FixedSizeList(a) => self.scalar_value_list_to_sql(a.values()), ScalarValue::List(a) => self.scalar_value_list_to_sql(a.values()), ScalarValue::LargeList(a) => self.scalar_value_list_to_sql(a.values()), @@ -1168,14 +1179,14 @@ impl Unparser<'_> { Ok(ast::Expr::Cast { kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(SingleQuotedString( + expr: Box::new(ast::Expr::value(SingleQuotedString( date.to_string(), ))), data_type: ast::DataType::Date, format: None, }) } - ScalarValue::Date32(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Date32(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Date64(Some(_)) => { let datetime = v .to_array()? @@ -1191,57 +1202,57 @@ impl Unparser<'_> { Ok(ast::Expr::Cast { kind: ast::CastKind::Cast, - expr: Box::new(ast::Expr::Value(SingleQuotedString( + expr: Box::new(ast::Expr::value(SingleQuotedString( datetime.to_string(), ))), data_type: self.ast_type_for_date64_in_cast(), format: None, }) } - ScalarValue::Date64(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Date64(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Time32Second(Some(_t)) => { self.handle_time::(v) } - ScalarValue::Time32Second(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Time32Second(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::Time32Millisecond(Some(_t)) => { self.handle_time::(v) } ScalarValue::Time32Millisecond(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::Time64Microsecond(Some(_t)) => { self.handle_time::(v) } ScalarValue::Time64Microsecond(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::Time64Nanosecond(Some(_t)) => { self.handle_time::(v) } - ScalarValue::Time64Nanosecond(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::Time64Nanosecond(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::TimestampSecond(Some(_ts), tz) => { self.handle_timestamp::(v, tz) } ScalarValue::TimestampSecond(None, _) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::TimestampMillisecond(Some(_ts), tz) => { self.handle_timestamp::(v, tz) } ScalarValue::TimestampMillisecond(None, _) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::TimestampMicrosecond(Some(_ts), tz) => { self.handle_timestamp::(v, tz) } ScalarValue::TimestampMicrosecond(None, _) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::TimestampNanosecond(Some(_ts), tz) => { self.handle_timestamp::(v, tz) } ScalarValue::TimestampNanosecond(None, _) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::IntervalYearMonth(Some(_)) | ScalarValue::IntervalDayTime(Some(_)) @@ -1249,33 +1260,33 @@ impl Unparser<'_> { self.interval_scalar_to_sql(v) } ScalarValue::IntervalYearMonth(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } - ScalarValue::IntervalDayTime(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::IntervalDayTime(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::IntervalMonthDayNano(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::DurationSecond(Some(_d)) => { not_impl_err!("Unsupported scalar: {v:?}") } - ScalarValue::DurationSecond(None) => Ok(ast::Expr::Value(ast::Value::Null)), + ScalarValue::DurationSecond(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::DurationMillisecond(Some(_d)) => { not_impl_err!("Unsupported scalar: {v:?}") } ScalarValue::DurationMillisecond(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::DurationMicrosecond(Some(_d)) => { not_impl_err!("Unsupported scalar: {v:?}") } ScalarValue::DurationMicrosecond(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::DurationNanosecond(Some(_d)) => { not_impl_err!("Unsupported scalar: {v:?}") } ScalarValue::DurationNanosecond(None) => { - Ok(ast::Expr::Value(ast::Value::Null)) + Ok(ast::Expr::value(ast::Value::Null)) } ScalarValue::Struct(_) => not_impl_err!("Unsupported scalar: {v:?}"), ScalarValue::Map(_) => not_impl_err!("Unsupported scalar: {v:?}"), @@ -1299,7 +1310,7 @@ impl Unparser<'_> { // MONTH only if months != 0 && days == 0 && microseconds == 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( months.to_string(), false, ))), @@ -1316,7 +1327,7 @@ impl Unparser<'_> { // DAY only if microseconds == 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( days.to_string(), false, ))), @@ -1334,7 +1345,7 @@ impl Unparser<'_> { if microseconds % 1_000_000 != 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( microseconds.to_string(), false, ))), @@ -1350,7 +1361,7 @@ impl Unparser<'_> { if secs % 60 != 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( secs.to_string(), false, ))), @@ -1366,7 +1377,7 @@ impl Unparser<'_> { if mins % 60 != 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( mins.to_string(), false, ))), @@ -1382,7 +1393,7 @@ impl Unparser<'_> { if hours % 24 != 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( hours.to_string(), false, ))), @@ -1397,7 +1408,7 @@ impl Unparser<'_> { let days = hours / 24; let interval = Interval { - value: Box::new(ast::Expr::Value(ast::Value::Number( + value: Box::new(ast::Expr::value(ast::Value::Number( days.to_string(), false, ))), @@ -1419,7 +1430,7 @@ impl Unparser<'_> { ); }; let interval = Interval { - value: Box::new(ast::Expr::Value(SingleQuotedString( + value: Box::new(ast::Expr::value(SingleQuotedString( result.to_uppercase(), ))), leading_field: None, @@ -1433,7 +1444,7 @@ impl Unparser<'_> { IntervalStyle::SQLStandard => match v { ScalarValue::IntervalYearMonth(Some(v)) => { let interval = Interval { - value: Box::new(ast::Expr::Value(SingleQuotedString( + value: Box::new(ast::Expr::value(SingleQuotedString( v.to_string(), ))), leading_field: Some(ast::DateTimeField::Month), @@ -1454,7 +1465,7 @@ impl Unparser<'_> { let millis = v.milliseconds % 1_000; let interval = Interval { - value: Box::new(ast::Expr::Value(SingleQuotedString(format!( + value: Box::new(ast::Expr::value(SingleQuotedString(format!( "{days} {hours}:{mins}:{secs}.{millis:3}" )))), leading_field: Some(ast::DateTimeField::Day), @@ -1467,7 +1478,7 @@ impl Unparser<'_> { ScalarValue::IntervalMonthDayNano(Some(v)) => { if v.months >= 0 && v.days == 0 && v.nanoseconds == 0 { let interval = Interval { - value: Box::new(ast::Expr::Value(SingleQuotedString( + value: Box::new(ast::Expr::value(SingleQuotedString( v.months.to_string(), ))), leading_field: Some(ast::DateTimeField::Month), @@ -1488,7 +1499,7 @@ impl Unparser<'_> { let millis = (v.nanoseconds % 1_000_000_000) / 1_000_000; let interval = Interval { - value: Box::new(ast::Expr::Value(SingleQuotedString( + value: Box::new(ast::Expr::value(SingleQuotedString( format!("{days} {hours}:{mins}:{secs}.{millis:03}"), ))), leading_field: Some(ast::DateTimeField::Day), @@ -1533,7 +1544,7 @@ impl Unparser<'_> { let args = self.function_args_to_sql(std::slice::from_ref(&unnest.expr))?; Ok(ast::Expr::Function(Function { - name: ObjectName(vec![Ident { + name: ObjectName::from(vec![Ident { value: "UNNEST".to_string(), quote_style: None, span: Span::empty(), @@ -1562,10 +1573,10 @@ impl Unparser<'_> { DataType::Int16 => Ok(ast::DataType::SmallInt(None)), DataType::Int32 => Ok(self.dialect.int32_cast_dtype()), DataType::Int64 => Ok(self.dialect.int64_cast_dtype()), - DataType::UInt8 => Ok(ast::DataType::UnsignedTinyInt(None)), - DataType::UInt16 => Ok(ast::DataType::UnsignedSmallInt(None)), - DataType::UInt32 => Ok(ast::DataType::UnsignedInteger(None)), - DataType::UInt64 => Ok(ast::DataType::UnsignedBigInt(None)), + DataType::UInt8 => Ok(ast::DataType::TinyIntUnsigned(None)), + DataType::UInt16 => Ok(ast::DataType::SmallIntUnsigned(None)), + DataType::UInt32 => Ok(ast::DataType::IntegerUnsigned(None)), + DataType::UInt64 => Ok(ast::DataType::BigIntUnsigned(None)), DataType::Float16 => { not_impl_err!("Unsupported DataType: conversion: {data_type:?}") } @@ -2538,7 +2549,7 @@ mod tests { let default_dialect = CustomDialectBuilder::new().build(); let mysql_dialect = CustomDialectBuilder::new() .with_int64_cast_dtype(ast::DataType::Custom( - ObjectName(vec![Ident::new("SIGNED")]), + ObjectName::from(vec![Ident::new("SIGNED")]), vec![], )) .build(); @@ -2566,7 +2577,7 @@ mod tests { let default_dialect = CustomDialectBuilder::new().build(); let mysql_dialect = CustomDialectBuilder::new() .with_int32_cast_dtype(ast::DataType::Custom( - ObjectName(vec![Ident::new("SIGNED")]), + ObjectName::from(vec![Ident::new("SIGNED")]), vec![], )) .build(); diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 0fa203c60b7b..447d42d1794d 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -49,7 +49,7 @@ use datafusion_expr::{ LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan, Unnest, UserDefinedLogicalNode, }; -use sqlparser::ast::{self, Ident, SetExpr, TableAliasColumnDef}; +use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef}; use std::sync::Arc; /// Convert a DataFusion [`LogicalPlan`] to [`ast::Statement`] @@ -356,7 +356,7 @@ impl Unparser<'_> { table_parts.push( self.new_ident_quoted_if_needs(scan.table_name.table().to_string()), ); - builder.name(ast::ObjectName(table_parts)); + builder.name(ast::ObjectName::from(table_parts)); relation.table(builder); Ok(()) @@ -471,7 +471,7 @@ impl Unparser<'_> { }; if let Some(fetch) = sort.fetch { - query_ref.limit(Some(ast::Expr::Value(ast::Value::Number( + query_ref.limit(Some(ast::Expr::value(ast::Value::Number( fetch.to_string(), false, )))); @@ -1048,11 +1048,13 @@ impl Unparser<'_> { } } - fn sorts_to_sql(&self, sort_exprs: &[SortExpr]) -> Result> { - sort_exprs - .iter() - .map(|sort_expr| self.sort_to_sql(sort_expr)) - .collect::>>() + fn sorts_to_sql(&self, sort_exprs: &[SortExpr]) -> Result { + Ok(OrderByKind::Expressions( + sort_exprs + .iter() + .map(|sort_expr| self.sort_to_sql(sort_expr)) + .collect::>>()?, + )) } fn join_operator_to_sql( @@ -1108,7 +1110,7 @@ impl Unparser<'_> { // this is represented as two columns like `[t1.id, t2.id]` // This code forms `id` (without relation name) let ident = self.new_ident_quoted_if_needs(left_name.to_string()); - object_names.push(ast::ObjectName(vec![ident])); + object_names.push(ast::ObjectName::from(vec![ident])); } // USING is only valid with matching column names; arbitrary expressions // are not allowed diff --git a/datafusion/sql/src/unparser/utils.rs b/datafusion/sql/src/unparser/utils.rs index f21fb2fcb49f..75038ccc4314 100644 --- a/datafusion/sql/src/unparser/utils.rs +++ b/datafusion/sql/src/unparser/utils.rs @@ -448,7 +448,7 @@ pub(crate) fn date_part_to_sql( }; return Ok(Some(ast::Expr::Function(ast::Function { - name: ast::ObjectName(vec![ast::Ident { + name: ast::ObjectName::from(vec![ast::Ident { value: "strftime".to_string(), quote_style: None, span: Span::empty(), @@ -457,7 +457,7 @@ pub(crate) fn date_part_to_sql( duplicate_treatment: None, args: vec![ ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr( - ast::Expr::Value(ast::Value::SingleQuotedString( + ast::Expr::value(ast::Value::SingleQuotedString( field.to_string(), )), )), diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 5af93a01e6c9..79f17a330bfe 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -368,7 +368,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { }, TestStatementWithDialect { sql: "SELECT j1_string from j1 join j2 on j1.j1_id = j2.j2_id order by j1_id", - expected: r#"SELECT j1.j1_string FROM j1 JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id ASC NULLS LAST"#, + expected: r#"SELECT j1.j1_string FROM j1 INNER JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id ASC NULLS LAST"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -393,7 +393,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { ) abc ORDER BY abc.j2_string", - expected: r#"SELECT abc.j1_string, abc.j2_string FROM (SELECT DISTINCT j1.j1_id, j1.j1_string, j2.j2_string FROM j1 JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, + expected: r#"SELECT abc.j1_string, abc.j2_string FROM (SELECT DISTINCT j1.j1_id, j1.j1_string, j2.j2_string FROM j1 INNER JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -410,7 +410,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { j1_id ) AS agg (id, string_count) ", - expected: r#"SELECT agg.string_count FROM (SELECT j1.j1_id, min(j2.j2_string) FROM j1 LEFT JOIN j2 ON (j1.j1_id = j2.j2_id) GROUP BY j1.j1_id) AS agg (id, string_count)"#, + expected: r#"SELECT agg.string_count FROM (SELECT j1.j1_id, min(j2.j2_string) FROM j1 LEFT OUTER JOIN j2 ON (j1.j1_id = j2.j2_id) GROUP BY j1.j1_id) AS agg (id, string_count)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -439,7 +439,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { ) abc ORDER BY abc.j2_string", - expected: r#"SELECT abc.j1_string, abc.j2_string FROM (SELECT j1.j1_id, j1.j1_string, j2.j2_string FROM j1 JOIN j2 ON (j1.j1_id = j2.j2_id) GROUP BY j1.j1_id, j1.j1_string, j2.j2_string ORDER BY j1.j1_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, + expected: r#"SELECT abc.j1_string, abc.j2_string FROM (SELECT j1.j1_id, j1.j1_string, j2.j2_string FROM j1 INNER JOIN j2 ON (j1.j1_id = j2.j2_id) GROUP BY j1.j1_id, j1.j1_string, j2.j2_string ORDER BY j1.j1_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -464,7 +464,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { ) abc ORDER BY j2_string", - expected: r#"SELECT abc.j1_string FROM (SELECT j1.j1_string, j2.j2_string FROM j1 JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, + expected: r#"SELECT abc.j1_string FROM (SELECT j1.j1_string, j2.j2_string FROM j1 INNER JOIN j2 ON (j1.j1_id = j2.j2_id) ORDER BY j1.j1_id DESC NULLS FIRST, j2.j2_id DESC NULLS FIRST LIMIT 10) AS abc ORDER BY abc.j2_string ASC NULLS LAST"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -562,7 +562,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) JOIN j1 ON u.c1 = j1.j1_id", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) JOIN j1 ON (u.c1 = j1.j1_id)"#, + expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -598,7 +598,7 @@ fn roundtrip_statement_with_dialect() -> Result<()> { }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) JOIN j1 ON u.c1 = j1.j1_id", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS u (c1) JOIN j1 ON (u.c1 = j1.j1_id)"#, + expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, @@ -1281,7 +1281,7 @@ fn test_join_with_table_scan_filters() -> Result<()> { let sql = plan_to_sql(&join_plan_with_filter)?; - let expected_sql = r#"SELECT * FROM left_table AS "left" JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND ("left"."name" LIKE 'some_name' AND (age > 10)))"#; + let expected_sql = r#"SELECT * FROM left_table AS "left" INNER JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND ("left"."name" LIKE 'some_name' AND (age > 10)))"#; assert_eq!(sql.to_string(), expected_sql); @@ -1296,7 +1296,7 @@ fn test_join_with_table_scan_filters() -> Result<()> { let sql = plan_to_sql(&join_plan_no_filter)?; - let expected_sql = r#"SELECT * FROM left_table AS "left" JOIN right_table ON "left".id = right_table.id AND ("left"."name" LIKE 'some_name' AND (age > 10))"#; + let expected_sql = r#"SELECT * FROM left_table AS "left" INNER JOIN right_table ON "left".id = right_table.id AND ("left"."name" LIKE 'some_name' AND (age > 10))"#; assert_eq!(sql.to_string(), expected_sql); @@ -1321,7 +1321,7 @@ fn test_join_with_table_scan_filters() -> Result<()> { let sql = plan_to_sql(&join_plan_multiple_filters)?; - let expected_sql = r#"SELECT * FROM left_table AS "left" JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND (("left"."name" LIKE 'some_name' AND (right_table."name" = 'before_join_filter_val')) AND (age > 10))) WHERE ("left"."name" = 'after_join_filter_val')"#; + let expected_sql = r#"SELECT * FROM left_table AS "left" INNER JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND (("left"."name" LIKE 'some_name' AND (right_table."name" = 'before_join_filter_val')) AND (age > 10))) WHERE ("left"."name" = 'after_join_filter_val')"#; assert_eq!(sql.to_string(), expected_sql); @@ -1351,7 +1351,7 @@ fn test_join_with_table_scan_filters() -> Result<()> { let sql = plan_to_sql(&join_plan_duplicated_filter)?; - let expected_sql = r#"SELECT * FROM left_table AS "left" JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND (("left"."name" LIKE 'some_name' AND (right_table.age > 10)) AND (right_table.age < 11)))"#; + let expected_sql = r#"SELECT * FROM left_table AS "left" INNER JOIN right_table ON "left".id = right_table.id AND (("left".id > 5) AND (("left"."name" LIKE 'some_name' AND (right_table.age > 10)) AND (right_table.age < 11)))"#; assert_eq!(sql.to_string(), expected_sql); From a25ccb3bdaf791fd3719bbdcae50c8460784e514 Mon Sep 17 00:00:00 2001 From: silezhou Date: Thu, 13 Mar 2025 12:28:24 +0000 Subject: [PATCH 02/10] merge upstream --- .github/workflows/dev_pr/labeler.yml | 17 +- .github/workflows/docs.yaml | 17 + .github/workflows/docs_pr.yaml | 30 +- .github/workflows/extended.yml | 8 +- .github/workflows/rust.yml | 137 +- Cargo.lock | 566 +- Cargo.toml | 70 +- benchmarks/README.md | 82 +- benchmarks/bench.sh | 162 +- benchmarks/src/cancellation.rs | 5 +- benchmarks/src/h2o.rs | 27 +- datafusion-cli/CONTRIBUTING.md | 75 + datafusion-cli/Cargo.toml | 10 +- datafusion-cli/src/exec.rs | 69 +- datafusion-cli/src/print_format.rs | 477 -- datafusion-cli/src/print_options.rs | 207 +- datafusion-cli/tests/cli_integration.rs | 129 +- .../tests/snapshots/aws_options.snap | 25 + .../snapshots/cli@load_local_csv.sql.snap | 26 + .../tests/snapshots/cli@load_s3_csv.sql.snap | 26 + .../tests/snapshots/cli@select.sql.snap | 23 + .../tests/snapshots/cli_format@automatic.snap | 21 + .../tests/snapshots/cli_format@csv.snap | 18 + .../tests/snapshots/cli_format@json.snap | 17 + .../tests/snapshots/cli_format@nd-json.snap | 17 + .../tests/snapshots/cli_format@table.snap | 21 + .../tests/snapshots/cli_format@tsv.snap | 18 + .../snapshots/cli_quick_test@backslash.snap | 17 + .../snapshots/cli_quick_test@batch_size.snap | 21 + .../tests/snapshots/cli_quick_test@files.snap | 19 + .../snapshots/cli_quick_test@statements.snap | 24 + .../{data/backslash.txt => sql/backslash.sql} | 0 .../tests/sql/integration/load_local_csv.sql | 6 + .../tests/sql/integration/load_s3_csv.sql | 5 + .../tests/{data/sql.txt => sql/select.sql} | 0 datafusion-examples/Cargo.toml | 2 +- .../examples/advanced_parquet_index.rs | 6 +- .../examples/csv_json_opener.rs | 6 +- datafusion-examples/examples/planner_api.rs | 20 +- datafusion-examples/examples/sql_dialect.rs | 8 +- datafusion/catalog-listing/src/helpers.rs | 11 + datafusion/catalog/src/session.rs | 14 +- datafusion/common/Cargo.toml | 4 +- datafusion/common/src/config.rs | 14 +- datafusion/common/src/dfschema.rs | 15 +- datafusion/common/src/scalar/mod.rs | 6 +- datafusion/core/Cargo.toml | 25 +- .../core/src/datasource/file_format/arrow.rs | 4 + .../core/src/datasource/file_format/avro.rs | 201 +- .../core/src/datasource/file_format/csv.rs | 884 +--- .../core/src/datasource/file_format/json.rs | 428 +- .../core/src/datasource/file_format/mod.rs | 362 +- .../src/datasource/file_format/options.rs | 15 +- .../src/datasource/file_format/parquet.rs | 1497 +----- .../core/src/datasource/listing/table.rs | 182 +- datafusion/core/src/datasource/memory.rs | 4 + datafusion/core/src/datasource/mod.rs | 278 +- .../datasource/physical_plan/arrow_file.rs | 5 +- .../core/src/datasource/physical_plan/avro.rs | 366 +- .../core/src/datasource/physical_plan/csv.rs | 856 +--- .../core/src/datasource/physical_plan/json.rs | 471 +- .../core/src/datasource/physical_plan/mod.rs | 144 +- .../{parquet/mod.rs => parquet.rs} | 613 +-- datafusion/core/src/datasource/statistics.rs | 75 - datafusion/core/src/datasource/view.rs | 11 +- datafusion/core/src/execution/context/csv.rs | 2 +- datafusion/core/src/execution/context/json.rs | 2 +- datafusion/core/src/execution/context/mod.rs | 4 +- .../core/src/execution/context/parquet.rs | 2 +- .../core/src/execution/session_state.rs | 247 +- .../src/execution/session_state_defaults.rs | 2 + datafusion/core/src/lib.rs | 28 +- datafusion/core/src/physical_planner.rs | 308 +- datafusion/core/src/test/mod.rs | 19 +- datafusion/core/src/test/object_store.rs | 132 +- datafusion/core/src/test_util/mod.rs | 42 +- datafusion/core/tests/core_integration.rs | 5 + .../core/tests/custom_sources_cases/mod.rs | 5 + .../provider_filter_pushdown.rs | 4 + .../tests/custom_sources_cases/statistics.rs | 4 + .../tests/dataframe/dataframe_functions.rs | 6 +- datafusion/core/tests/dataframe/mod.rs | 383 +- datafusion/core/tests/expr_api/mod.rs | 66 + .../aggregation_fuzzer/data_generator.rs | 23 +- .../fuzz_cases/aggregation_fuzzer/mod.rs | 20 + datafusion/core/tests/memory_limit/mod.rs | 25 + .../tests/parquet/external_access_plan.rs | 2 +- .../core/tests/parquet/schema_coercion.rs | 15 +- .../enforce_distribution.rs | 1989 +++++--- .../physical_optimizer/enforce_sorting.rs | 154 +- .../physical_optimizer/join_selection.rs | 9 + .../physical_optimizer/projection_pushdown.rs | 96 +- .../physical_optimizer/sanity_checker.rs | 10 +- .../tests/physical_optimizer/test_utils.rs | 87 +- datafusion/core/tests/serde/mod.rs | 34 + datafusion/core/tests/sql/explain_analyze.rs | 7 +- .../tests/user_defined/user_defined_plan.rs | 4 + .../user_defined_scalar_functions.rs | 10 +- datafusion/datasource-avro/Cargo.toml | 60 + datafusion/datasource-avro/LICENSE.txt | 1 + datafusion/datasource-avro/NOTICE.txt | 1 + datafusion/datasource-avro/README.md | 26 + .../src}/avro_to_arrow/arrow_array_reader.rs | 36 +- .../src}/avro_to_arrow/mod.rs | 22 +- .../src}/avro_to_arrow/reader.rs | 18 +- .../src}/avro_to_arrow/schema.rs | 10 +- datafusion/datasource-avro/src/file_format.rs | 160 + datafusion/datasource-avro/src/mod.rs | 30 + datafusion/datasource-avro/src/source.rs | 282 + datafusion/datasource-csv/Cargo.toml | 60 + datafusion/datasource-csv/LICENSE.txt | 1 + datafusion/datasource-csv/NOTICE.txt | 1 + datafusion/datasource-csv/README.md | 26 + datafusion/datasource-csv/src/file_format.rs | 740 +++ datafusion/datasource-csv/src/mod.rs | 38 + datafusion/datasource-csv/src/source.rs | 786 +++ datafusion/datasource-json/Cargo.toml | 56 + datafusion/datasource-json/LICENSE.txt | 1 + datafusion/datasource-json/NOTICE.txt | 1 + datafusion/datasource-json/README.md | 26 + datafusion/datasource-json/src/file_format.rs | 420 ++ datafusion/datasource-json/src/mod.rs | 21 + datafusion/datasource-json/src/source.rs | 440 ++ datafusion/datasource-parquet/Cargo.toml | 65 + datafusion/datasource-parquet/LICENSE.txt | 1 + datafusion/datasource-parquet/NOTICE.txt | 1 + datafusion/datasource-parquet/README.md | 26 + .../src}/access_plan.rs | 2 +- .../datasource-parquet/src/file_format.rs | 1403 +++++ .../src}/metrics.rs | 2 +- datafusion/datasource-parquet/src/mod.rs | 547 ++ .../src}/opener.rs | 18 +- .../src}/page_filter.rs | 2 +- .../src}/reader.rs | 4 +- .../src}/row_filter.rs | 14 +- .../src}/row_group_filter.rs | 9 +- .../src}/source.rs | 44 +- .../src}/writer.rs | 2 +- datafusion/datasource/Cargo.toml | 7 +- datafusion/datasource/src/decoder.rs | 191 + datafusion/datasource/src/display.rs | 4 +- datafusion/datasource/src/file.rs | 4 +- datafusion/datasource/src/file_scan_config.rs | 55 +- datafusion/datasource/src/memory.rs | 21 +- datafusion/datasource/src/mod.rs | 113 +- .../src}/schema_adapter.rs | 222 +- datafusion/datasource/src/source.rs | 31 +- datafusion/datasource/src/test_util.rs | 15 +- datafusion/expr-common/src/signature.rs | 5 + datafusion/expr/src/expr.rs | 37 +- datafusion/expr/src/expr_fn.rs | 4 + datafusion/expr/src/expr_rewriter/mod.rs | 1 + datafusion/expr/src/expr_schema.rs | 2 + datafusion/expr/src/logical_plan/builder.rs | 1 + datafusion/expr/src/logical_plan/extension.rs | 18 - datafusion/expr/src/logical_plan/plan.rs | 140 +- datafusion/expr/src/tree_node.rs | 4 + .../expr/src/type_coercion/functions.rs | 9 +- datafusion/expr/src/udf.rs | 48 +- datafusion/expr/src/utils.rs | 170 +- datafusion/ffi/Cargo.toml | 2 +- datafusion/ffi/src/execution_plan.rs | 24 +- datafusion/functions-aggregate/src/count.rs | 51 +- datafusion/functions-aggregate/src/min_max.rs | 2 +- datafusion/functions-aggregate/src/planner.rs | 2 + datafusion/functions-nested/benches/map.rs | 13 +- datafusion/functions-nested/src/extract.rs | 33 +- datafusion/functions-nested/src/lib.rs | 2 + datafusion/functions-nested/src/max.rs | 137 + datafusion/functions-nested/src/replace.rs | 7 +- datafusion/functions-nested/src/resize.rs | 32 +- datafusion/functions-nested/src/sort.rs | 51 +- datafusion/functions-nested/src/string.rs | 17 +- datafusion/functions-window/src/planner.rs | 2 + .../functions/benches/character_length.rs | 36 +- datafusion/functions/benches/chr.rs | 15 +- datafusion/functions/benches/cot.rs | 25 +- datafusion/functions/benches/date_bin.rs | 14 +- datafusion/functions/benches/date_trunc.rs | 15 +- datafusion/functions/benches/encoding.rs | 45 +- datafusion/functions/benches/isnan.rs | 25 +- datafusion/functions/benches/iszero.rs | 25 +- datafusion/functions/benches/make_date.rs | 37 +- datafusion/functions/benches/nullif.rs | 14 +- datafusion/functions/benches/pad.rs | 70 +- datafusion/functions/benches/random.rs | 25 +- datafusion/functions/benches/reverse.rs | 30 +- datafusion/functions/benches/signum.rs | 25 +- datafusion/functions/benches/strpos.rs | 31 +- datafusion/functions/benches/substr.rs | 66 +- datafusion/functions/benches/substr_index.rs | 12 +- datafusion/functions/benches/to_char.rs | 24 +- datafusion/functions/benches/to_timestamp.rs | 53 +- datafusion/functions/benches/trunc.rs | 25 +- .../functions/src/core/union_extract.rs | 4 +- datafusion/functions/src/datetime/to_char.rs | 58 +- datafusion/functions/src/regex/regexplike.rs | 27 +- datafusion/functions/src/string/bit_length.rs | 17 +- datafusion/functions/src/string/contains.rs | 85 +- datafusion/functions/src/string/ends_with.rs | 43 +- .../functions/src/string/levenshtein.rs | 150 +- datafusion/functions/src/string/lower.rs | 14 +- .../functions/src/string/octet_length.rs | 17 +- datafusion/functions/src/string/replace.rs | 81 +- datafusion/functions/src/string/upper.rs | 14 +- datafusion/functions/src/unicode/initcap.rs | 14 +- datafusion/functions/src/unicode/strpos.rs | 22 +- datafusion/macros/Cargo.toml | 2 +- .../src/analyzer/expand_wildcard_rule.rs | 332 -- .../src/analyzer/inline_table_scan.rs | 8 +- datafusion/optimizer/src/analyzer/mod.rs | 5 - .../optimizer/src/analyzer/type_coercion.rs | 3 + .../optimizer/src/common_subexpr_eliminate.rs | 2 + datafusion/optimizer/src/decorrelate.rs | 17 +- .../src/eliminate_group_by_constant.rs | 10 +- datafusion/optimizer/src/lib.rs | 3 - datafusion/optimizer/src/optimizer.rs | 59 +- datafusion/optimizer/src/push_down_filter.rs | 87 + .../simplify_expressions/expr_simplifier.rs | 115 +- .../optimizer/src/simplify_expressions/mod.rs | 1 + .../simplify_expressions/simplify_exprs.rs | 2 - .../unwrap_cast.rs} | 461 +- .../src/simplify_expressions/utils.rs | 19 + datafusion/optimizer/src/utils.rs | 40 - .../optimizer/tests/optimizer_integration.rs | 40 +- datafusion/physical-expr/src/aggregate.rs | 11 +- .../physical-expr/src/equivalence/mod.rs | 32 + .../src/equivalence/properties.rs | 4544 ----------------- .../src/equivalence/properties/dependency.rs | 1774 +++++++ .../src/equivalence/properties/joins.rs | 301 ++ .../src/equivalence/properties/mod.rs | 1639 ++++++ .../src/equivalence/properties/union.rs | 927 ++++ datafusion/physical-expr/src/lib.rs | 4 - .../physical-expr/src/scalar_function.rs | 36 +- .../src/enforce_sorting/sort_pushdown.rs | 162 +- .../physical-optimizer/src/optimizer.rs | 8 +- .../src/output_requirements.rs | 12 +- datafusion/physical-optimizer/src/utils.rs | 4 + .../physical-plan/src/aggregates/mod.rs | 118 +- datafusion/physical-plan/src/analyze.rs | 4 + .../physical-plan/src/coalesce_batches.rs | 4 + .../physical-plan/src/coalesce_partitions.rs | 4 + datafusion/physical-plan/src/common.rs | 2 +- datafusion/physical-plan/src/display.rs | 565 +- datafusion/physical-plan/src/empty.rs | 4 + .../physical-plan/src/execution_plan.rs | 2 +- datafusion/physical-plan/src/explain.rs | 4 + datafusion/physical-plan/src/filter.rs | 3 + datafusion/physical-plan/src/insert.rs | 1 + .../physical-plan/src/joins/cross_join.rs | 4 + .../physical-plan/src/joins/hash_join.rs | 13 + .../src/joins/nested_loop_join.rs | 7 + .../src/joins/sort_merge_join.rs | 19 +- .../src/joins/symmetric_hash_join.rs | 14 + datafusion/physical-plan/src/lib.rs | 3 +- datafusion/physical-plan/src/limit.rs | 10 + datafusion/physical-plan/src/memory.rs | 4 + .../physical-plan/src/placeholder_row.rs | 5 + datafusion/physical-plan/src/projection.rs | 11 + .../physical-plan/src/recursive_query.rs | 4 + datafusion/physical-plan/src/render_tree.rs | 230 + .../physical-plan/src/repartition/mod.rs | 12 + .../physical-plan/src/sorts/partial_sort.rs | 9 + datafusion/physical-plan/src/sorts/sort.rs | 49 +- .../src/sorts/sort_preserving_merge.rs | 19 + datafusion/physical-plan/src/spill.rs | 143 +- datafusion/physical-plan/src/stream.rs | 25 +- datafusion/physical-plan/src/streaming.rs | 12 + datafusion/physical-plan/src/test.rs | 12 +- datafusion/physical-plan/src/test/exec.rs | 24 + datafusion/physical-plan/src/union.rs | 8 + datafusion/physical-plan/src/unnest.rs | 3 + datafusion/physical-plan/src/values.rs | 4 + .../src/windows/bounded_window_agg_exec.rs | 11 + .../src/windows/window_agg_exec.rs | 8 + datafusion/physical-plan/src/work_table.rs | 3 + datafusion/proto-common/gen/src/main.rs | 3 - datafusion/proto/Cargo.toml | 1 + datafusion/proto/gen/src/main.rs | 4 - .../proto/src/logical_plan/from_proto.rs | 1 + datafusion/proto/src/logical_plan/mod.rs | 16 +- datafusion/proto/src/logical_plan/to_proto.rs | 1 + datafusion/proto/src/physical_plan/mod.rs | 25 +- .../tests/cases/roundtrip_physical_plan.rs | 113 +- datafusion/sql/src/expr/mod.rs | 2 + datafusion/sql/src/parser.rs | 165 +- datafusion/sql/src/planner.rs | 18 +- datafusion/sql/src/select.rs | 90 +- datafusion/sql/src/unparser/expr.rs | 23 +- datafusion/sql/src/unparser/plan.rs | 11 +- datafusion/sql/src/utils.rs | 2 + datafusion/sql/tests/cases/plan_to_sql.rs | 123 +- datafusion/sql/tests/sql_integration.rs | 527 +- datafusion/sqllogictest/Cargo.toml | 7 +- datafusion/sqllogictest/test_files/alias.slt | 59 + datafusion/sqllogictest/test_files/array.slt | 156 +- datafusion/sqllogictest/test_files/copy.slt | 17 + datafusion/sqllogictest/test_files/cte.slt | 146 + datafusion/sqllogictest/test_files/ddl.slt | 28 + .../sqllogictest/test_files/explain.slt | 11 +- .../sqllogictest/test_files/explain_tree.slt | 1643 ++++++ datafusion/sqllogictest/test_files/expr.slt | 13 + .../sqllogictest/test_files/group_by.slt | 29 + .../test_files/information_schema.slt | 8 +- datafusion/sqllogictest/test_files/joins.slt | 266 + datafusion/sqllogictest/test_files/order.slt | 61 +- .../sqllogictest/test_files/prepare.slt | 15 + datafusion/sqllogictest/test_files/scalar.slt | 2 +- datafusion/sqllogictest/test_files/select.slt | 13 + .../sqllogictest/test_files/simplify_expr.slt | 34 + .../test_files/string/string_view.slt | 14 +- .../sqllogictest/test_files/subquery.slt | 58 + .../sqllogictest/test_files/timestamps.slt | 20 + .../sqllogictest/test_files/wildcard.slt | 15 + datafusion/sqllogictest/test_files/window.slt | 26 +- datafusion/substrait/Cargo.toml | 2 +- .../substrait/src/logical_plan/consumer.rs | 37 +- .../substrait/src/logical_plan/producer.rs | 35 +- .../tests/cases/roundtrip_logical_plan.rs | 66 + datafusion/substrait/tests/utils.rs | 1 + datafusion/wasmtest/Cargo.toml | 2 +- dev/changelog/46.0.0.md | 421 ++ dev/release/README.md | 2 + dev/release/verify-release-candidate.sh | 39 +- docs/build.sh | 2 +- docs/source/conf.py | 4 - .../api-health.md | 8 + docs/source/contributor-guide/howtos.md | 4 +- docs/source/contributor-guide/index.md | 2 +- docs/source/contributor-guide/testing.md | 12 + docs/source/index.rst | 27 +- .../library-user-guide/query-optimizer.md | 4 +- docs/source/library-user-guide/upgrading.md | 215 + docs/source/user-guide/configs.md | 7 +- docs/source/user-guide/sql/operators.md | 63 +- .../source/user-guide/sql/scalar_functions.md | 33 + 336 files changed, 22361 insertions(+), 15131 deletions(-) create mode 100644 datafusion-cli/CONTRIBUTING.md create mode 100644 datafusion-cli/tests/snapshots/aws_options.snap create mode 100644 datafusion-cli/tests/snapshots/cli@load_local_csv.sql.snap create mode 100644 datafusion-cli/tests/snapshots/cli@load_s3_csv.sql.snap create mode 100644 datafusion-cli/tests/snapshots/cli@select.sql.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@automatic.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@csv.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@json.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@nd-json.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@table.snap create mode 100644 datafusion-cli/tests/snapshots/cli_format@tsv.snap create mode 100644 datafusion-cli/tests/snapshots/cli_quick_test@backslash.snap create mode 100644 datafusion-cli/tests/snapshots/cli_quick_test@batch_size.snap create mode 100644 datafusion-cli/tests/snapshots/cli_quick_test@files.snap create mode 100644 datafusion-cli/tests/snapshots/cli_quick_test@statements.snap rename datafusion-cli/tests/{data/backslash.txt => sql/backslash.sql} (100%) create mode 100644 datafusion-cli/tests/sql/integration/load_local_csv.sql create mode 100644 datafusion-cli/tests/sql/integration/load_s3_csv.sql rename datafusion-cli/tests/{data/sql.txt => sql/select.sql} (100%) rename datafusion/core/src/datasource/physical_plan/{parquet/mod.rs => parquet.rs} (72%) create mode 100644 datafusion/core/tests/serde/mod.rs create mode 100644 datafusion/datasource-avro/Cargo.toml create mode 120000 datafusion/datasource-avro/LICENSE.txt create mode 120000 datafusion/datasource-avro/NOTICE.txt create mode 100644 datafusion/datasource-avro/README.md rename datafusion/{core/src/datasource => datasource-avro/src}/avro_to_arrow/arrow_array_reader.rs (99%) rename datafusion/{core/src/datasource => datasource-avro/src}/avro_to_arrow/mod.rs (67%) rename datafusion/{core/src/datasource => datasource-avro/src}/avro_to_arrow/reader.rs (95%) rename datafusion/{core/src/datasource => datasource-avro/src}/avro_to_arrow/schema.rs (98%) create mode 100644 datafusion/datasource-avro/src/file_format.rs create mode 100644 datafusion/datasource-avro/src/mod.rs create mode 100644 datafusion/datasource-avro/src/source.rs create mode 100644 datafusion/datasource-csv/Cargo.toml create mode 120000 datafusion/datasource-csv/LICENSE.txt create mode 120000 datafusion/datasource-csv/NOTICE.txt create mode 100644 datafusion/datasource-csv/README.md create mode 100644 datafusion/datasource-csv/src/file_format.rs create mode 100644 datafusion/datasource-csv/src/mod.rs create mode 100644 datafusion/datasource-csv/src/source.rs create mode 100644 datafusion/datasource-json/Cargo.toml create mode 120000 datafusion/datasource-json/LICENSE.txt create mode 120000 datafusion/datasource-json/NOTICE.txt create mode 100644 datafusion/datasource-json/README.md create mode 100644 datafusion/datasource-json/src/file_format.rs create mode 100644 datafusion/datasource-json/src/mod.rs create mode 100644 datafusion/datasource-json/src/source.rs create mode 100644 datafusion/datasource-parquet/Cargo.toml create mode 120000 datafusion/datasource-parquet/LICENSE.txt create mode 120000 datafusion/datasource-parquet/NOTICE.txt create mode 100644 datafusion/datasource-parquet/README.md rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/access_plan.rs (99%) create mode 100644 datafusion/datasource-parquet/src/file_format.rs rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/metrics.rs (99%) create mode 100644 datafusion/datasource-parquet/src/mod.rs rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/opener.rs (95%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/page_filter.rs (99%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/reader.rs (98%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/row_filter.rs (98%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/row_group_filter.rs (99%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/source.rs (94%) rename datafusion/{core/src/datasource/physical_plan/parquet => datasource-parquet/src}/writer.rs (98%) create mode 100644 datafusion/datasource/src/decoder.rs rename datafusion/{core/src/datasource => datasource/src}/schema_adapter.rs (69%) create mode 100644 datafusion/functions-nested/src/max.rs delete mode 100644 datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs rename datafusion/optimizer/src/{unwrap_cast_in_comparison.rs => simplify_expressions/unwrap_cast.rs} (79%) delete mode 100755 datafusion/physical-expr/src/equivalence/properties.rs create mode 100644 datafusion/physical-expr/src/equivalence/properties/dependency.rs create mode 100644 datafusion/physical-expr/src/equivalence/properties/joins.rs create mode 100644 datafusion/physical-expr/src/equivalence/properties/mod.rs create mode 100644 datafusion/physical-expr/src/equivalence/properties/union.rs create mode 100644 datafusion/physical-plan/src/render_tree.rs create mode 100644 datafusion/sqllogictest/test_files/alias.slt create mode 100644 datafusion/sqllogictest/test_files/explain_tree.slt create mode 100644 datafusion/sqllogictest/test_files/simplify_expr.slt create mode 100644 dev/changelog/46.0.0.md rename docs/source/{library-user-guide => contributor-guide}/api-health.md (92%) create mode 100644 docs/source/library-user-guide/upgrading.md diff --git a/.github/workflows/dev_pr/labeler.yml b/.github/workflows/dev_pr/labeler.yml index 4e44e47f5968..5e8cd19ab5ad 100644 --- a/.github/workflows/dev_pr/labeler.yml +++ b/.github/workflows/dev_pr/labeler.yml @@ -29,11 +29,20 @@ sql: logical-expr: - changed-files: - - any-glob-to-any-file: ['datafusion/expr/**/*'] + - any-glob-to-any-file: ['datafusion/expr/**/*', 'datafusion/expr-common/**/*'] + +ffi: + - changed-files: + - any-glob-to-any-file: ['datafusion/ffi/**/*'] physical-expr: - changed-files: - - any-glob-to-any-file: ['datafusion/physical-expr/**/*', 'datafusion/physical-expr-common/**/*', 'datafusion/physical-expr-aggregate/**/*', 'datafusion/physical-plan/**/*'] + - any-glob-to-any-file: ['datafusion/physical-expr/**/*', 'datafusion/physical-expr-common/**/*', 'datafusion/physical-expr-aggregate/**/*'] + +physical-plan: + - changed-files: + - any-glob-to-any-file: [datafusion/physical-plan/**/*'] + catalog: - changed-files: @@ -47,6 +56,10 @@ execution: - changed-files: - any-glob-to-any-file: ['datafusion/execution/**/*'] +datasource: + - changed-files: + - any-glob-to-any-file: ['datafusion/datasource/**/*'] + functions: - changed-files: - any-glob-to-any-file: ['datafusion/functions/**/*', 'datafusion/functions-aggregate/**/*', 'datafusion/functions-aggregate-common', 'datafusion/functions-nested'] diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index 0b43339f57a6..5f1b2c139598 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -1,3 +1,20 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + on: push: branches: diff --git a/.github/workflows/docs_pr.yaml b/.github/workflows/docs_pr.yaml index 3fad08643aa2..d3c901c5b71b 100644 --- a/.github/workflows/docs_pr.yaml +++ b/.github/workflows/docs_pr.yaml @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +# Tests for Docs that runs on PRs name: Docs concurrency: @@ -48,7 +49,34 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable - - name: Run doctests + - name: Run doctests (embedded rust examples) run: cargo test --doc --features avro,json - name: Verify Working Directory Clean run: git diff --exit-code + + # Test doc build + linux-test-doc-build: + name: Test doc build + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: true + fetch-depth: 1 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install doc dependencies + run: | + set -x + python3 -m venv venv + source venv/bin/activate + pip install -r docs/requirements.txt + - name: Build docs html and check for warnings + run: | + set -x + source venv/bin/activate + cd docs + ./build.sh # fails on errors + diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index 3f882d7a3a82..9ee72653b23a 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -36,6 +36,7 @@ jobs: linux-build-lib: name: linux build test runs-on: ubuntu-latest + # note: do not use amd/rust container to preserve disk space steps: - uses: actions/checkout@v4 with: @@ -45,7 +46,7 @@ jobs: run: | curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y source $HOME/.cargo/env - rustup default stable + rustup toolchain install - name: Install Protobuf Compiler run: sudo apt-get install -y protobuf-compiler - name: Prepare cargo build @@ -58,6 +59,7 @@ jobs: name: cargo test 'extended_tests' (amd64) needs: linux-build-lib runs-on: ubuntu-latest + # note: do not use amd/rust container to preserve disk space steps: - uses: actions/checkout@v4 with: @@ -69,7 +71,7 @@ jobs: run: | curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y source $HOME/.cargo/env - rustup default stable + rustup toolchain install - name: Install Protobuf Compiler run: sudo apt-get install -y protobuf-compiler # For debugging, test binaries can be large. @@ -122,7 +124,7 @@ jobs: rust-version: stable - name: Run sqllogictest run: | - cargo test --profile release-nonlto --test sqllogictests -- --include-sqlite + cargo test --features backtrace --profile release-nonlto --test sqllogictests -- --include-sqlite cargo clean diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 99aaa7d6f290..dfdc13057d4e 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -66,9 +66,12 @@ jobs: # the changes to `Cargo.lock` after building with the updated manifest. cargo check --profile ci --workspace --all-targets --features integration-tests --locked - # cargo check common, functions and substrait with no default features - linux-cargo-check-no-default-features: - name: cargo check no default features + # Check datafusion-common features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. + linux-datafusion-common-features: + name: cargo check datafusion-common features needs: linux-build-lib runs-on: ubuntu-latest container: @@ -79,28 +82,68 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable - - name: Check datafusion without default features - # Some of the test binaries require the parquet feature still - #run: cargo check --all-targets --no-default-features -p datafusion - run: cargo check --profile ci --no-default-features -p datafusion - - - name: Check datafusion-common without default features + - name: Check datafusion-common (no-default-features) run: cargo check --profile ci --all-targets --no-default-features -p datafusion-common + # Note: don't check other feature flags as datafusion-common is not typically used standalone - - name: Check datafusion-functions without default features - run: cargo check --profile ci --all-targets --no-default-features -p datafusion-functions - - - name: Check datafusion-substrait without default features + # Check datafusion-substrait features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. + linux-datafusion-substrait-features: + name: cargo check datafusion-substrait features + needs: linux-build-lib + runs-on: ubuntu-latest + container: + image: amd64/rust + steps: + - uses: actions/checkout@v4 + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: stable + - name: Check datafusion-substrait (no-default-features) run: cargo check --profile ci --all-targets --no-default-features -p datafusion-substrait + - name: Check datafusion-substrait (physical) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-substrait --features=physical + - name: Install cmake + run: | + # note the builder setup runs apt-get update / installs protobuf compiler + apt-get install -y cmake + - name: Check datafusion-substrait (protoc) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-substrait --features=protoc - - name: Check workspace in debug mode - run: cargo check --profile ci --all-targets --workspace - - - name: Check workspace with additional features - run: cargo check --profile ci --workspace --benches --features avro,json,integration-tests - - # cargo check datafusion to ensure that the datafusion crate can be built with only a - # subset of the function packages enabled. + # Check datafusion-proto features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. + linux-datafusion-proto-features: + name: cargo check datafusion-proto features + needs: linux-build-lib + runs-on: ubuntu-latest + container: + image: amd64/rust + steps: + - uses: actions/checkout@v4 + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: stable + - name: Check datafusion-proto (no-default-features) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-proto + # fails due to https://github.com/apache/datafusion/issues/15157 + #- name: Check datafusion-proto (json) + # run: cargo check --profile ci --all-targets --no-default-features -p datafusion-proto --features=json + - name: Check datafusion-proto (parquet) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-proto --features=parquet + - name: Check datafusion-proto (avro) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-proto --features=avro + + + # Check datafusion crate features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. linux-cargo-check-datafusion: name: cargo check datafusion needs: linux-build-lib @@ -113,6 +156,11 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable + - name: Check datafusion (no-default-features) + # Some of the test binaries require the parquet feature still + #run: cargo check --all-targets --no-default-features -p datafusion + run: cargo check --profile ci --no-default-features -p datafusion + - name: Check datafusion (nested_expressions) run: cargo check --profile ci --no-default-features --features=nested_expressions -p datafusion @@ -134,8 +182,10 @@ jobs: - name: Check datafusion (string_expressions) run: cargo check --profile ci --no-default-features --features=string_expressions -p datafusion - # cargo check datafusion-functions to ensure that the datafusion-functions crate can be built with - # only a subset of the function packages enabled. + # Check datafusion-functions crate features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. linux-cargo-check-datafusion-functions: name: cargo check functions needs: linux-build-lib @@ -148,6 +198,9 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable + - name: Check datafusion-functions (no-default-features) + run: cargo check --profile ci --all-targets --no-default-features -p datafusion-functions + - name: Check datafusion-functions (crypto) run: cargo check --profile ci --all-targets --no-default-features --features=crypto_expressions -p datafusion-functions @@ -171,21 +224,41 @@ jobs: name: cargo test (amd64) needs: linux-build-lib runs-on: ubuntu-latest - container: - image: amd64/rust steps: - uses: actions/checkout@v4 with: submodules: true fetch-depth: 1 - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - with: - rust-version: stable + run: rustup toolchain install stable + - name: Install Protobuf Compiler + run: sudo apt-get install -y protobuf-compiler + - name: Setup Minio - S3-compatible storage + run: | + docker run -d --name minio-container \ + -p 9000:9000 \ + -e MINIO_ROOT_USER=TEST-DataFusionLogin -e MINIO_ROOT_PASSWORD=TEST-DataFusionPassword \ + -v $(pwd)/datafusion/core/tests/data:/source quay.io/minio/minio \ + server /data + docker exec minio-container /bin/sh -c "\ + mc ready local + mc alias set localminio http://localhost:9000 TEST-DataFusionLogin TEST-DataFusionPassword && \ + mc mb localminio/data && \ + mc cp -r /source/* localminio/data" - name: Run tests (excluding doctests) + env: + RUST_BACKTRACE: 1 + AWS_ENDPOINT: http://127.0.0.1:9000 + AWS_ACCESS_KEY_ID: TEST-DataFusionLogin + AWS_SECRET_ACCESS_KEY: TEST-DataFusionPassword + TEST_STORAGE_INTEGRATION: 1 + AWS_ALLOW_HTTP: true run: cargo test --profile ci --exclude datafusion-examples --exclude ffi_example_table_provider --exclude datafusion-benchmarks --workspace --lib --tests --bins --features avro,json,backtrace,integration-tests - name: Verify Working Directory Clean run: git diff --exit-code + - name: Minio Output + if: ${{ !cancelled() }} + run: docker logs minio-container linux-test-example: name: cargo examples (amd64) @@ -259,6 +332,10 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable + - name: Install dependencies + run: | + apt-get update -qq + apt-get install -y -qq clang - name: Install wasm-pack run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh - name: Build with wasm-pack @@ -295,7 +372,7 @@ jobs: export RUST_MIN_STACK=20971520 export TPCH_DATA=`realpath datafusion/sqllogictest/test_files/tpch/data` cargo test plan_q --package datafusion-benchmarks --profile ci --features=ci -- --test-threads=1 - INCLUDE_TPCH=true cargo test --profile ci --package datafusion-sqllogictest --test sqllogictests + INCLUDE_TPCH=true cargo test --features backtrace --profile ci --package datafusion-sqllogictest --test sqllogictests - name: Verify Working Directory Clean run: git diff --exit-code @@ -331,7 +408,7 @@ jobs: - name: Run sqllogictest run: | cd datafusion/sqllogictest - PG_COMPAT=true PG_URI="postgresql://postgres:postgres@$POSTGRES_HOST:$POSTGRES_PORT/db_test" cargo test --profile ci --features=postgres --test sqllogictests + PG_COMPAT=true PG_URI="postgresql://postgres:postgres@$POSTGRES_HOST:$POSTGRES_PORT/db_test" cargo test --features backtrace --profile ci --features=postgres --test sqllogictests env: # use postgres for the host here because we have specified a container for the job POSTGRES_HOST: postgres diff --git a/Cargo.lock b/Cargo.lock index e005fb08877c..cf0e884a2f42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -453,6 +453,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "85934a9d0261e0fa5d4e2a5295107d743b543a6e0484a835d4b8db2da15306f9" dependencies = [ "bitflags 2.8.0", + "serde", ] [[package]] @@ -516,11 +517,11 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.18" +version = "0.4.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df895a515f70646414f4b45c0b79082783b80552b373a68283012928df56f522" +checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" dependencies = [ - "bzip2 0.4.4", + "bzip2 0.5.1", "flate2", "futures-core", "memchr", @@ -548,7 +549,7 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -570,18 +571,18 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] name = "async-trait" -version = "0.1.86" +version = "0.1.87" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "644dd749086bf3771a2fbc5f256fdb982d53f011c7d5d560304eafeecebce79d" +checksum = "d556ec1359574147ec0c4fc5eb525f3f23263a592b1a9c07e0a75b427de55c97" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -607,9 +608,9 @@ checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" [[package]] name = "aws-config" -version = "1.5.17" +version = "1.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "490aa7465ee685b2ced076bb87ef654a47724a7844e2c7d3af4e749ce5b875dd" +checksum = "90aff65e86db5fe300752551c1b015ef72b708ac54bded8ef43d0d53cb7cb0b1" dependencies = [ "aws-credential-types", "aws-runtime", @@ -617,7 +618,7 @@ dependencies = [ "aws-sdk-ssooidc", "aws-sdk-sts", "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.61.1", "aws-smithy-json", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -656,7 +657,7 @@ dependencies = [ "aws-credential-types", "aws-sigv4", "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.60.12", "aws-smithy-runtime", "aws-smithy-runtime-api", "aws-smithy-types", @@ -674,14 +675,14 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.60.0" +version = "1.61.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60186fab60b24376d3e33b9ff0a43485f99efd470e3b75a9160c849741d63d56" +checksum = "e65ff295979977039a25f5a0bf067a64bc5e6aa38f3cef4037cf42516265553c" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.61.1", "aws-smithy-json", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -696,14 +697,14 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.61.0" +version = "1.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7033130ce1ee13e6018905b7b976c915963755aef299c1521897679d6cd4f8ef" +checksum = "91430a60f754f235688387b75ee798ef00cfd09709a582be2b7525ebb5306d4f" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.61.1", "aws-smithy-json", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -718,14 +719,14 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.61.0" +version = "1.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5c1cac7677179d622b4448b0d31bcb359185295dc6fca891920cfb17e2b5156" +checksum = "9276e139d39fff5a0b0c984fc2d30f970f9a202da67234f948fda02e5bea1dbe" dependencies = [ "aws-credential-types", "aws-runtime", "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.61.1", "aws-smithy-json", "aws-smithy-query", "aws-smithy-runtime", @@ -746,7 +747,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9bfe75fad52793ce6dec0dc3d4b1f388f038b5eb866c8d4d7f3a8e21b5ea5051" dependencies = [ "aws-credential-types", - "aws-smithy-http", + "aws-smithy-http 0.60.12", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", @@ -793,6 +794,26 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-smithy-http" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6f276f21c7921fe902826618d1423ae5bf74cf8c1b8472aee8434f3dfd31824" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "bytes-utils", + "futures-core", + "http 0.2.12", + "http-body 0.4.6", + "once_cell", + "percent-encoding", + "pin-project-lite", + "pin-utils", + "tracing", +] + [[package]] name = "aws-smithy-json" version = "0.61.2" @@ -819,7 +840,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d526a12d9ed61fadefda24abe2e682892ba288c2018bcb38b1b4c111d13f6d92" dependencies = [ "aws-smithy-async", - "aws-smithy-http", + "aws-smithy-http 0.60.12", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", @@ -1087,7 +1108,7 @@ dependencies = [ "serde_json", "serde_repr", "serde_urlencoded", - "thiserror 2.0.11", + "thiserror 2.0.12", "tokio", "tokio-util", "tower-service", @@ -1126,7 +1147,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -1197,9 +1218,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.10.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f61dac84819c6588b558454b194026eb1f09c293b9036ae9b159e74e73ab6cf9" +checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" [[package]] name = "bytes-utils" @@ -1346,9 +1367,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.5.30" +version = "4.5.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92b7b18d71fad5313a1e320fa9897994228ce274b60faa4d694fe0ea89cd9e6d" +checksum = "027bb0d98429ae334a8698531da7077bdf906419543a35a55c2cb1b66437d767" dependencies = [ "clap_builder", "clap_derive", @@ -1356,9 +1377,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.30" +version = "4.5.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a35db2071778a7344791a4fb4f95308b5673d219dee3ae348b86642574ecc90c" +checksum = "5589e0cba072e0f3d23791efac0fd8627b49c829c196a492e88168e6a669d863" dependencies = [ "anstream", "anstyle", @@ -1375,7 +1396,7 @@ dependencies = [ "heck 0.5.0", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -1550,7 +1571,7 @@ dependencies = [ "anes", "cast", "ciborium", - "clap 4.5.30", + "clap 4.5.31", "criterion-plot", "futures", "is-terminal", @@ -1657,7 +1678,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a2785755761f3ddc1492979ce1e48d2c00d09311c39e4466429188f3dd6501" dependencies = [ "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -1681,7 +1702,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -1692,7 +1713,7 @@ checksum = "d336a2a514f6ccccaa3e09b02d41d35330c07ddf03a62165fcec10bb561c7806" dependencies = [ "darling_core", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -1717,9 +1738,8 @@ dependencies = [ [[package]] name = "datafusion" -version = "45.0.0" +version = "46.0.0" dependencies = [ - "apache-avro", "arrow", "arrow-ipc", "arrow-schema", @@ -1735,6 +1755,10 @@ dependencies = [ "datafusion-common", "datafusion-common-runtime", "datafusion-datasource", + "datafusion-datasource-avro", + "datafusion-datasource-csv", + "datafusion-datasource-json", + "datafusion-datasource-parquet", "datafusion-doc", "datafusion-execution", "datafusion-expr", @@ -1759,7 +1783,6 @@ dependencies = [ "itertools 0.14.0", "log", "nix", - "num-traits", "object_store", "parking_lot", "parquet", @@ -1783,7 +1806,7 @@ dependencies = [ [[package]] name = "datafusion-benchmarks" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "datafusion", @@ -1807,7 +1830,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-trait", @@ -1826,7 +1849,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-trait", @@ -1847,19 +1870,21 @@ dependencies = [ [[package]] name = "datafusion-cli" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "assert_cmd", "async-trait", "aws-config", "aws-credential-types", - "clap 4.5.30", + "clap 4.5.31", "ctor", "datafusion", "dirs", "env_logger", "futures", + "insta", + "insta-cmd", "mimalloc", "object_store", "parking_lot", @@ -1874,7 +1899,7 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "apache-avro", @@ -1900,7 +1925,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "45.0.0" +version = "46.0.0" dependencies = [ "log", "tokio", @@ -1908,7 +1933,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-compression", @@ -1930,6 +1955,7 @@ dependencies = [ "itertools 0.14.0", "log", "object_store", + "parquet", "rand 0.8.5", "tempfile", "tokio", @@ -1939,13 +1965,114 @@ dependencies = [ "zstd", ] +[[package]] +name = "datafusion-datasource-avro" +version = "46.0.0" +dependencies = [ + "apache-avro", + "arrow", + "async-trait", + "bytes", + "chrono", + "datafusion-catalog", + "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "futures", + "num-traits", + "object_store", + "rstest", + "serde_json", + "tokio", +] + +[[package]] +name = "datafusion-datasource-csv" +version = "46.0.0" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "rand 0.8.5", + "regex", + "tokio", + "url", +] + +[[package]] +name = "datafusion-datasource-json" +version = "46.0.0" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "futures", + "object_store", + "serde_json", + "tokio", +] + +[[package]] +name = "datafusion-datasource-parquet" +version = "46.0.0" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "chrono", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-optimizer", + "datafusion-physical-plan", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "parking_lot", + "parquet", + "rand 0.8.5", + "tokio", +] + [[package]] name = "datafusion-doc" -version = "45.0.0" +version = "46.0.0" [[package]] name = "datafusion-examples" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "arrow-flight", @@ -1972,7 +2099,7 @@ dependencies = [ [[package]] name = "datafusion-execution" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "chrono", @@ -1990,7 +2117,7 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "chrono", @@ -2011,7 +2138,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2022,7 +2149,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" -version = "45.0.0" +version = "46.0.0" dependencies = [ "abi_stable", "arrow", @@ -2040,7 +2167,7 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "arrow-buffer", @@ -2069,7 +2196,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "arrow", @@ -2090,7 +2217,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "arrow", @@ -2103,7 +2230,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "arrow-ord", @@ -2124,7 +2251,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-trait", @@ -2138,7 +2265,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2154,7 +2281,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2162,16 +2289,16 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "45.0.0" +version = "46.0.0" dependencies = [ "datafusion-expr", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] name = "datafusion-optimizer" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-trait", @@ -2195,7 +2322,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "arrow", @@ -2219,7 +2346,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "arrow", @@ -2231,7 +2358,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2249,7 +2376,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "45.0.0" +version = "46.0.0" dependencies = [ "ahash 0.8.11", "arrow", @@ -2283,7 +2410,7 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "chrono", @@ -2306,7 +2433,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2319,7 +2446,7 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "bigdecimal", @@ -2342,14 +2469,14 @@ dependencies = [ [[package]] name = "datafusion-sqllogictest" -version = "45.0.0" +version = "46.0.0" dependencies = [ "arrow", "async-trait", "bigdecimal", "bytes", "chrono", - "clap 4.5.30", + "clap 4.5.31", "datafusion", "env_logger", "futures", @@ -2366,14 +2493,14 @@ dependencies = [ "tempfile", "testcontainers", "testcontainers-modules", - "thiserror 2.0.11", + "thiserror 2.0.12", "tokio", "tokio-postgres", ] [[package]] name = "datafusion-substrait" -version = "45.0.0" +version = "46.0.0" dependencies = [ "async-recursion", "async-trait", @@ -2392,7 +2519,7 @@ dependencies = [ [[package]] name = "datafusion-wasmtest" -version = "45.0.0" +version = "46.0.0" dependencies = [ "chrono", "console_error_panic_hook", @@ -2465,7 +2592,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -2500,7 +2627,7 @@ dependencies = [ "enum-ordinalize", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -2538,7 +2665,7 @@ checksum = "0d28318a75d4aead5c4db25382e8ef717932d0346600cacae6357eb5941bc5ff" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -2553,14 +2680,14 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.6" +version = "0.11.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dcaee3d8e3cfc3fd92428d477bc97fc29ec8716d180c0d74c643bb26166660e0" +checksum = "c3716d7a920fb4fac5d84e9d4bce8ceb321e9414b4409da61b07b75c1e3d0697" dependencies = [ "anstream", "anstyle", "env_filter", - "humantime", + "jiff", "log", ] @@ -2622,7 +2749,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e5768da2206272c81ef0b5e951a41862938a6070da63bcea197899942d3b947" dependencies = [ "cfg-if", - "rustix", + "rustix 0.38.44", "windows-sys 0.52.0", ] @@ -2795,7 +2922,7 @@ checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -2906,6 +3033,19 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" +[[package]] +name = "globset" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54a1028dfc5f5df5da8a56a73e6c153c9a9708ec57232470703592a3f18e49f5" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata", + "regex-syntax", +] + [[package]] name = "h2" version = "0.3.26" @@ -3383,7 +3523,7 @@ checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -3454,6 +3594,34 @@ version = "2.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5" +[[package]] +name = "insta" +version = "1.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50259abbaa67d11d2bcafc7ba1d094ed7a0c70e3ce893f0d0997f73558cb3084" +dependencies = [ + "console", + "globset", + "linked-hash-map", + "once_cell", + "pin-project", + "regex", + "serde", + "similar", + "walkdir", +] + +[[package]] +name = "insta-cmd" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffeeefa927925cced49ccb01bf3e57c9d4cd132df21e576eb9415baeab2d3de6" +dependencies = [ + "insta", + "serde", + "serde_json", +] + [[package]] name = "integer-encoding" version = "3.0.4" @@ -3516,6 +3684,30 @@ version = "1.0.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +[[package]] +name = "jiff" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d699bc6dfc879fb1bf9bdff0d4c56f0884fc6f0d0eb0fba397a6d00cd9a6b85e" +dependencies = [ + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde", +] + +[[package]] +name = "jiff-static" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d16e75759ee0aa64c57a56acbf43916987b20c77373cb7e808979e02b93c9f9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.100", +] + [[package]] name = "jobserver" version = "0.1.32" @@ -3607,9 +3799,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.170" +version = "0.2.171" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "875b3680cb2f8f71bdcf9a30f38d48282f5d3c95cbf9b3fa57269bb5d5c06828" +checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" [[package]] name = "libflate" @@ -3680,16 +3872,28 @@ checksum = "5297962ef19edda4ce33aaa484386e0a5b3d7f2f4e037cbeee00503ef6b29d33" dependencies = [ "anstream", "anstyle", - "clap 4.5.30", + "clap 4.5.31", "escape8259", ] +[[package]] +name = "linked-hash-map" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" + [[package]] name = "linux-raw-sys" version = "0.4.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" +[[package]] +name = "linux-raw-sys" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9" + [[package]] name = "litemap" version = "0.7.4" @@ -4113,7 +4317,7 @@ dependencies = [ "regex", "regex-syntax", "structmeta", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4239,7 +4443,7 @@ checksum = "f6e859e6e5bd50440ab63c47e3ebabc90f26251f7c73c3d3e837b74a1cc3fa67" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4294,6 +4498,15 @@ version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6" +[[package]] +name = "portable-atomic-util" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507" +dependencies = [ + "portable-atomic", +] + [[package]] name = "postgres-derive" version = "0.4.6" @@ -4303,7 +4516,7 @@ dependencies = [ "heck 0.5.0", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4389,7 +4602,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6924ced06e1f7dfe3fa48d57b9f74f55d8915f5036121bef647ef4b204895fac" dependencies = [ "proc-macro2", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4460,7 +4673,7 @@ dependencies = [ "prost", "prost-types", "regex", - "syn 2.0.98", + "syn 2.0.100", "tempfile", ] @@ -4474,7 +4687,7 @@ dependencies = [ "itertools 0.14.0", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4526,9 +4739,9 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.23.4" +version = "0.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57fe09249128b3173d092de9523eaa75136bf7ba85e0d69eca241c7939c933cc" +checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" dependencies = [ "cfg-if", "indoc", @@ -4544,9 +4757,9 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.23.4" +version = "0.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1cd3927b5a78757a0d71aa9dff669f903b1eb64b54142a9bd9f757f8fde65fd7" +checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" dependencies = [ "once_cell", "target-lexicon", @@ -4554,9 +4767,9 @@ dependencies = [ [[package]] name = "pyo3-ffi" -version = "0.23.4" +version = "0.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dab6bb2102bd8f991e7749f130a70d05dd557613e39ed2deeee8e9ca0c4d548d" +checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" dependencies = [ "libc", "pyo3-build-config", @@ -4564,27 +4777,27 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.23.4" +version = "0.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91871864b353fd5ffcb3f91f2f703a22a9797c91b9ab497b1acac7b07ae509c7" +checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] name = "pyo3-macros-backend" -version = "0.23.4" +version = "0.23.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43abc3b80bc20f3facd86cd3c60beed58c3e2aa26213f3cda368de39c60a27e4" +checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" dependencies = [ "heck 0.5.0", "proc-macro2", "pyo3-build-config", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4616,7 +4829,7 @@ dependencies = [ "rustc-hash", "rustls 0.23.23", "socket2", - "thiserror 2.0.11", + "thiserror 2.0.12", "tokio", "tracing", ] @@ -4635,7 +4848,7 @@ dependencies = [ "rustls 0.23.23", "rustls-pki-types", "slab", - "thiserror 2.0.11", + "thiserror 2.0.12", "tinyvec", "tracing", "web-time", @@ -4788,7 +5001,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -4817,7 +5030,7 @@ checksum = "dd6f9d3d47bdd2ad6945c5015a226ec6155d0bcdfd8f7cd29f86b71f8de99d2b" dependencies = [ "getrandom 0.2.15", "libredox", - "thiserror 2.0.11", + "thiserror 2.0.12", ] [[package]] @@ -4937,9 +5150,9 @@ dependencies = [ [[package]] name = "ring" -version = "0.17.9" +version = "0.17.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e75ec5e92c4d8aede845126adc388046234541629e76029599ed35a003c7ed24" +checksum = "70ac5d832aa16abd7d1def883a8545280c20a60f523a370aa3a9617c2b8550ee" dependencies = [ "cc", "cfg-if", @@ -5010,7 +5223,7 @@ dependencies = [ "regex", "relative-path", "rustc_version", - "syn 2.0.98", + "syn 2.0.100", "unicode-ident", ] @@ -5022,7 +5235,7 @@ checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14" dependencies = [ "quote", "rand 0.8.5", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5072,7 +5285,20 @@ dependencies = [ "bitflags 2.8.0", "errno", "libc", - "linux-raw-sys", + "linux-raw-sys 0.4.15", + "windows-sys 0.59.0", +] + +[[package]] +name = "rustix" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7178faa4b75a30e269c71e61c353ce2748cf3d76f0c44c393f4e60abf49b825" +dependencies = [ + "bitflags 2.8.0", + "errno", + "libc", + "linux-raw-sys 0.9.2", "windows-sys 0.59.0", ] @@ -5247,7 +5473,7 @@ dependencies = [ "proc-macro2", "quote", "serde_derive_internals", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5310,9 +5536,9 @@ dependencies = [ [[package]] name = "semver" -version = "1.0.25" +version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f79dfe2d285b0488816f30e700a7438c5a73d816b5b7d3ac72fbc48b0d185e03" +checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" dependencies = [ "serde", ] @@ -5349,7 +5575,7 @@ checksum = "f09503e191f4e797cb8aac08e9a4a4695c5edf6a2e70e376d961ddd5c969f82b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5360,14 +5586,14 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] name = "serde_json" -version = "1.0.139" +version = "1.0.140" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44f86c3acccc9c65b153fe1b85a3be07fe5515274ec9f0653b4a0875731c72a6" +checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373" dependencies = [ "itoa", "memchr", @@ -5383,7 +5609,7 @@ checksum = "6c64451ba24fc7a6a2d60fc75dd9c83c90903b19028d4eff35e88fc1e86564e9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5395,7 +5621,7 @@ dependencies = [ "proc-macro2", "quote", "serde", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5437,7 +5663,7 @@ dependencies = [ "darling", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5530,7 +5756,7 @@ dependencies = [ "heck 0.5.0", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5569,9 +5795,9 @@ dependencies = [ [[package]] name = "sqllogictest" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f1c93848602f92e5925690d4805ccbc1ccdb61bee7d4ae79ad6862b542a539c" +checksum = "17b2f0b80fc250ed3fdd82fc88c0ada5ad62ee1ed5314ac5474acfa52082f518" dependencies = [ "async-trait", "educe", @@ -5588,15 +5814,15 @@ dependencies = [ "similar", "subst", "tempfile", - "thiserror 2.0.11", + "thiserror 2.0.12", "tracing", ] [[package]] name = "sqlparser" -version = "0.54.0" +version = "0.55.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c66e3b7374ad4a6af849b08b3e7a6eda0edbd82f0fd59b57e22671bf16979899" +checksum = "c4521174166bac1ff04fe16ef4524c70144cd29682a45978978ca3d7f4e0be11" dependencies = [ "log", "recursive", @@ -5611,7 +5837,7 @@ checksum = "da5fc6819faabb412da764b99d3b713bb55083c11e7e0c00144d386cd6a1939c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5665,7 +5891,7 @@ dependencies = [ "proc-macro2", "quote", "structmeta-derive", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5676,7 +5902,7 @@ checksum = "152a0b65a590ff6c3da95cabe2353ee04e6167c896b28e3b14478c2636c922fc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5728,7 +5954,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5741,7 +5967,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5756,9 +5982,9 @@ dependencies = [ [[package]] name = "substrait" -version = "0.53.2" +version = "0.54.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6fac3d70185423235f37b889764e184b81a5af4bb7c95833396ee9bd92577e1b" +checksum = "93890ad613de815a5b76e38bc4a934b4012ebe197717c9dd6a17f7af8cf33dae" dependencies = [ "heck 0.5.0", "pbjson", @@ -5775,7 +6001,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "syn 2.0.98", + "syn 2.0.100", "typify", "walkdir", ] @@ -5799,9 +6025,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.98" +version = "2.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1" +checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" dependencies = [ "proc-macro2", "quote", @@ -5825,7 +6051,7 @@ checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -5856,15 +6082,15 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" [[package]] name = "tempfile" -version = "3.17.1" +version = "3.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e5a0acb1f3f55f65cc4a866c361b2fb2a0ff6366785ae6fbb5f85df07ba230" +checksum = "2c317e0a526ee6120d8dabad239c8dadca62b24b6f168914bbbc8e2fb1f0e567" dependencies = [ "cfg-if", "fastrand", "getrandom 0.3.1", "once_cell", - "rustix", + "rustix 1.0.2", "windows-sys 0.59.0", ] @@ -5906,7 +6132,7 @@ dependencies = [ "serde", "serde_json", "serde_with", - "thiserror 2.0.11", + "thiserror 2.0.12", "tokio", "tokio-stream", "tokio-tar", @@ -5943,11 +6169,11 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.11" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d452f284b73e6d76dd36758a0c8684b1d5be31f92b89d07fd5822175732206fc" +checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ - "thiserror-impl 2.0.11", + "thiserror-impl 2.0.12", ] [[package]] @@ -5958,18 +6184,18 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] name = "thiserror-impl" -version = "2.0.11" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26afc1baea8a989337eeb52b6e72a039780ce45c3edfcc9c5b9d112feeb173c2" +checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6084,7 +6310,7 @@ checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6285,7 +6511,7 @@ checksum = "395ae124c09f9e6918a2310af6038fba074bcf474ac352496d5910dd59a2226d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6351,7 +6577,7 @@ checksum = "f9534daa9fd3ed0bd911d462a37f172228077e7abf18c18a5f67199d959205f8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6385,8 +6611,8 @@ dependencies = [ "semver", "serde", "serde_json", - "syn 2.0.98", - "thiserror 2.0.11", + "syn 2.0.100", + "thiserror 2.0.12", "unicode-ident", ] @@ -6403,7 +6629,7 @@ dependencies = [ "serde", "serde_json", "serde_tokenstream", - "syn 2.0.98", + "syn 2.0.100", "typify-impl", ] @@ -6601,7 +6827,7 @@ dependencies = [ "log", "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", "wasm-bindgen-shared", ] @@ -6636,7 +6862,7 @@ checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -6671,7 +6897,7 @@ checksum = "17d5042cc5fa009658f9a7333ef24291b1291a25b6382dd68862a7f3b969f69b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6788,7 +7014,7 @@ checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -6799,7 +7025,7 @@ checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -7035,8 +7261,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e105d177a3871454f754b33bb0ee637ecaaac997446375fd3e5d43a2ed00c909" dependencies = [ "libc", - "linux-raw-sys", - "rustix", + "linux-raw-sys 0.4.15", + "rustix 0.38.44", ] [[package]] @@ -7074,7 +7300,7 @@ checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", "synstructure", ] @@ -7105,7 +7331,7 @@ checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -7116,7 +7342,7 @@ checksum = "76331675d372f91bf8d17e13afbd5fe639200b73d01f0fc748bb059f9cca2db7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] @@ -7136,7 +7362,7 @@ checksum = "595eed982f7d355beb85837f651fa22e90b3c044842dc7f2c2842c086f295808" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", "synstructure", ] @@ -7165,7 +7391,7 @@ checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.98", + "syn 2.0.100", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 40064d045382..df03f03557f0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,12 +16,16 @@ # under the License. [workspace] -exclude = ["dev/depcheck"] members = [ "datafusion/common", "datafusion/common-runtime", "datafusion/catalog", "datafusion/catalog-listing", + "datafusion/datasource", + "datafusion/datasource-avro", + "datafusion/datasource-csv", + "datafusion/datasource-json", + "datafusion/datasource-parquet", "datafusion/core", "datafusion/expr", "datafusion/expr-common", @@ -57,6 +61,7 @@ members = [ "datafusion/macros", "datafusion/doc", ] +exclude = ["dev/depcheck"] resolver = "2" [workspace.package] @@ -69,7 +74,7 @@ repository = "https://github.com/apache/datafusion" # Define Minimum Supported Rust Version (MSRV) rust-version = "1.82.0" # Define DataFusion version -version = "45.0.0" +version = "46.0.0" [workspace.dependencies] # We turn off default-features for some dependencies here so the workspaces which inherit them can @@ -80,6 +85,7 @@ version = "45.0.0" ahash = { version = "0.8", default-features = false, features = [ "runtime-rng", ] } +apache-avro = { version = "0.17", default-features = false } arrow = { version = "54.2.1", features = [ "prettyprint", "chrono-tz", @@ -93,40 +99,44 @@ arrow-ipc = { version = "54.2.0", default-features = false, features = [ ] } arrow-ord = { version = "54.1.0", default-features = false } arrow-schema = { version = "54.1.0", default-features = false } -async-trait = "0.1.73" +async-trait = "0.1.87" bigdecimal = "0.4.7" bytes = "1.10" chrono = { version = "0.4.38", default-features = false } criterion = "0.5.1" ctor = "0.2.9" dashmap = "6.0.1" -datafusion = { path = "datafusion/core", version = "45.0.0", default-features = false } -datafusion-catalog = { path = "datafusion/catalog", version = "45.0.0" } -datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "45.0.0" } -datafusion-common = { path = "datafusion/common", version = "45.0.0", default-features = false } -datafusion-common-runtime = { path = "datafusion/common-runtime", version = "45.0.0" } -datafusion-datasource = { path = "datafusion/datasource", version = "45.0.0", default-features = false } -datafusion-doc = { path = "datafusion/doc", version = "45.0.0" } -datafusion-execution = { path = "datafusion/execution", version = "45.0.0" } -datafusion-expr = { path = "datafusion/expr", version = "45.0.0" } -datafusion-expr-common = { path = "datafusion/expr-common", version = "45.0.0" } -datafusion-ffi = { path = "datafusion/ffi", version = "45.0.0" } -datafusion-functions = { path = "datafusion/functions", version = "45.0.0" } -datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "45.0.0" } -datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "45.0.0" } -datafusion-functions-nested = { path = "datafusion/functions-nested", version = "45.0.0" } -datafusion-functions-table = { path = "datafusion/functions-table", version = "45.0.0" } -datafusion-functions-window = { path = "datafusion/functions-window", version = "45.0.0" } -datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "45.0.0" } -datafusion-macros = { path = "datafusion/macros", version = "45.0.0" } -datafusion-optimizer = { path = "datafusion/optimizer", version = "45.0.0", default-features = false } -datafusion-physical-expr = { path = "datafusion/physical-expr", version = "45.0.0", default-features = false } -datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "45.0.0", default-features = false } -datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "45.0.0" } -datafusion-physical-plan = { path = "datafusion/physical-plan", version = "45.0.0" } -datafusion-proto = { path = "datafusion/proto", version = "45.0.0" } -datafusion-proto-common = { path = "datafusion/proto-common", version = "45.0.0" } -datafusion-sql = { path = "datafusion/sql", version = "45.0.0" } +datafusion = { path = "datafusion/core", version = "46.0.0", default-features = false } +datafusion-catalog = { path = "datafusion/catalog", version = "46.0.0" } +datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "46.0.0" } +datafusion-common = { path = "datafusion/common", version = "46.0.0", default-features = false } +datafusion-common-runtime = { path = "datafusion/common-runtime", version = "46.0.0" } +datafusion-datasource = { path = "datafusion/datasource", version = "46.0.0", default-features = false } +datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "46.0.0", default-features = false } +datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "46.0.0", default-features = false } +datafusion-datasource-json = { path = "datafusion/datasource-json", version = "46.0.0", default-features = false } +datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "46.0.0", default-features = false } +datafusion-doc = { path = "datafusion/doc", version = "46.0.0" } +datafusion-execution = { path = "datafusion/execution", version = "46.0.0" } +datafusion-expr = { path = "datafusion/expr", version = "46.0.0" } +datafusion-expr-common = { path = "datafusion/expr-common", version = "46.0.0" } +datafusion-ffi = { path = "datafusion/ffi", version = "46.0.0" } +datafusion-functions = { path = "datafusion/functions", version = "46.0.0" } +datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "46.0.0" } +datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "46.0.0" } +datafusion-functions-nested = { path = "datafusion/functions-nested", version = "46.0.0" } +datafusion-functions-table = { path = "datafusion/functions-table", version = "46.0.0" } +datafusion-functions-window = { path = "datafusion/functions-window", version = "46.0.0" } +datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "46.0.0" } +datafusion-macros = { path = "datafusion/macros", version = "46.0.0" } +datafusion-optimizer = { path = "datafusion/optimizer", version = "46.0.0", default-features = false } +datafusion-physical-expr = { path = "datafusion/physical-expr", version = "46.0.0", default-features = false } +datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "46.0.0", default-features = false } +datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "46.0.0" } +datafusion-physical-plan = { path = "datafusion/physical-plan", version = "46.0.0" } +datafusion-proto = { path = "datafusion/proto", version = "46.0.0" } +datafusion-proto-common = { path = "datafusion/proto-common", version = "46.0.0" } +datafusion-sql = { path = "datafusion/sql", version = "46.0.0" } doc-comment = "0.3" env_logger = "0.11" futures = "0.3" diff --git a/benchmarks/README.md b/benchmarks/README.md index 2954f42c25db..f17d6b5a07b6 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -195,13 +195,13 @@ metadata (number of cores, DataFusion version, etc.). $ git checkout main # generate an output script in /tmp/output_main $ mkdir -p /tmp/output_main -$ cargo run --release --bin tpch -- benchmark datafusion --iterations 5 --path ./data --format parquet -o /tmp/output_main +$ cargo run --release --bin tpch -- benchmark datafusion --iterations 5 --path ./data --format parquet -o /tmp/output_main/tpch.json # generate an output script in /tmp/output_branch $ mkdir -p /tmp/output_branch $ git checkout my_branch -$ cargo run --release --bin tpch -- benchmark datafusion --iterations 5 --path ./data --format parquet -o /tmp/output_branch +$ cargo run --release --bin tpch -- benchmark datafusion --iterations 5 --path ./data --format parquet -o /tmp/output_branch/tpch.json # compare the results: -./compare.py /tmp/output_main/tpch-summary--1679330119.json /tmp/output_branch/tpch-summary--1679328405.json +./compare.py /tmp/output_main/tpch.json /tmp/output_branch/tpch.json ``` This will produce output like: @@ -329,7 +329,37 @@ Your benchmark should create and use an instance of `BenchmarkRun` defined in `b # Benchmarks -The output of `dfbench` help includes a description of each benchmark, which is reproduced here for convenience +The output of `dfbench` help includes a description of each benchmark, which is reproduced here for convenience. + +## Cancellation + +Test performance of cancelling queries +Queries in DataFusion should stop executing "quickly" after they are +cancelled (the output stream is dropped). + +The queries are executed on a synthetic dataset generated during +the benchmark execution that is an anonymized version of a +real-world data set. + +The query is an anonymized version of a real-world query, and the +test starts the query then cancels it and reports how long it takes +for the runtime to fully exit. + +Example output: + +``` +Using 7 files found on disk +Starting to load data into in-memory object store +Done loading data into in-memory object store +in main, sleeping +Starting spawned +Creating logical plan... +Creating physical plan... +Executing physical plan... +Getting results... +cancelling thread +done dropping runtime in 83.531417ms +``` ## ClickBench @@ -513,5 +543,49 @@ For example, to run query 1 with the small data generated above: cargo run --release --bin dfbench -- h2o --path ./benchmarks/data/h2o/G1_1e7_1e7_100_0.csv --query 1 ``` +## h2o benchmarks for join + +### Generate data for h2o benchmarks +There are three options for generating data for h2o benchmarks: `small`, `medium`, and `big`. The data is generated in the `data` directory. + +1. Generate small data (4 table files, the largest is 1e7 rows) +```bash +./bench.sh data h2o_small_join +``` + + +2. Generate medium data (4 table files, the largest is 1e8 rows) +```bash +./bench.sh data h2o_medium_join +``` + +3. Generate large data (4 table files, the largest is 1e9 rows) +```bash +./bench.sh data h2o_big_join +``` + +### Run h2o benchmarks +There are three options for running h2o benchmarks: `small`, `medium`, and `big`. +1. Run small data benchmark +```bash +./bench.sh run h2o_small_join +``` + +2. Run medium data benchmark +```bash +./bench.sh run h2o_medium_join +``` + +3. Run large data benchmark +```bash +./bench.sh run h2o_big_join +``` + +4. Run a specific query with a specific join data paths, the data paths are including 4 table files. + +For example, to run query 1 with the small data generated above: +```bash +cargo run --release --bin dfbench -- h2o --join-paths ./benchmarks/data/h2o/J1_1e7_NA_0.csv,./benchmarks/data/h2o/J1_1e7_1e1_0.csv,./benchmarks/data/h2o/J1_1e7_1e4_0.csv,./benchmarks/data/h2o/J1_1e7_1e7_NA.csv --queries-path ./benchmarks/queries/h2o/join.sql --query 1 +``` [1]: http://www.tpc.org/tpch/ [2]: https://www1.nyc.gov/site/tlc/about/tlc-trip-record-data.page diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 1e56c96479c3..43d3d78c7e1f 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -81,9 +81,12 @@ clickbench_1: ClickBench queries against a single parquet file clickbench_partitioned: ClickBench queries against a partitioned (100 files) parquet clickbench_extended: ClickBench \"inspired\" queries against a single parquet (DataFusion specific) external_aggr: External aggregation benchmark -h2o_small: h2oai benchmark with small dataset (1e7 rows), default file format is csv -h2o_medium: h2oai benchmark with medium dataset (1e8 rows), default file format is csv -h2o_big: h2oai benchmark with large dataset (1e9 rows), default file format is csv +h2o_small: h2oai benchmark with small dataset (1e7 rows) for groupby, default file format is csv +h2o_medium: h2oai benchmark with medium dataset (1e8 rows) for groupby, default file format is csv +h2o_big: h2oai benchmark with large dataset (1e9 rows) for groupby, default file format is csv +h2o_small_join: h2oai benchmark with small dataset (1e7 rows) for join, default file format is csv +h2o_medium_join: h2oai benchmark with medium dataset (1e8 rows) for join, default file format is csv +h2o_big_join: h2oai benchmark with large dataset (1e9 rows) for join, default file format is csv imdb: Join Order Benchmark (JOB) using the IMDB dataset converted to parquet ********** @@ -150,6 +153,9 @@ main() { data_h2o "SMALL" data_h2o "MEDIUM" data_h2o "BIG" + data_h2o_join "SMALL" + data_h2o_join "MEDIUM" + data_h2o_join "BIG" data_clickbench_1 data_clickbench_partitioned data_imdb @@ -189,6 +195,15 @@ main() { h2o_big) data_h2o "BIG" "CSV" ;; + h2o_small_join) + data_h2o_join "SMALL" "CSV" + ;; + h2o_medium_join) + data_h2o_join "MEDIUM" "CSV" + ;; + h2o_big_join) + data_h2o_join "BIG" "CSV" + ;; external_aggr) # same data as for tpch data_tpch "1" @@ -242,6 +257,9 @@ main() { run_h2o "SMALL" "PARQUET" "groupby" run_h2o "MEDIUM" "PARQUET" "groupby" run_h2o "BIG" "PARQUET" "groupby" + run_h2o_join "SMALL" "PARQUET" "join" + run_h2o_join "MEDIUM" "PARQUET" "join" + run_h2o_join "BIG" "PARQUET" "join" run_imdb run_external_aggr ;; @@ -287,6 +305,15 @@ main() { h2o_big) run_h2o "BIG" "CSV" "groupby" ;; + h2o_small_join) + run_h2o_join "SMALL" "CSV" "join" + ;; + h2o_medium_join) + run_h2o_join "MEDIUM" "CSV" "join" + ;; + h2o_big_join) + run_h2o_join "BIG" "CSV" "join" + ;; external_aggr) run_external_aggr ;; @@ -687,7 +714,82 @@ data_h2o() { deactivate } -## todo now only support groupby, after https://github.com/mrpowers-io/falsa/issues/21 done, we can add support for join +data_h2o_join() { + # Default values for size and data format + SIZE=${1:-"SMALL"} + DATA_FORMAT=${2:-"CSV"} + + # Function to compare Python versions + version_ge() { + [ "$(printf '%s\n' "$1" "$2" | sort -V | head -n1)" = "$2" ] + } + + export PYO3_USE_ABI3_FORWARD_COMPATIBILITY=1 + + # Find the highest available Python version (3.10 or higher) + REQUIRED_VERSION="3.10" + PYTHON_CMD=$(command -v python3 || true) + + if [ -n "$PYTHON_CMD" ]; then + PYTHON_VERSION=$($PYTHON_CMD -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") + if version_ge "$PYTHON_VERSION" "$REQUIRED_VERSION"; then + echo "Found Python version $PYTHON_VERSION, which is suitable." + else + echo "Python version $PYTHON_VERSION found, but version $REQUIRED_VERSION or higher is required." + PYTHON_CMD="" + fi + fi + + # Search for suitable Python versions if the default is unsuitable + if [ -z "$PYTHON_CMD" ]; then + # Loop through all available Python3 commands on the system + for CMD in $(compgen -c | grep -E '^python3(\.[0-9]+)?$'); do + if command -v "$CMD" &> /dev/null; then + PYTHON_VERSION=$($CMD -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") + if version_ge "$PYTHON_VERSION" "$REQUIRED_VERSION"; then + PYTHON_CMD="$CMD" + echo "Found suitable Python version: $PYTHON_VERSION ($CMD)" + break + fi + fi + done + fi + + # If no suitable Python version found, exit with an error + if [ -z "$PYTHON_CMD" ]; then + echo "Python 3.10 or higher is required. Please install it." + return 1 + fi + + echo "Using Python command: $PYTHON_CMD" + + # Install falsa and other dependencies + echo "Installing falsa..." + + # Set virtual environment directory + VIRTUAL_ENV="${PWD}/venv" + + # Create a virtual environment using the detected Python command + $PYTHON_CMD -m venv "$VIRTUAL_ENV" + + # Activate the virtual environment and install dependencies + source "$VIRTUAL_ENV/bin/activate" + + # Ensure 'falsa' is installed (avoid unnecessary reinstall) + pip install --quiet --upgrade falsa + + # Create directory if it doesn't exist + H2O_DIR="${DATA_DIR}/h2o" + mkdir -p "${H2O_DIR}" + + # Generate h2o test data + echo "Generating h2o test data in ${H2O_DIR} with size=${SIZE} and format=${DATA_FORMAT}" + falsa join --path-prefix="${H2O_DIR}" --size "${SIZE}" --data-format "${DATA_FORMAT}" + + # Deactivate virtual environment after completion + deactivate +} + run_h2o() { # Default values for size and data format SIZE=${1:-"SMALL"} @@ -700,7 +802,7 @@ run_h2o() { RESULTS_FILE="${RESULTS_DIR}/h2o.json" echo "RESULTS_FILE: ${RESULTS_FILE}" - echo "Running h2o benchmark..." + echo "Running h2o groupby benchmark..." # Set the file name based on the size case "$SIZE" in @@ -730,6 +832,56 @@ run_h2o() { -o "${RESULTS_FILE}" } +run_h2o_join() { + # Default values for size and data format + SIZE=${1:-"SMALL"} + DATA_FORMAT=${2:-"CSV"} + DATA_FORMAT=$(echo "$DATA_FORMAT" | tr '[:upper:]' '[:lower:]') + RUN_Type=${3:-"join"} + + # Data directory and results file path + H2O_DIR="${DATA_DIR}/h2o" + RESULTS_FILE="${RESULTS_DIR}/h2o_join.json" + + echo "RESULTS_FILE: ${RESULTS_FILE}" + echo "Running h2o join benchmark..." + + # Set the file name based on the size + case "$SIZE" in + "SMALL") + X_TABLE_FILE_NAME="J1_1e7_NA_0.${DATA_FORMAT}" + SMALL_TABLE_FILE_NAME="J1_1e7_1e1_0.${DATA_FORMAT}" + MEDIUM_TABLE_FILE_NAME="J1_1e7_1e4_0.${DATA_FORMAT}" + LARGE_TABLE_FILE_NAME="J1_1e7_1e7_NA.${DATA_FORMAT}" + ;; + "MEDIUM") + X_TABLE_FILE_NAME="J1_1e8_NA_0.${DATA_FORMAT}" + SMALL_TABLE_FILE_NAME="J1_1e8_1e2_0.${DATA_FORMAT}" + MEDIUM_TABLE_FILE_NAME="J1_1e8_1e5_0.${DATA_FORMAT}" + LARGE_TABLE_FILE_NAME="J1_1e8_1e8_NA.${DATA_FORMAT}" + ;; + "BIG") + X_TABLE_FILE_NAME="J1_1e9_NA_0.${DATA_FORMAT}" + SMALL_TABLE_FILE_NAME="J1_1e9_1e3_0.${DATA_FORMAT}" + MEDIUM_TABLE_FILE_NAME="J1_1e9_1e6_0.${DATA_FORMAT}" + LARGE_TABLE_FILE_NAME="J1_1e9_1e9_NA.${DATA_FORMAT}" + ;; + *) + echo "Invalid size. Valid options are SMALL, MEDIUM, or BIG." + return 1 + ;; + esac + + # Set the query file name based on the RUN_Type + QUERY_FILE="${SCRIPT_DIR}/queries/h2o/${RUN_Type}.sql" + + $CARGO_COMMAND --bin dfbench -- h2o \ + --iterations 3 \ + --join-paths "${H2O_DIR}/${X_TABLE_FILE_NAME},${H2O_DIR}/${SMALL_TABLE_FILE_NAME},${H2O_DIR}/${MEDIUM_TABLE_FILE_NAME},${H2O_DIR}/${LARGE_TABLE_FILE_NAME}" \ + --queries-path "${QUERY_FILE}" \ + -o "${RESULTS_FILE}" +} + # Runs the external aggregation benchmark run_external_aggr() { # Use TPC-H SF1 dataset diff --git a/benchmarks/src/cancellation.rs b/benchmarks/src/cancellation.rs index 3c3ca424a308..f5740bdc96e0 100644 --- a/benchmarks/src/cancellation.rs +++ b/benchmarks/src/cancellation.rs @@ -47,6 +47,9 @@ use tokio_util::sync::CancellationToken; /// Test performance of cancelling queries /// +/// Queries in DataFusion should stop executing "quickly" after they are +/// cancelled (the output stream is dropped). +/// /// The queries are executed on a synthetic dataset generated during /// the benchmark execution that is an anonymized version of a /// real-world data set. @@ -97,7 +100,7 @@ impl RunOpt { println!("Done loading data into in-memory object store"); let mut rundata = BenchmarkRun::new(); - rundata.start_new_case("Arglebargle"); + rundata.start_new_case("Cancellation"); for i in 0..self.common.iterations { let elapsed = run_test(self.wait_time, Arc::clone(&store))?; diff --git a/benchmarks/src/h2o.rs b/benchmarks/src/h2o.rs index eae7f67f1d62..cc463e70d74a 100644 --- a/benchmarks/src/h2o.rs +++ b/benchmarks/src/h2o.rs @@ -53,6 +53,16 @@ pub struct RunOpt { )] path: PathBuf, + /// Path to data files (parquet or csv), using , to separate the paths + /// Default value is the small files for join x table, small table, medium table, big table files in the h2o benchmark + /// This is the small csv file case + #[structopt( + short = "join-paths", + long = "join-paths", + default_value = "benchmarks/data/h2o/J1_1e7_NA_0.csv,benchmarks/data/h2o/J1_1e7_1e1_0.csv,benchmarks/data/h2o/J1_1e7_1e4_0.csv,benchmarks/data/h2o/J1_1e7_1e7_NA.csv" + )] + join_paths: String, + /// If present, write results json here #[structopt(parse(from_os_str), short = "o", long = "output")] output_path: Option, @@ -71,8 +81,16 @@ impl RunOpt { let rt_builder = self.common.runtime_env_builder()?; let ctx = SessionContext::new_with_config_rt(config, rt_builder.build_arc()?); - // Register data - self.register_data(&ctx).await?; + if self.queries_path.to_str().unwrap().contains("join") { + let join_paths: Vec<&str> = self.join_paths.split(',').collect(); + let table_name: Vec<&str> = vec!["x", "small", "medium", "large"]; + for (i, path) in join_paths.iter().enumerate() { + ctx.register_csv(table_name[i], path, Default::default()) + .await?; + } + } else if self.queries_path.to_str().unwrap().contains("groupby") { + self.register_data(&ctx).await?; + } let iterations = self.common.iterations; let mut benchmark_run = BenchmarkRun::new(); @@ -81,17 +99,22 @@ impl RunOpt { let sql = queries.get_query(query_id)?; println!("Q{query_id}: {sql}"); + let mut millis = Vec::with_capacity(iterations); for i in 1..=iterations { let start = Instant::now(); let results = ctx.sql(sql).await?.collect().await?; let elapsed = start.elapsed(); let ms = elapsed.as_secs_f64() * 1000.0; + millis.push(ms); let row_count: usize = results.iter().map(|b| b.num_rows()).sum(); println!( "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows" ); benchmark_run.write_iter(elapsed, row_count); } + let avg = millis.iter().sum::() / millis.len() as f64; + println!("Query {query_id} avg time: {avg:.2} ms"); + if self.common.debug { ctx.sql(sql).await?.explain(false, false)?.show().await?; } diff --git a/datafusion-cli/CONTRIBUTING.md b/datafusion-cli/CONTRIBUTING.md new file mode 100644 index 000000000000..4b464dffc57c --- /dev/null +++ b/datafusion-cli/CONTRIBUTING.md @@ -0,0 +1,75 @@ + + +# Development instructions + +## Running Tests + +Tests can be run using `cargo` + +```shell +cargo test +``` + +## Running Storage Integration Tests + +By default, storage integration tests are not run. To run them you will need to set `TEST_STORAGE_INTEGRATION=1` and +then provide the necessary configuration for that object store. + +For some of the tests, [snapshots](https://datafusion.apache.org/contributor-guide/testing.html#snapshot-testing) are used. + +### AWS + +To test the S3 integration against [Minio](https://github.com/minio/minio) + +First start up a container with Minio and load test files. + +```shell +docker run -d \ + --name datafusion-test-minio \ + -p 9000:9000 \ + -e MINIO_ROOT_USER=TEST-DataFusionLogin \ + -e MINIO_ROOT_PASSWORD=TEST-DataFusionPassword \ + -v $(pwd)/../datafusion/core/tests/data:/source \ + quay.io/minio/minio server /data + +docker exec datafusion-test-minio /bin/sh -c "\ + mc ready local + mc alias set localminio http://localhost:9000 TEST-DataFusionLogin TEST-DataFusionPassword && \ + mc mb localminio/data && \ + mc cp -r /source/* localminio/data" +``` + +Setup environment + +```shell +export TEST_STORAGE_INTEGRATION=1 +export AWS_ACCESS_KEY_ID=TEST-DataFusionLogin +export AWS_SECRET_ACCESS_KEY=TEST-DataFusionPassword +export AWS_ENDPOINT=http://127.0.0.1:9000 +export AWS_ALLOW_HTTP=true +``` + +Note that `AWS_ENDPOINT` is set without slash at the end. + +Run tests + +```shell +cargo test +``` diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 0fcc8cecd390..258fd995a73e 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -30,12 +30,16 @@ rust-version = { workspace = true } [package.metadata.docs.rs] all-features = true +[features] +default = [] +backtrace = ["datafusion/backtrace"] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } -aws-config = "1.5.17" +aws-config = "1.5.18" aws-credential-types = "1.2.0" -clap = { version = "4.5.30", features = ["derive", "cargo"] } +clap = { version = "4.5.31", features = ["derive", "cargo"] } datafusion = { workspace = true, features = [ "avro", "crypto_expressions", @@ -63,5 +67,7 @@ url = { workspace = true } [dev-dependencies] assert_cmd = "2.0" ctor = { workspace = true } +insta = { version = "1.41.1", features = ["glob", "filters"] } +insta-cmd = "0.6.0" predicates = "3.0" rstest = { workspace = true } diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index 690f51315dd0..0f4d70c1cca9 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -26,6 +26,12 @@ use crate::{ object_storage::get_object_store, print_options::{MaxRows, PrintOptions}, }; +use futures::StreamExt; +use std::collections::HashMap; +use std::fs::File; +use std::io::prelude::*; +use std::io::BufReader; + use datafusion::common::instant::Instant; use datafusion::common::{plan_datafusion_err, plan_err}; use datafusion::config::ConfigFileType; @@ -35,15 +41,13 @@ use datafusion::logical_expr::{DdlStatement, LogicalPlan}; use datafusion::physical_plan::execution_plan::EmissionType; use datafusion::physical_plan::{execute_stream, ExecutionPlanProperties}; use datafusion::sql::parser::{DFParser, Statement}; -use datafusion::sql::sqlparser; use datafusion::sql::sqlparser::dialect::dialect_from_str; + +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::physical_plan::spill::get_record_batch_memory_size; +use datafusion::sql::sqlparser; use rustyline::error::ReadlineError; use rustyline::Editor; -use std::collections::HashMap; -use std::fs::File; -use std::io::prelude::*; -use std::io::BufReader; -use std::sync::Arc; use tokio::signal; /// run and execute SQL statements and commands, against a context with the given print options @@ -225,17 +229,18 @@ pub(super) async fn exec_and_print( for statement in statements { let adjusted = AdjustedPrintOptions::new(print_options.clone()).with_statement(&statement); + let plan = create_plan(ctx, statement).await?; let adjusted = adjusted.with_plan(&plan); let df = ctx.execute_logical_plan(plan).await?; let physical_plan = df.create_physical_plan().await?; - let is_unbounded = physical_plan.boundedness().is_unbounded(); - let mut stream = execute_stream(Arc::clone(&physical_plan), task_ctx.clone())?; + // Track memory usage for the query result if it's bounded + let mut reservation = + MemoryConsumer::new("DataFusion-Cli").register(task_ctx.memory_pool()); - // Both bounded and unbounded streams are streaming prints - if is_unbounded { + if physical_plan.boundedness().is_unbounded() { if physical_plan.pipeline_behavior() == EmissionType::Final { return plan_err!( "The given query can generate a valid result only once \ @@ -244,43 +249,37 @@ pub(super) async fn exec_and_print( } // As the input stream comes, we can generate results. // However, memory safety is not guaranteed. - print_options - .print_stream(MaxRows::Unlimited, stream, now) - .await?; + let stream = execute_stream(physical_plan, task_ctx.clone())?; + print_options.print_stream(stream, now).await?; } else { // Bounded stream; collected results size is limited by the maxrows option let schema = physical_plan.schema(); + let mut stream = execute_stream(physical_plan, task_ctx.clone())?; + let mut results = vec![]; + let mut row_count = 0_usize; let max_rows = match print_options.maxrows { MaxRows::Unlimited => usize::MAX, MaxRows::Limited(n) => n, }; - let stdout = std::io::stdout(); - let mut writer = stdout.lock(); - - // If we don't want to print the table, we should use the streaming print same as above - if print_options.format != PrintFormat::Table - && print_options.format != PrintFormat::Automatic - { - print_options - .print_stream(print_options.maxrows, stream, now) - .await?; - continue; + while let Some(batch) = stream.next().await { + let batch = batch?; + let curr_num_rows = batch.num_rows(); + // Stop collecting results if the number of rows exceeds the limit + // results batch should include the last batch that exceeds the limit + if row_count < max_rows + curr_num_rows { + // Try to grow the reservation to accommodate the batch in memory + reservation.try_grow(get_record_batch_memory_size(&batch))?; + results.push(batch); + } + row_count += curr_num_rows; } - - // into_inner will finalize the print options to table if it's automatic adjusted .into_inner() - .print_table_batch( - print_options, - schema, - &mut stream, - max_rows, - &mut writer, - now, - ) - .await?; + .print_batches(schema, &results, now, row_count)?; + reservation.free(); } } + Ok(()) } diff --git a/datafusion-cli/src/print_format.rs b/datafusion-cli/src/print_format.rs index ed3f03781c43..1fc949593512 100644 --- a/datafusion-cli/src/print_format.rs +++ b/datafusion-cli/src/print_format.rs @@ -23,10 +23,8 @@ use crate::print_options::MaxRows; use arrow::csv::writer::WriterBuilder; use arrow::datatypes::SchemaRef; -use arrow::error::ArrowError; use arrow::json::{ArrayWriter, LineDelimitedWriter}; use arrow::record_batch::RecordBatch; -use arrow::util::display::{ArrayFormatter, ValueFormatter}; use arrow::util::pretty::pretty_format_batches_with_options; use datafusion::common::format::DEFAULT_CLI_FORMAT_OPTIONS; use datafusion::error::Result; @@ -211,145 +209,6 @@ impl PrintFormat { } Ok(()) } - - #[allow(clippy::too_many_arguments)] - pub fn process_batch( - &self, - batch: &RecordBatch, - schema: SchemaRef, - preview_batches: &mut Vec, - preview_row_count: &mut usize, - preview_limit: usize, - precomputed_widths: &mut Option>, - header_printed: &mut bool, - writer: &mut dyn std::io::Write, - ) -> Result<()> { - if precomputed_widths.is_none() { - preview_batches.push(batch.clone()); - *preview_row_count += batch.num_rows(); - if *preview_row_count >= preview_limit { - let widths = - Self::compute_column_widths(self, preview_batches, schema.clone())?; - *precomputed_widths = Some(widths.clone()); - Self::print_header(self, &schema, &widths, writer)?; - *header_printed = true; - for preview_batch in preview_batches.drain(..) { - Self::print_batch_with_widths(self, &preview_batch, &widths, writer)?; - } - } - } else { - let widths = precomputed_widths.as_ref().unwrap(); - if !*header_printed { - Self::print_header(self, &schema, widths, writer)?; - *header_printed = true; - } - Self::print_batch_with_widths(self, batch, widths, writer)?; - } - Ok(()) - } - - pub fn compute_column_widths( - &self, - batches: &Vec, - schema: SchemaRef, - ) -> Result> { - let mut widths: Vec = - schema.fields().iter().map(|f| f.name().len()).collect(); - for batch in batches { - let formatters = batch - .columns() - .iter() - .map(|c| ArrayFormatter::try_new(c.as_ref(), &DEFAULT_CLI_FORMAT_OPTIONS)) - .collect::, ArrowError>>()?; - for row in 0..batch.num_rows() { - for (i, formatter) in formatters.iter().enumerate() { - let cell = formatter.value(row); - widths[i] = widths[i].max(cell.to_string().len()); - } - } - } - Ok(widths) - } - - pub fn print_header( - &self, - schema: &SchemaRef, - widths: &[usize], - writer: &mut dyn std::io::Write, - ) -> Result<()> { - Self::print_border(widths, writer)?; - - let header: Vec = schema - .fields() - .iter() - .enumerate() - .map(|(i, field)| Self::pad_cell(field.name(), widths[i])) - .collect(); - writeln!(writer, "| {} |", header.join(" | "))?; - - Self::print_border(widths, writer)?; - Ok(()) - } - - pub fn print_batch_with_widths( - &self, - batch: &RecordBatch, - widths: &[usize], - writer: &mut dyn std::io::Write, - ) -> Result<()> { - let formatters = batch - .columns() - .iter() - .map(|c| ArrayFormatter::try_new(c.as_ref(), &DEFAULT_CLI_FORMAT_OPTIONS)) - .collect::, ArrowError>>()?; - for row in 0..batch.num_rows() { - let cells: Vec = formatters - .iter() - .enumerate() - .map(|(i, formatter)| Self::pad_value(&formatter.value(row), widths[i])) - .collect(); - writeln!(writer, "| {} |", cells.join(" | "))?; - } - Ok(()) - } - - pub fn print_dotted_line( - &self, - widths: &[usize], - writer: &mut dyn std::io::Write, - ) -> Result<()> { - let cells: Vec = widths - .iter() - .map(|&w| format!(" {: Result<()> { - let cells: Vec = widths.iter().map(|&w| "-".repeat(w + 2)).collect(); - writeln!(writer, "+{}+", cells.join("+"))?; - Ok(()) - } - - fn print_border(widths: &[usize], writer: &mut dyn std::io::Write) -> Result<()> { - let cells: Vec = widths.iter().map(|&w| "-".repeat(w + 2)).collect(); - writeln!(writer, "+{}+", cells.join("+"))?; - Ok(()) - } - - fn pad_cell(cell: &str, width: usize) -> String { - format!("{: String { - let s = formatter.try_to_string().unwrap_or_default(); - format!("{: = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_batch_with_same_widths() { - let batch = three_column_batch(); - let widths = vec![1, 1, 1]; - let mut writer = Vec::new(); - let format = PrintFormat::Table; - format - .print_batch_with_widths(&batch, &widths, &mut writer) - .unwrap(); - let expected = &["| 1 | 4 | 7 |", "| 2 | 5 | 8 |", "| 3 | 6 | 9 |"]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_batch_with_different_widths() { - let batch = three_column_batch_with_widths(); - let widths = vec![7, 5, 6]; - let mut writer = Vec::new(); - let format = PrintFormat::Table; - format - .print_batch_with_widths(&batch, &widths, &mut writer) - .unwrap(); - let expected = &[ - "| 1 | 42222 | 7 |", - "| 2222222 | 5 | 8 |", - "| 3 | 6 | 922222 |", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_dotted_line() { - let widths = vec![1, 1, 1]; - let mut writer = Vec::new(); - let format = PrintFormat::Table; - format.print_dotted_line(&widths, &mut writer).unwrap(); - let expected = &["| . | . | . |"]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_bottom_border() { - let widths = vec![1, 1, 1]; - let mut writer = Vec::new(); - let format = PrintFormat::Table; - format.print_bottom_border(&widths, &mut writer).unwrap(); - let expected = &["+---+---+---+"]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_batches_with_maxrows() { - let batch = one_column_batch(); - let schema = one_column_schema(); - let format = PrintFormat::Table; - - // should print out entire output with no truncation if unlimited or - // limit greater than number of batches or equal to the number of batches - for max_rows in [MaxRows::Unlimited, MaxRows::Limited(5), MaxRows::Limited(3)] { - let mut writer = Vec::new(); - format - .print_batches( - &mut writer, - schema.clone(), - &[batch.clone()], - max_rows, - true, - ) - .unwrap(); - let expected = &[ - "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - // should truncate output if limit is less than number of batches - let mut writer = Vec::new(); - format - .print_batches( - &mut writer, - schema.clone(), - &[batch.clone()], - MaxRows::Limited(1), - true, - ) - .unwrap(); - let expected = &[ - "+---+", "| a |", "+---+", "| 1 |", "| . |", "| . |", "| . |", "+---+", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - // test print_batch with different batch widths - // and preview count is less than the first batch - #[test] - fn test_print_batches_with_preview_count_less_than_first_batch() { - let batch = three_column_batch_with_widths(); - let schema = three_column_schema(); - let format = PrintFormat::Table; - let preview_limit = 2; - let mut preview_batches = Vec::new(); - let mut preview_row_count = 0; - let mut precomputed_widths = None; - let mut header_printed = false; - let mut writer = Vec::new(); - - format - .process_batch( - &batch, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - let expected = &[ - "+---------+-------+--------+", - "| a | b | c |", - "+---------+-------+--------+", - "| 1 | 42222 | 7 |", - "| 2222222 | 5 | 8 |", - "| 3 | 6 | 922222 |", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_batches_with_preview_and_later_batches() { - let batch1 = three_column_batch(); - let batch2 = three_column_batch_with_widths(); - let schema = three_column_schema(); - let format = PrintFormat::Table; - // preview limit is less than the first batch - // so the second batch if it's width is greater than the first batch, it will be unformatted - let preview_limit = 2; - let mut preview_batches = Vec::new(); - let mut preview_row_count = 0; - let mut precomputed_widths = None; - let mut header_printed = false; - let mut writer = Vec::new(); - - format - .process_batch( - &batch1, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - format - .process_batch( - &batch2, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - format - .process_batch( - &batch1, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - let expected = &[ - "+---+---+---+", - "| a | b | c |", - "+---+---+---+", - "| 1 | 4 | 7 |", - "| 2 | 5 | 8 |", - "| 3 | 6 | 9 |", - "| 1 | 42222 | 7 |", - "| 2222222 | 5 | 8 |", - "| 3 | 6 | 922222 |", - "| 1 | 4 | 7 |", - "| 2 | 5 | 8 |", - "| 3 | 6 | 9 |", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - - #[test] - fn test_print_batches_with_preview_cover_later_batches() { - let batch1 = three_column_batch(); - let batch2 = three_column_batch_with_widths(); - let schema = three_column_schema(); - let format = PrintFormat::Table; - // preview limit is greater than the first batch - let preview_limit = 4; - let mut preview_batches = Vec::new(); - let mut preview_row_count = 0; - let mut precomputed_widths = None; - let mut header_printed = false; - let mut writer = Vec::new(); - - format - .process_batch( - &batch1, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - format - .process_batch( - &batch2, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - format - .process_batch( - &batch1, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - &mut writer, - ) - .unwrap(); - - let expected = &[ - "+---------+-------+--------+", - "| a | b | c |", - "+---------+-------+--------+", - "| 1 | 4 | 7 |", - "| 2 | 5 | 8 |", - "| 3 | 6 | 9 |", - "| 1 | 42222 | 7 |", - "| 2222222 | 5 | 8 |", - "| 3 | 6 | 922222 |", - "| 1 | 4 | 7 |", - "| 2 | 5 | 8 |", - "| 3 | 6 | 9 |", - ]; - let binding = String::from_utf8(writer.clone()).unwrap(); - let actual: Vec<_> = binding.trim_end().split('\n').collect(); - assert_eq!(actual, expected); - } - #[derive(Debug)] struct PrintBatchesTest { format: PrintFormat, @@ -1136,19 +672,6 @@ mod tests { .unwrap() } - /// Return a batch with three columns and three rows, but with different widths - fn three_column_batch_with_widths() -> RecordBatch { - RecordBatch::try_new( - three_column_schema(), - vec![ - Arc::new(Int32Array::from(vec![1, 2222222, 3])), - Arc::new(Int32Array::from(vec![42222, 5, 6])), - Arc::new(Int32Array::from(vec![7, 8, 922222])), - ], - ) - .unwrap() - } - /// Return a schema with one column fn one_column_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index 8dd7ca9c8081..9557e783e8a7 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -29,7 +29,6 @@ use datafusion::common::DataFusionError; use datafusion::error::Result; use datafusion::physical_plan::RecordBatchStream; -use datafusion::execution::SendableRecordBatchStream; use futures::StreamExt; #[derive(Debug, Clone, PartialEq, Copy)] @@ -75,6 +74,27 @@ pub struct PrintOptions { pub color: bool, } +// Returns the query execution details formatted +fn get_execution_details_formatted( + row_count: usize, + maxrows: MaxRows, + query_start_time: Instant, +) -> String { + let nrows_shown_msg = match maxrows { + MaxRows::Limited(nrows) if nrows < row_count => { + format!("(First {nrows} displayed. Use --maxrows to adjust)") + } + _ => String::new(), + }; + + format!( + "{} row(s) fetched. {}\nElapsed {:.3} seconds.\n", + row_count, + nrows_shown_msg, + query_start_time.elapsed().as_secs_f64() + ) +} + impl PrintOptions { /// Print the batches to stdout using the specified format pub fn print_batches( @@ -90,7 +110,7 @@ impl PrintOptions { self.format .print_batches(&mut writer, schema, batches, self.maxrows, true)?; - let formatted_exec_details = self.get_execution_details_formatted( + let formatted_exec_details = get_execution_details_formatted( row_count, if self.format == PrintFormat::Table { self.maxrows @@ -107,125 +127,9 @@ impl PrintOptions { Ok(()) } - pub async fn print_table_batch( - &self, - print_options: &PrintOptions, - schema: SchemaRef, - stream: &mut SendableRecordBatchStream, - max_rows: usize, - writer: &mut dyn std::io::Write, - now: Instant, - ) -> Result<()> { - let preview_limit: usize = 1000; - let mut preview_batches: Vec = vec![]; - let mut preview_row_count = 0_usize; - let mut total_count = 0_usize; - let mut precomputed_widths: Option> = None; - let mut header_printed = false; - let mut max_rows_reached = false; - - while let Some(batch) = stream.next().await { - let batch = batch?; - let batch_rows = batch.num_rows(); - - if !max_rows_reached && total_count < max_rows { - if total_count + batch_rows > max_rows { - let needed = max_rows - total_count; - let batch_to_print = batch.slice(0, needed); - print_options.format.process_batch( - &batch_to_print, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - writer, - )?; - if precomputed_widths.is_none() { - let widths = print_options - .format - .compute_column_widths(&preview_batches, schema.clone())?; - precomputed_widths = Some(widths.clone()); - if !header_printed { - print_options - .format - .print_header(&schema, &widths, writer)?; - } - for preview_batch in preview_batches.drain(..) { - print_options.format.print_batch_with_widths( - &preview_batch, - &widths, - writer, - )?; - } - } - if let Some(ref widths) = precomputed_widths { - for _ in 0..3 { - print_options.format.print_dotted_line(widths, writer)?; - } - print_options.format.print_bottom_border(widths, writer)?; - } - max_rows_reached = true; - } else { - print_options.format.process_batch( - &batch, - schema.clone(), - &mut preview_batches, - &mut preview_row_count, - preview_limit, - &mut precomputed_widths, - &mut header_printed, - writer, - )?; - } - } - - total_count += batch_rows; - } - - if !max_rows_reached { - if precomputed_widths.is_none() && !preview_batches.is_empty() { - let widths = print_options - .format - .compute_column_widths(&preview_batches, schema.clone())?; - precomputed_widths = Some(widths); - if !header_printed { - print_options.format.print_header( - &schema, - precomputed_widths.as_ref().unwrap(), - writer, - )?; - } - for preview_batch in preview_batches.drain(..) { - print_options.format.print_batch_with_widths( - &preview_batch, - precomputed_widths.as_ref().unwrap(), - writer, - )?; - } - } - if let Some(ref widths) = precomputed_widths { - print_options.format.print_bottom_border(widths, writer)?; - } - } - - let formatted_exec_details = print_options.get_execution_details_formatted( - total_count, - print_options.maxrows, - now, - ); - if !print_options.quiet { - writeln!(writer, "{}", formatted_exec_details)?; - } - - Ok(()) - } - /// Print the stream to stdout using the specified format pub async fn print_stream( &self, - max_rows: MaxRows, mut stream: Pin>, query_start_time: Instant, ) -> Result<()> { @@ -235,49 +139,30 @@ impl PrintOptions { )); }; - let max_count = match self.maxrows { - MaxRows::Unlimited => usize::MAX, - MaxRows::Limited(n) => n, - }; - let stdout = std::io::stdout(); let mut writer = stdout.lock(); let mut row_count = 0_usize; let mut with_header = true; - let mut max_rows_reached = false; while let Some(maybe_batch) = stream.next().await { let batch = maybe_batch?; - let curr_batch_rows = batch.num_rows(); - if !max_rows_reached && row_count < max_count { - if row_count + curr_batch_rows > max_count { - let needed = max_count - row_count; - let batch_to_print = batch.slice(0, needed); - self.format.print_batches( - &mut writer, - batch.schema(), - &[batch_to_print], - max_rows, - with_header, - )?; - max_rows_reached = true; - } else { - self.format.print_batches( - &mut writer, - batch.schema(), - &[batch], - max_rows, - with_header, - )?; - } - } - row_count += curr_batch_rows; + row_count += batch.num_rows(); + self.format.print_batches( + &mut writer, + batch.schema(), + &[batch], + MaxRows::Unlimited, + with_header, + )?; with_header = false; } - let formatted_exec_details = - self.get_execution_details_formatted(row_count, max_rows, query_start_time); + let formatted_exec_details = get_execution_details_formatted( + row_count, + MaxRows::Unlimited, + query_start_time, + ); if !self.quiet { writeln!(writer, "{formatted_exec_details}")?; @@ -285,26 +170,4 @@ impl PrintOptions { Ok(()) } - - // Returns the query execution details formatted - pub fn get_execution_details_formatted( - &self, - row_count: usize, - maxrows: MaxRows, - query_start_time: Instant, - ) -> String { - let nrows_shown_msg = match maxrows { - MaxRows::Limited(nrows) if nrows < row_count => { - format!("(First {nrows} displayed. Use --maxrows to adjust)") - } - _ => String::new(), - }; - - format!( - "{} row(s) fetched. {}\nElapsed {:.3} seconds.\n", - row_count, - nrows_shown_msg, - query_start_time.elapsed().as_secs_f64() - ) - } } diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 6d7b8685a7dd..a54a920e97bb 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -17,10 +17,24 @@ use std::process::Command; -use assert_cmd::prelude::{CommandCargoExt, OutputAssertExt}; -use predicates::prelude::predicate; use rstest::rstest; +use insta::{glob, Settings}; +use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; +use std::{env, fs}; + +fn cli() -> Command { + Command::new(get_cargo_bin("datafusion-cli")) +} + +fn make_settings() -> Settings { + let mut settings = Settings::clone_current(); + settings.set_prepend_module_to_snapshot(false); + settings.add_filter(r"Elapsed .* seconds\.", "[ELAPSED]"); + settings.add_filter(r"DataFusion CLI v.*", "[CLI_VERSION]"); + settings +} + #[cfg(test)] #[ctor::ctor] fn init() { @@ -28,41 +42,106 @@ fn init() { let _ = env_logger::try_init(); } -// Disabled due to https://github.com/apache/datafusion/issues/10793 -#[cfg(not(target_family = "windows"))] #[rstest] -#[case::exec_from_commands( - ["--command", "select 1", "--format", "json", "-q"], - "[{\"Int64(1)\":1}]\n" -)] #[case::exec_multiple_statements( - ["--command", "select 1; select 2;", "--format", "json", "-q"], - "[{\"Int64(1)\":1}]\n[{\"Int64(2)\":2}]\n" + "statements", + ["--command", "select 1; select 2;", "-q"], )] #[case::exec_backslash( - ["--file", "tests/data/backslash.txt", "--format", "json", "-q"], - "[{\"Utf8(\\\"\\\\\\\")\":\"\\\\\",\"Utf8(\\\"\\\\\\\\\\\")\":\"\\\\\\\\\",\"Utf8(\\\"\\\\\\\\\\\\\\\\\\\\\\\")\":\"\\\\\\\\\\\\\\\\\\\\\",\"Utf8(\\\"dsdsds\\\\\\\\\\\\\\\\\\\")\":\"dsdsds\\\\\\\\\\\\\\\\\",\"Utf8(\\\"\\\\t\\\")\":\"\\\\t\",\"Utf8(\\\"\\\\0\\\")\":\"\\\\0\",\"Utf8(\\\"\\\\n\\\")\":\"\\\\n\"}]\n" + "backslash", + ["--file", "tests/sql/backslash.sql", "--format", "json", "-q"], )] #[case::exec_from_files( - ["--file", "tests/data/sql.txt", "--format", "json", "-q"], - "[{\"Int64(1)\":1}]\n" + "files", + ["--file", "tests/sql/select.sql", "-q"], )] #[case::set_batch_size( - ["--command", "show datafusion.execution.batch_size", "--format", "json", "-q", "-b", "1"], - "[{\"name\":\"datafusion.execution.batch_size\",\"value\":\"1\"}]\n" -)] - -/// Add case fixed issue: https://github.com/apache/datafusion/issues/14920 -#[case::exec_from_commands( - ["--command", "SELECT * FROM generate_series(1, 5) t1(v1) ORDER BY v1 DESC;", "--format", "table", "-q"], - "+----+\n| v1 |\n+----+\n| 5 |\n| 4 |\n| 3 |\n| 2 |\n| 1 |\n+----+\n" + "batch_size", + ["--command", "show datafusion.execution.batch_size", "-q", "-b", "1"], )] #[test] fn cli_quick_test<'a>( + #[case] snapshot_name: &'a str, #[case] args: impl IntoIterator, - #[case] expected: &str, ) { - let mut cmd = Command::cargo_bin("datafusion-cli").unwrap(); + let mut settings = make_settings(); + settings.set_snapshot_suffix(snapshot_name); + let _bound = settings.bind_to_scope(); + + let mut cmd = cli(); cmd.args(args); - cmd.assert().stdout(predicate::eq(expected)); + + assert_cmd_snapshot!(cmd); +} + +#[rstest] +#[case("csv")] +#[case("tsv")] +#[case("table")] +#[case("json")] +#[case("nd-json")] +#[case("automatic")] +#[test] +fn test_cli_format<'a>(#[case] format: &'a str) { + let mut settings = make_settings(); + settings.set_snapshot_suffix(format); + let _bound = settings.bind_to_scope(); + + let mut cmd = cli(); + cmd.args(["--command", "select 1", "-q", "--format", format]); + + assert_cmd_snapshot!(cmd); +} + +#[tokio::test] +async fn test_cli() { + if env::var("TEST_STORAGE_INTEGRATION").is_err() { + eprintln!("Skipping external storages integration tests"); + return; + } + + let settings = make_settings(); + let _bound = settings.bind_to_scope(); + + glob!("sql/integration/*.sql", |path| { + let input = fs::read_to_string(path).unwrap(); + assert_cmd_snapshot!(cli().pass_stdin(input)) + }); +} + +#[tokio::test] +async fn test_aws_options() { + // Separate test is needed to pass aws as options in sql and not via env + + if env::var("TEST_STORAGE_INTEGRATION").is_err() { + eprintln!("Skipping external storages integration tests"); + return; + } + + let settings = make_settings(); + let _bound = settings.bind_to_scope(); + + let access_key_id = + env::var("AWS_ACCESS_KEY_ID").expect("AWS_ACCESS_KEY_ID is not set"); + let secret_access_key = + env::var("AWS_SECRET_ACCESS_KEY").expect("AWS_SECRET_ACCESS_KEY is not set"); + let endpoint_url = env::var("AWS_ENDPOINT").expect("AWS_ENDPOINT is not set"); + + let input = format!( + r#"CREATE EXTERNAL TABLE CARS +STORED AS CSV +LOCATION 's3://data/cars.csv' +OPTIONS( + 'aws.access_key_id' '{}', + 'aws.secret_access_key' '{}', + 'aws.endpoint' '{}', + 'aws.allow_http' 'true' +); + +SELECT * FROM CARS limit 1; +"#, + access_key_id, secret_access_key, endpoint_url + ); + + assert_cmd_snapshot!(cli().env_clear().pass_stdin(input)); } diff --git a/datafusion-cli/tests/snapshots/aws_options.snap b/datafusion-cli/tests/snapshots/aws_options.snap new file mode 100644 index 000000000000..283cf57bc662 --- /dev/null +++ b/datafusion-cli/tests/snapshots/aws_options.snap @@ -0,0 +1,25 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: [] + stdin: "CREATE EXTERNAL TABLE CARS\nSTORED AS CSV\nLOCATION 's3://data/cars.csv'\nOPTIONS(\n 'aws.access_key_id' 'TEST-DataFusionLogin',\n 'aws.secret_access_key' 'TEST-DataFusionPassword',\n 'aws.endpoint' 'http://127.0.0.1:9000',\n 'aws.allow_http' 'true'\n);\n\nSELECT * FROM CARS limit 1;\n" +--- +success: true +exit_code: 0 +----- stdout ----- +[CLI_VERSION] +0 row(s) fetched. +[ELAPSED] + ++-----+-------+---------------------+ +| car | speed | time | ++-----+-------+---------------------+ +| red | 20.0 | 1996-04-12T12:05:03 | ++-----+-------+---------------------+ +1 row(s) fetched. +[ELAPSED] + +\q + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli@load_local_csv.sql.snap b/datafusion-cli/tests/snapshots/cli@load_local_csv.sql.snap new file mode 100644 index 000000000000..029d5f8d5b9f --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli@load_local_csv.sql.snap @@ -0,0 +1,26 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: [] + stdin: "CREATE EXTERNAL TABLE CARS\nSTORED AS CSV\nLOCATION '../datafusion/core/tests/data/cars.csv'\nOPTIONS ('has_header' 'TRUE');\n\nSELECT * FROM CARS limit 1;" +input_file: tests/sql/load_local_csv.sql +--- +success: true +exit_code: 0 +----- stdout ----- +[CLI_VERSION] +0 row(s) fetched. +[ELAPSED] + ++-----+-------+---------------------+ +| car | speed | time | ++-----+-------+---------------------+ +| red | 20.0 | 1996-04-12T12:05:03 | ++-----+-------+---------------------+ +1 row(s) fetched. +[ELAPSED] + +\q + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli@load_s3_csv.sql.snap b/datafusion-cli/tests/snapshots/cli@load_s3_csv.sql.snap new file mode 100644 index 000000000000..858989621a1f --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli@load_s3_csv.sql.snap @@ -0,0 +1,26 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: [] + stdin: "CREATE EXTERNAL TABLE CARS\nSTORED AS CSV\nLOCATION 's3://data/cars.csv';\n\nSELECT * FROM CARS limit 1;" +input_file: tests/sql/load_s3_csv.sql +--- +success: true +exit_code: 0 +----- stdout ----- +[CLI_VERSION] +0 row(s) fetched. +[ELAPSED] + ++-----+-------+---------------------+ +| car | speed | time | ++-----+-------+---------------------+ +| red | 20.0 | 1996-04-12T12:05:03 | ++-----+-------+---------------------+ +1 row(s) fetched. +[ELAPSED] + +\q + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli@select.sql.snap b/datafusion-cli/tests/snapshots/cli@select.sql.snap new file mode 100644 index 000000000000..c137d9fe2b13 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli@select.sql.snap @@ -0,0 +1,23 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: [] + stdin: select 1; +input_file: tests/sql/select.sql +--- +success: true +exit_code: 0 +----- stdout ----- +[CLI_VERSION] ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ +1 row(s) fetched. +[ELAPSED] + +\q + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@automatic.snap b/datafusion-cli/tests/snapshots/cli_format@automatic.snap new file mode 100644 index 000000000000..2591f493e90a --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@automatic.snap @@ -0,0 +1,21 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - automatic +--- +success: true +exit_code: 0 +----- stdout ----- ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@csv.snap b/datafusion-cli/tests/snapshots/cli_format@csv.snap new file mode 100644 index 000000000000..c41b042298eb --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@csv.snap @@ -0,0 +1,18 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - csv +--- +success: true +exit_code: 0 +----- stdout ----- +Int64(1) +1 + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@json.snap b/datafusion-cli/tests/snapshots/cli_format@json.snap new file mode 100644 index 000000000000..8f804a337cce --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@json.snap @@ -0,0 +1,17 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - json +--- +success: true +exit_code: 0 +----- stdout ----- +[{"Int64(1)":1}] + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@nd-json.snap b/datafusion-cli/tests/snapshots/cli_format@nd-json.snap new file mode 100644 index 000000000000..7b4ce1e2530c --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@nd-json.snap @@ -0,0 +1,17 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - nd-json +--- +success: true +exit_code: 0 +----- stdout ----- +{"Int64(1)":1} + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@table.snap b/datafusion-cli/tests/snapshots/cli_format@table.snap new file mode 100644 index 000000000000..99914182462a --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@table.snap @@ -0,0 +1,21 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - table +--- +success: true +exit_code: 0 +----- stdout ----- ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_format@tsv.snap b/datafusion-cli/tests/snapshots/cli_format@tsv.snap new file mode 100644 index 000000000000..968268c31dd5 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_format@tsv.snap @@ -0,0 +1,18 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1 + - "-q" + - "--format" + - tsv +--- +success: true +exit_code: 0 +----- stdout ----- +Int64(1) +1 + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_quick_test@backslash.snap b/datafusion-cli/tests/snapshots/cli_quick_test@backslash.snap new file mode 100644 index 000000000000..c01699146aa8 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_quick_test@backslash.snap @@ -0,0 +1,17 @@ +--- +source: datafusion-cli/tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--file" + - tests/sql/backslash.sql + - "--format" + - json + - "-q" +--- +success: true +exit_code: 0 +----- stdout ----- +[{"Utf8(\"\\\")":"\\","Utf8(\"\\\\\")":"\\\\","Utf8(\"\\\\\\\\\\\")":"\\\\\\\\\\","Utf8(\"dsdsds\\\\\\\\\")":"dsdsds\\\\\\\\","Utf8(\"\\t\")":"\\t","Utf8(\"\\0\")":"\\0","Utf8(\"\\n\")":"\\n"}] + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_quick_test@batch_size.snap b/datafusion-cli/tests/snapshots/cli_quick_test@batch_size.snap new file mode 100644 index 000000000000..c27d527df0b6 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_quick_test@batch_size.snap @@ -0,0 +1,21 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - show datafusion.execution.batch_size + - "-q" + - "-b" + - "1" +--- +success: true +exit_code: 0 +----- stdout ----- ++---------------------------------+-------+ +| name | value | ++---------------------------------+-------+ +| datafusion.execution.batch_size | 1 | ++---------------------------------+-------+ + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_quick_test@files.snap b/datafusion-cli/tests/snapshots/cli_quick_test@files.snap new file mode 100644 index 000000000000..7c44e41729a1 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_quick_test@files.snap @@ -0,0 +1,19 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--file" + - tests/sql/select.sql + - "-q" +--- +success: true +exit_code: 0 +----- stdout ----- ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ + +----- stderr ----- diff --git a/datafusion-cli/tests/snapshots/cli_quick_test@statements.snap b/datafusion-cli/tests/snapshots/cli_quick_test@statements.snap new file mode 100644 index 000000000000..3b975bb6a927 --- /dev/null +++ b/datafusion-cli/tests/snapshots/cli_quick_test@statements.snap @@ -0,0 +1,24 @@ +--- +source: tests/cli_integration.rs +info: + program: datafusion-cli + args: + - "--command" + - select 1; select 2; + - "-q" +--- +success: true +exit_code: 0 +----- stdout ----- ++----------+ +| Int64(1) | ++----------+ +| 1 | ++----------+ ++----------+ +| Int64(2) | ++----------+ +| 2 | ++----------+ + +----- stderr ----- diff --git a/datafusion-cli/tests/data/backslash.txt b/datafusion-cli/tests/sql/backslash.sql similarity index 100% rename from datafusion-cli/tests/data/backslash.txt rename to datafusion-cli/tests/sql/backslash.sql diff --git a/datafusion-cli/tests/sql/integration/load_local_csv.sql b/datafusion-cli/tests/sql/integration/load_local_csv.sql new file mode 100644 index 000000000000..8920c48c5f5f --- /dev/null +++ b/datafusion-cli/tests/sql/integration/load_local_csv.sql @@ -0,0 +1,6 @@ +CREATE EXTERNAL TABLE CARS +STORED AS CSV +LOCATION '../datafusion/core/tests/data/cars.csv' +OPTIONS ('has_header' 'TRUE'); + +SELECT * FROM CARS limit 1; \ No newline at end of file diff --git a/datafusion-cli/tests/sql/integration/load_s3_csv.sql b/datafusion-cli/tests/sql/integration/load_s3_csv.sql new file mode 100644 index 000000000000..10c2e38b9764 --- /dev/null +++ b/datafusion-cli/tests/sql/integration/load_s3_csv.sql @@ -0,0 +1,5 @@ +CREATE EXTERNAL TABLE CARS +STORED AS CSV +LOCATION 's3://data/cars.csv'; + +SELECT * FROM CARS limit 1; \ No newline at end of file diff --git a/datafusion-cli/tests/data/sql.txt b/datafusion-cli/tests/sql/select.sql similarity index 100% rename from datafusion-cli/tests/data/sql.txt rename to datafusion-cli/tests/sql/select.sql diff --git a/datafusion-examples/Cargo.toml b/datafusion-examples/Cargo.toml index ea3139adac3d..d2bbdd78e3f2 100644 --- a/datafusion-examples/Cargo.toml +++ b/datafusion-examples/Cargo.toml @@ -61,7 +61,7 @@ async-trait = { workspace = true } bytes = { workspace = true } dashmap = { workspace = true } # note only use main datafusion crate for examples -datafusion = { workspace = true, default-features = true, features = ["avro"] } +datafusion = { workspace = true, default-features = true } datafusion-proto = { workspace = true } env_logger = { workspace = true } futures = { workspace = true } diff --git a/datafusion-examples/examples/advanced_parquet_index.rs b/datafusion-examples/examples/advanced_parquet_index.rs index bb1cf3c8f78d..d6cf61c61d73 100644 --- a/datafusion-examples/examples/advanced_parquet_index.rs +++ b/datafusion-examples/examples/advanced_parquet_index.rs @@ -23,8 +23,6 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; -use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; -use arrow::datatypes::SchemaRef; use datafusion::catalog::Session; use datafusion::common::{ internal_datafusion_err, DFSchema, DataFusionError, Result, ScalarValue, @@ -32,7 +30,7 @@ use datafusion::common::{ use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; use datafusion::datasource::physical_plan::{ - parquet::ParquetFileReaderFactory, FileMeta, FileScanConfig, ParquetSource, + FileMeta, FileScanConfig, ParquetFileReaderFactory, ParquetSource, }; use datafusion::datasource::TableProvider; use datafusion::execution::object_store::ObjectStoreUrl; @@ -53,6 +51,8 @@ use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; +use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow::datatypes::SchemaRef; use async_trait::async_trait; use bytes::Bytes; use futures::future::BoxFuture; diff --git a/datafusion-examples/examples/csv_json_opener.rs b/datafusion-examples/examples/csv_json_opener.rs index 574137afe5c9..6dc38a436a0c 100644 --- a/datafusion-examples/examples/csv_json_opener.rs +++ b/datafusion-examples/examples/csv_json_opener.rs @@ -18,15 +18,15 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::datasource::physical_plan::JsonSource; use datafusion::{ assert_batches_eq, - datasource::physical_plan::FileSource, datasource::{ file_format::file_compression_type::FileCompressionType, listing::PartitionedFile, object_store::ObjectStoreUrl, - physical_plan::{CsvSource, FileScanConfig, FileStream, JsonOpener}, + physical_plan::{ + CsvSource, FileScanConfig, FileSource, FileStream, JsonOpener, JsonSource, + }, }, error::Result, physical_plan::metrics::ExecutionPlanMetricsSet, diff --git a/datafusion-examples/examples/planner_api.rs b/datafusion-examples/examples/planner_api.rs index e52f0d78682f..41110a3e0a9c 100644 --- a/datafusion-examples/examples/planner_api.rs +++ b/datafusion-examples/examples/planner_api.rs @@ -17,7 +17,7 @@ use datafusion::error::Result; use datafusion::logical_expr::{LogicalPlan, PlanType}; -use datafusion::physical_plan::displayable; +use datafusion::physical_plan::{displayable, DisplayFormatType}; use datafusion::physical_planner::DefaultPhysicalPlanner; use datafusion::prelude::*; @@ -78,7 +78,11 @@ async fn to_physical_plan_in_one_api_demo( println!( "Physical plan direct from logical plan:\n\n{}\n\n", displayable(physical_plan.as_ref()) - .to_stringified(false, PlanType::InitialPhysicalPlan) + .to_stringified( + false, + PlanType::InitialPhysicalPlan, + DisplayFormatType::Default + ) .plan ); @@ -120,7 +124,11 @@ async fn to_physical_plan_step_by_step_demo( println!( "Final physical plan:\n\n{}\n\n", displayable(physical_plan.as_ref()) - .to_stringified(false, PlanType::InitialPhysicalPlan) + .to_stringified( + false, + PlanType::InitialPhysicalPlan, + DisplayFormatType::Default + ) .plan ); @@ -135,7 +143,11 @@ async fn to_physical_plan_step_by_step_demo( println!( "Optimized physical plan:\n\n{}\n\n", displayable(physical_plan.as_ref()) - .to_stringified(false, PlanType::InitialPhysicalPlan) + .to_stringified( + false, + PlanType::InitialPhysicalPlan, + DisplayFormatType::Default + ) .plan ); diff --git a/datafusion-examples/examples/sql_dialect.rs b/datafusion-examples/examples/sql_dialect.rs index 16aa5be02635..12141847ca36 100644 --- a/datafusion-examples/examples/sql_dialect.rs +++ b/datafusion-examples/examples/sql_dialect.rs @@ -19,7 +19,7 @@ use std::fmt::Display; use datafusion::error::Result; use datafusion::sql::{ - parser::{CopyToSource, CopyToStatement, DFParser, Statement}, + parser::{CopyToSource, CopyToStatement, DFParser, DFParserBuilder, Statement}, sqlparser::{keywords::Keyword, parser::ParserError, tokenizer::Token}, }; @@ -46,9 +46,9 @@ struct MyParser<'a> { df_parser: DFParser<'a>, } -impl MyParser<'_> { - fn new(sql: &str) -> Result { - let df_parser = DFParser::new(sql)?; +impl<'a> MyParser<'a> { + fn new(sql: &'a str) -> Result { + let df_parser = DFParserBuilder::new(sql).build()?; Ok(Self { df_parser }) } diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index cf475263535a..9ac8423042d3 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -103,6 +103,8 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool { // - AGGREGATE and WINDOW should not end up in filter conditions, except maybe in some edge cases // - Can `Wildcard` be considered as a `Literal`? // - ScalarVariable could be `applicable`, but that would require access to the context + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::AggregateFunction { .. } | Expr::WindowFunction { .. } | Expr::Wildcard { .. } @@ -532,6 +534,7 @@ pub fn describe_partition(partition: &Partition) -> (&str, usize, Vec<&str>) { #[cfg(test)] mod tests { use async_trait::async_trait; + use datafusion_common::config::TableOptions; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnv; use futures::FutureExt; @@ -1068,5 +1071,13 @@ mod tests { fn as_any(&self) -> &dyn Any { unimplemented!() } + + fn table_options(&self) -> &TableOptions { + unimplemented!() + } + + fn table_options_mut(&mut self) -> &mut TableOptions { + unimplemented!() + } } } diff --git a/datafusion/catalog/src/session.rs b/datafusion/catalog/src/session.rs index db49529ac43f..9dd870e43568 100644 --- a/datafusion/catalog/src/session.rs +++ b/datafusion/catalog/src/session.rs @@ -16,7 +16,7 @@ // under the License. use async_trait::async_trait; -use datafusion_common::config::ConfigOptions; +use datafusion_common::config::{ConfigOptions, TableOptions}; use datafusion_common::{DFSchema, Result}; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnv; @@ -120,6 +120,18 @@ pub trait Session: Send + Sync { fn execution_props(&self) -> &ExecutionProps; fn as_any(&self) -> &dyn Any; + + /// Return the table options + fn table_options(&self) -> &TableOptions; + + /// return the TableOptions options with its extensions + fn default_table_options(&self) -> TableOptions { + self.table_options() + .combine_with_session_config(self.config_options()) + } + + /// Returns a mutable reference to [`TableOptions`] + fn table_options_mut(&mut self) -> &mut TableOptions; } /// Create a new task context instance from Session diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 76f07be95c60..a607f796fc9c 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -58,12 +58,12 @@ base64 = "0.22.1" half = { workspace = true } hashbrown = { workspace = true } indexmap = { workspace = true } -libc = "0.2.170" +libc = "0.2.171" log = { workspace = true } object_store = { workspace = true, optional = true } parquet = { workspace = true, optional = true, default-features = true } paste = "1.0.15" -pyo3 = { version = "0.23.3", optional = true } +pyo3 = { version = "0.23.5", optional = true } recursive = { workspace = true, optional = true } sqlparser = { workspace = true } tokio = { workspace = true } diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 2ea5b3701550..b0f17630c910 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -252,10 +252,18 @@ config_namespace! { /// string length and thus DataFusion can not enforce such limits. pub support_varchar_with_length: bool, default = true + /// If true, `VARCHAR` is mapped to `Utf8View` during SQL planning. + /// If false, `VARCHAR` is mapped to `Utf8` during SQL planning. + /// Default is false. + pub map_varchar_to_utf8view: bool, default = false + /// When set to true, the source locations relative to the original SQL - /// query (i.e. [`Span`](sqlparser::tokenizer::Span)) will be collected + /// query (i.e. [`Span`](https://docs.rs/sqlparser/latest/sqlparser/tokenizer/struct.Span.html)) will be collected /// and recorded in the logical plan nodes. pub collect_spans: bool, default = false + + /// Specifies the recursion depth limit when parsing complex SQL Queries + pub recursion_limit: usize, default = 50 } } @@ -708,6 +716,10 @@ config_namespace! { /// When set to true, the explain statement will print schema information pub show_schema: bool, default = false + + /// Display format of explain. Default is "indent". + /// When set to "tree", it will print the plan in a tree-rendered format. + pub format: String, default = "indent".to_string() } } diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index 99fb179c76a3..65bb40810f18 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -159,22 +159,9 @@ impl DFSchema { } /// Create a new `DFSchema` from a list of Arrow [Field]s - #[allow(deprecated)] pub fn from_unqualified_fields( fields: Fields, metadata: HashMap, - ) -> Result { - Self::from_unqualifed_fields(fields, metadata) - } - - /// Create a new `DFSchema` from a list of Arrow [Field]s - #[deprecated( - since = "40.0.0", - note = "Please use `from_unqualified_fields` instead (this one's name is a typo). This method is subject to be removed soon" - )] - pub fn from_unqualifed_fields( - fields: Fields, - metadata: HashMap, ) -> Result { let field_count = fields.len(); let schema = Arc::new(Schema::new_with_metadata(fields, metadata)); @@ -1047,7 +1034,7 @@ impl SchemaExt for Schema { .iter() .zip(other.fields().iter()) .try_for_each(|(f1, f2)| { - if f1.name() != f2.name() || !DFSchema::datatype_is_logically_equal(f1.data_type(), f2.data_type()) { + if f1.name() != f2.name() || (!DFSchema::datatype_is_logically_equal(f1.data_type(), f2.data_type()) && !can_cast_types(f2.data_type(), f1.data_type())) { _plan_err!( "Inserting query schema mismatch: Expected table field '{}' with type {:?}, \ but got '{}' with type {:?}.", diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 9059ae07e648..367f359ae742 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -2764,8 +2764,10 @@ impl ScalarValue { Ok(scalars) } - // TODO: Support more types after other ScalarValue is wrapped with ArrayRef - /// Get raw data (inner array) inside ScalarValue + #[deprecated( + since = "46.0.0", + note = "This function is obsolete. Use `to_array` instead" + )] pub fn raw_data(&self) -> Result { match self { ScalarValue::List(arr) => Ok(arr.to_owned()), diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 438e2600a66d..fd1fd4164da0 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -40,9 +40,15 @@ nested_expressions = ["datafusion-functions-nested"] # This feature is deprecated. Use the `nested_expressions` feature instead. array_expressions = ["nested_expressions"] # Used to enable the avro format -avro = ["apache-avro", "num-traits", "datafusion-common/avro", "datafusion-datasource/avro"] +avro = ["datafusion-common/avro", "datafusion-datasource-avro"] backtrace = ["datafusion-common/backtrace"] -compression = ["xz2", "bzip2", "flate2", "zstd", "datafusion-datasource/compression"] +compression = [ + "xz2", + "bzip2", + "flate2", + "zstd", + "datafusion-datasource/compression", +] crypto_expressions = ["datafusion-functions/crypto_expressions"] datetime_expressions = ["datafusion-functions/datetime_expressions"] default = [ @@ -61,7 +67,7 @@ encoding_expressions = ["datafusion-functions/encoding_expressions"] # Used for testing ONLY: causes all values to hash to the same value (test for collisions) force_hash_collisions = ["datafusion-physical-plan/force_hash_collisions", "datafusion-common/force_hash_collisions"] math_expressions = ["datafusion-functions/math_expressions"] -parquet = ["datafusion-common/parquet", "dep:parquet"] +parquet = ["datafusion-common/parquet", "dep:parquet", "datafusion-datasource-parquet"] pyarrow = ["datafusion-common/pyarrow", "parquet"] regex_expressions = [ "datafusion-functions/regex_expressions", @@ -73,7 +79,12 @@ recursive_protection = [ "datafusion-physical-optimizer/recursive_protection", "datafusion-sql/recursive_protection", ] -serde = ["dep:serde"] +serde = [ + "dep:serde", + # Enable `#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]` + # statements in `arrow-schema` crate + "arrow-schema/serde", +] string_expressions = ["datafusion-functions/string_expressions"] unicode_expressions = [ "datafusion-sql/unicode_expressions", @@ -82,7 +93,6 @@ unicode_expressions = [ extended_tests = [] [dependencies] -apache-avro = { version = "0.17", optional = true } arrow = { workspace = true } arrow-ipc = { workspace = true } arrow-schema = { workspace = true } @@ -95,6 +105,10 @@ datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } +datafusion-datasource-avro = { workspace = true, optional = true } +datafusion-datasource-csv = { workspace = true } +datafusion-datasource-json = { workspace = true } +datafusion-datasource-parquet = { workspace = true, optional = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } @@ -114,7 +128,6 @@ flate2 = { version = "1.1.0", optional = true } futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } -num-traits = { version = "0.2", optional = true } object_store = { workspace = true } parking_lot = { workspace = true } parquet = { workspace = true, optional = true, default-features = true } diff --git a/datafusion/core/src/datasource/file_format/arrow.rs b/datafusion/core/src/datasource/file_format/arrow.rs index 3614b788af90..6835d9a6da2a 100644 --- a/datafusion/core/src/datasource/file_format/arrow.rs +++ b/datafusion/core/src/datasource/file_format/arrow.rs @@ -301,6 +301,10 @@ impl DisplayAs for ArrowFileSink { FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; write!(f, ")") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/src/datasource/file_format/avro.rs b/datafusion/core/src/datasource/file_format/avro.rs index e7314e839bf2..3723a728a08f 100644 --- a/datafusion/core/src/datasource/file_format/avro.rs +++ b/datafusion/core/src/datasource/file_format/avro.rs @@ -15,163 +15,31 @@ // specific language governing permissions and limitations // under the License. -//! [`AvroFormat`] Apache Avro [`FileFormat`] abstractions - -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::sync::Arc; - -use super::file_compression_type::FileCompressionType; -use super::FileFormat; -use super::FileFormatFactory; -use crate::datasource::avro_to_arrow::read_avro_schema_from_reader; -use crate::datasource::physical_plan::AvroSource; -use crate::error::Result; -use crate::physical_plan::ExecutionPlan; -use crate::physical_plan::Statistics; - -use arrow::datatypes::Schema; -use arrow::datatypes::SchemaRef; -use async_trait::async_trait; -use datafusion_catalog::Session; -use datafusion_common::internal_err; -use datafusion_common::parsers::CompressionTypeVariant; -use datafusion_common::GetExt; -use datafusion_common::DEFAULT_AVRO_EXTENSION; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_physical_expr::PhysicalExpr; -use object_store::{GetResultPayload, ObjectMeta, ObjectStore}; - -#[derive(Default)] -/// Factory struct used to create [AvroFormat] -pub struct AvroFormatFactory; - -impl AvroFormatFactory { - /// Creates an instance of [AvroFormatFactory] - pub fn new() -> Self { - Self {} - } -} - -impl FileFormatFactory for AvroFormatFactory { - fn create( - &self, - _state: &dyn Session, - _format_options: &HashMap, - ) -> Result> { - Ok(Arc::new(AvroFormat)) - } +//! Re-exports the [`datafusion_datasource_avro::file_format`] module, and contains tests for it. - fn default(&self) -> Arc { - Arc::new(AvroFormat) - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -impl fmt::Debug for AvroFormatFactory { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("AvroFormatFactory").finish() - } -} - -impl GetExt for AvroFormatFactory { - fn get_ext(&self) -> String { - // Removes the dot, i.e. ".parquet" -> "parquet" - DEFAULT_AVRO_EXTENSION[1..].to_string() - } -} - -/// Avro `FileFormat` implementation. -#[derive(Default, Debug)] -pub struct AvroFormat; - -#[async_trait] -impl FileFormat for AvroFormat { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_ext(&self) -> String { - AvroFormatFactory::new().get_ext() - } - - fn get_ext_with_compression( - &self, - file_compression_type: &FileCompressionType, - ) -> Result { - let ext = self.get_ext(); - match file_compression_type.get_variant() { - CompressionTypeVariant::UNCOMPRESSED => Ok(ext), - _ => internal_err!("Avro FileFormat does not support compression."), - } - } - - async fn infer_schema( - &self, - _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = vec![]; - for object in objects { - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - GetResultPayload::File(mut file, _) => { - read_avro_schema_from_reader(&mut file)? - } - GetResultPayload::Stream(_) => { - // TODO: Fetching entire file to get schema is potentially wasteful - let data = r.bytes().await?; - read_avro_schema_from_reader(&mut data.as_ref())? - } - }; - schemas.push(schema); - } - let merged_schema = Schema::try_merge(schemas)?; - Ok(Arc::new(merged_schema)) - } - - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, - table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) - } - - async fn create_physical_plan( - &self, - _state: &dyn Session, - conf: FileScanConfig, - _filters: Option<&Arc>, - ) -> Result> { - Ok(conf.with_source(self.file_source()).build()) - } - - fn file_source(&self) -> Arc { - Arc::new(AvroSource::new()) - } -} +pub use datafusion_datasource_avro::file_format::*; #[cfg(test)] -#[cfg(feature = "avro")] mod tests { - use super::*; - use crate::datasource::file_format::test_util::scan_format; - use crate::physical_plan::collect; - use crate::prelude::{SessionConfig, SessionContext}; + use std::sync::Arc; + + use crate::{ + datasource::file_format::test_util::scan_format, prelude::SessionContext, + }; use arrow::array::{as_string_array, Array}; - use datafusion_common::cast::{ - as_binary_array, as_boolean_array, as_float32_array, as_float64_array, - as_int32_array, as_timestamp_microsecond_array, + use datafusion_catalog::Session; + use datafusion_common::{ + assert_batches_eq, + cast::{ + as_binary_array, as_boolean_array, as_float32_array, as_float64_array, + as_int32_array, as_timestamp_microsecond_array, + }, + test_util, Result, }; + + use datafusion_datasource_avro::AvroFormat; + use datafusion_execution::config::SessionConfig; + use datafusion_physical_plan::{collect, ExecutionPlan}; use futures::StreamExt; #[tokio::test] @@ -260,7 +128,7 @@ mod tests { "| 1 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30312f30312f3039 | 31 | 2009-01-01T00:01:00 |", "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+"]; - crate::assert_batches_eq!(expected, &batches); + assert_batches_eq!(expected, &batches); Ok(()) } @@ -510,36 +378,9 @@ mod tests { projection: Option>, limit: Option, ) -> Result> { - let testdata = crate::test_util::arrow_test_data(); + let testdata = test_util::arrow_test_data(); let store_root = format!("{testdata}/avro"); let format = AvroFormat {}; scan_format(state, &format, &store_root, file_name, projection, limit).await } } - -#[cfg(test)] -#[cfg(not(feature = "avro"))] -mod tests { - use super::*; - - use super::super::test_util::scan_format; - use crate::error::DataFusionError; - use crate::prelude::SessionContext; - - #[tokio::test] - async fn test() -> Result<()> { - let session_ctx = SessionContext::new(); - let state = session_ctx.state(); - let format = AvroFormat {}; - let testdata = crate::test_util::arrow_test_data(); - let filename = "avro/alltypes_plain.avro"; - let result = scan_format(&state, &format, &testdata, filename, None, None).await; - assert!(matches!( - result, - Err(DataFusionError::NotImplemented(msg)) - if msg == *"cannot read avro schema without the 'avro' feature enabled" - )); - - Ok(()) - } -} diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index 45ad3e8c1c30..ca22a85d35c1 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -15,764 +15,182 @@ // specific language governing permissions and limitations // under the License. -//! [`CsvFormat`], Comma Separated Value (CSV) [`FileFormat`] abstractions - -use std::any::Any; -use std::collections::{HashMap, HashSet}; -use std::fmt::{self, Debug}; -use std::sync::Arc; - -use super::write::orchestration::spawn_writer_tasks_and_join; -use super::{ - Decoder, DecoderDeserializer, FileFormat, FileFormatFactory, - DEFAULT_SCHEMA_INFER_MAX_RECORD, -}; -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::write::demux::DemuxedStreamReceiver; -use crate::datasource::file_format::write::BatchSerializer; -use crate::datasource::physical_plan::{CsvSource, FileSink, FileSinkConfig}; -use crate::error::Result; -use crate::execution::context::SessionState; -use crate::physical_plan::insert::{DataSink, DataSinkExec}; -use crate::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, Statistics, -}; - -use arrow::array::RecordBatch; -use arrow::csv::WriterBuilder; -use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; -use arrow::error::ArrowError; -use datafusion_catalog::Session; -use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions}; -use datafusion_common::file_options::csv_writer::CsvWriterOptions; -use datafusion_common::{ - exec_err, not_impl_err, DataFusionError, GetExt, DEFAULT_CSV_EXTENSION, -}; -use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::display::FileGroupDisplay; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_execution::TaskContext; -use datafusion_expr::dml::InsertOp; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr_common::sort_expr::LexRequirement; - -use async_trait::async_trait; -use bytes::{Buf, Bytes}; -use futures::stream::BoxStream; -use futures::{pin_mut, Stream, StreamExt, TryStreamExt}; -use object_store::{delimited::newline_delimited_stream, ObjectMeta, ObjectStore}; -use regex::Regex; - -#[derive(Default)] -/// Factory struct used to create [CsvFormatFactory] -pub struct CsvFormatFactory { - /// the options for csv file read - pub options: Option, -} - -impl CsvFormatFactory { - /// Creates an instance of [CsvFormatFactory] - pub fn new() -> Self { - Self { options: None } - } - - /// Creates an instance of [CsvFormatFactory] with customized default options - pub fn new_with_options(options: CsvOptions) -> Self { - Self { - options: Some(options), - } - } -} - -impl Debug for CsvFormatFactory { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CsvFormatFactory") - .field("options", &self.options) - .finish() - } -} - -impl FileFormatFactory for CsvFormatFactory { - fn create( - &self, - state: &dyn Session, - format_options: &HashMap, - ) -> Result> { - let state = state.as_any().downcast_ref::().unwrap(); - let csv_options = match &self.options { - None => { - let mut table_options = state.default_table_options(); - table_options.set_config_format(ConfigFileType::CSV); - table_options.alter_with_string_hash_map(format_options)?; - table_options.csv - } - Some(csv_options) => { - let mut csv_options = csv_options.clone(); - for (k, v) in format_options { - csv_options.set(k, v)?; - } - csv_options - } - }; - - Ok(Arc::new(CsvFormat::default().with_options(csv_options))) - } - - fn default(&self) -> Arc { - Arc::new(CsvFormat::default()) - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -impl GetExt for CsvFormatFactory { - fn get_ext(&self) -> String { - // Removes the dot, i.e. ".parquet" -> "parquet" - DEFAULT_CSV_EXTENSION[1..].to_string() - } -} - -/// Character Separated Value `FileFormat` implementation. -#[derive(Debug, Default)] -pub struct CsvFormat { - options: CsvOptions, -} - -impl CsvFormat { - /// Return a newline delimited stream from the specified file on - /// Stream, decompressing if necessary - /// Each returned `Bytes` has a whole number of newline delimited rows - async fn read_to_delimited_chunks<'a>( - &self, - store: &Arc, - object: &ObjectMeta, - ) -> BoxStream<'a, Result> { - // stream to only read as many rows as needed into memory - let stream = store - .get(&object.location) - .await - .map_err(DataFusionError::ObjectStore); - let stream = match stream { - Ok(stream) => self - .read_to_delimited_chunks_from_stream( - stream - .into_stream() - .map_err(DataFusionError::ObjectStore) - .boxed(), - ) - .await - .map_err(DataFusionError::from) - .left_stream(), - Err(e) => { - futures::stream::once(futures::future::ready(Err(e))).right_stream() - } - }; - stream.boxed() - } - - async fn read_to_delimited_chunks_from_stream<'a>( - &self, - stream: BoxStream<'a, Result>, - ) -> BoxStream<'a, Result> { - let file_compression_type: FileCompressionType = self.options.compression.into(); - let decoder = file_compression_type.convert_stream(stream); - let steam = match decoder { - Ok(decoded_stream) => { - newline_delimited_stream(decoded_stream.map_err(|e| match e { - DataFusionError::ObjectStore(e) => e, - err => object_store::Error::Generic { - store: "read to delimited chunks failed", - source: Box::new(err), - }, - })) - .map_err(DataFusionError::from) - .left_stream() - } - Err(e) => { - futures::stream::once(futures::future::ready(Err(e))).right_stream() - } - }; - steam.boxed() - } - - /// Set the csv options - pub fn with_options(mut self, options: CsvOptions) -> Self { - self.options = options; - self - } - - /// Retrieve the csv options - pub fn options(&self) -> &CsvOptions { - &self.options - } - - /// Set a limit in terms of records to scan to infer the schema - /// - default to `DEFAULT_SCHEMA_INFER_MAX_RECORD` - pub fn with_schema_infer_max_rec(mut self, max_rec: usize) -> Self { - self.options.schema_infer_max_rec = Some(max_rec); - self - } - - /// Set true to indicate that the first line is a header. - /// - default to true - pub fn with_has_header(mut self, has_header: bool) -> Self { - self.options.has_header = Some(has_header); - self - } - - /// Set the regex to use for null values in the CSV reader. - /// - default to treat empty values as null. - pub fn with_null_regex(mut self, null_regex: Option) -> Self { - self.options.null_regex = null_regex; - self - } - - /// Returns `Some(true)` if the first line is a header, `Some(false)` if - /// it is not, and `None` if it is not specified. - pub fn has_header(&self) -> Option { - self.options.has_header - } - - /// Lines beginning with this byte are ignored. - pub fn with_comment(mut self, comment: Option) -> Self { - self.options.comment = comment; - self - } - - /// The character separating values within a row. - /// - default to ',' - pub fn with_delimiter(mut self, delimiter: u8) -> Self { - self.options.delimiter = delimiter; - self - } - - /// The quote character in a row. - /// - default to '"' - pub fn with_quote(mut self, quote: u8) -> Self { - self.options.quote = quote; - self - } - - /// The escape character in a row. - /// - default is None - pub fn with_escape(mut self, escape: Option) -> Self { - self.options.escape = escape; - self - } +//! Re-exports the [`datafusion_datasource_csv::file_format`] module, and contains tests for it. +pub use datafusion_datasource_csv::file_format::*; - /// The character used to indicate the end of a row. - /// - default to None (CRLF) - pub fn with_terminator(mut self, terminator: Option) -> Self { - self.options.terminator = terminator; - self - } - - /// Specifies whether newlines in (quoted) values are supported. - /// - /// Parsing newlines in quoted values may be affected by execution behaviour such as - /// parallel file scanning. Setting this to `true` ensures that newlines in values are - /// parsed successfully, which may reduce performance. - /// - /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting. - pub fn with_newlines_in_values(mut self, newlines_in_values: bool) -> Self { - self.options.newlines_in_values = Some(newlines_in_values); - self - } - - /// Set a `FileCompressionType` of CSV - /// - defaults to `FileCompressionType::UNCOMPRESSED` - pub fn with_file_compression_type( - mut self, - file_compression_type: FileCompressionType, - ) -> Self { - self.options.compression = file_compression_type.into(); - self - } - - /// The delimiter character. - pub fn delimiter(&self) -> u8 { - self.options.delimiter - } - - /// The quote character. - pub fn quote(&self) -> u8 { - self.options.quote - } - - /// The escape character. - pub fn escape(&self) -> Option { - self.options.escape - } -} - -#[derive(Debug)] -pub(crate) struct CsvDecoder { - inner: arrow::csv::reader::Decoder, -} - -impl CsvDecoder { - pub(crate) fn new(decoder: arrow::csv::reader::Decoder) -> Self { - Self { inner: decoder } - } -} - -impl Decoder for CsvDecoder { - fn decode(&mut self, buf: &[u8]) -> Result { - self.inner.decode(buf) - } - - fn flush(&mut self) -> Result, ArrowError> { - self.inner.flush() - } - - fn can_flush_early(&self) -> bool { - self.inner.capacity() == 0 - } -} - -impl Debug for CsvSerializer { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CsvSerializer") - .field("header", &self.header) - .finish() - } -} +#[cfg(test)] +mod tests { + use std::fmt::{self, Display}; + use std::ops::Range; + use std::sync::{Arc, Mutex}; -impl From for DecoderDeserializer { - fn from(decoder: arrow::csv::reader::Decoder) -> Self { - DecoderDeserializer::new(CsvDecoder::new(decoder)) - } -} + use super::*; -#[async_trait] -impl FileFormat for CsvFormat { - fn as_any(&self) -> &dyn Any { - self - } + use crate::datasource::file_format::test_util::scan_format; + use crate::datasource::listing::ListingOptions; + use crate::execution::session_state::SessionStateBuilder; + use crate::prelude::{CsvReadOptions, SessionConfig, SessionContext}; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use datafusion_catalog::Session; + use datafusion_common::cast::as_string_array; + use datafusion_common::internal_err; + use datafusion_common::stats::Precision; + use datafusion_common::test_util::arrow_test_data; + use datafusion_common::{assert_batches_eq, Result}; + use datafusion_datasource::decoder::{ + BatchDeserializer, DecoderDeserializer, DeserializerOutput, + }; + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource::file_format::FileFormat; + use datafusion_datasource::write::BatchSerializer; + use datafusion_expr::{col, lit}; + use datafusion_physical_plan::{collect, ExecutionPlan}; - fn get_ext(&self) -> String { - CsvFormatFactory::new().get_ext() - } + use arrow::array::{ + BooleanArray, Float64Array, Int32Array, RecordBatch, StringArray, + }; + use arrow::compute::concat_batches; + use arrow::csv::ReaderBuilder; + use arrow::util::pretty::pretty_format_batches; + use async_trait::async_trait; + use bytes::Bytes; + use chrono::DateTime; + use futures::stream::BoxStream; + use futures::StreamExt; + use object_store::local::LocalFileSystem; + use object_store::path::Path; + use object_store::{ + Attributes, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, + ObjectMeta, ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, + }; + use regex::Regex; + use rstest::*; - fn get_ext_with_compression( - &self, - file_compression_type: &FileCompressionType, - ) -> Result { - let ext = self.get_ext(); - Ok(format!("{}{}", ext, file_compression_type.get_ext())) + /// Mock ObjectStore to provide an variable stream of bytes on get + /// Able to keep track of how many iterations of the provided bytes were repeated + #[derive(Debug)] + struct VariableStream { + bytes_to_repeat: Bytes, + max_iterations: usize, + iterations_detected: Arc>, } - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = vec![]; - - let mut records_to_read = self - .options - .schema_infer_max_rec - .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); - - for object in objects { - let stream = self.read_to_delimited_chunks(store, object).await; - let (schema, records_read) = self - .infer_schema_from_stream(state, records_to_read, stream) - .await - .map_err(|err| { - DataFusionError::Context( - format!("Error when processing CSV file {}", &object.location), - Box::new(err), - ) - })?; - records_to_read -= records_read; - schemas.push(schema); - if records_to_read == 0 { - break; - } + impl Display for VariableStream { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "VariableStream") } - - let merged_schema = Schema::try_merge(schemas)?; - Ok(Arc::new(merged_schema)) - } - - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, - table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) - } - - async fn create_physical_plan( - &self, - state: &dyn Session, - mut conf: FileScanConfig, - _filters: Option<&Arc>, - ) -> Result> { - conf.file_compression_type = self.options.compression.into(); - // Consult configuration options for default values - let has_header = self - .options - .has_header - .unwrap_or(state.config_options().catalog.has_header); - let newlines_in_values = self - .options - .newlines_in_values - .unwrap_or(state.config_options().catalog.newlines_in_values); - conf.new_lines_in_values = newlines_in_values; - - let source = Arc::new( - CsvSource::new(has_header, self.options.delimiter, self.options.quote) - .with_escape(self.options.escape) - .with_terminator(self.options.terminator) - .with_comment(self.options.comment), - ); - Ok(conf.with_source(source).build()) } - async fn create_writer_physical_plan( - &self, - input: Arc, - state: &dyn Session, - conf: FileSinkConfig, - order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for CSV"); + #[async_trait] + impl ObjectStore for VariableStream { + async fn put_opts( + &self, + _location: &Path, + _payload: PutPayload, + _opts: PutOptions, + ) -> object_store::Result { + unimplemented!() } - // `has_header` and `newlines_in_values` fields of CsvOptions may inherit - // their values from session from configuration settings. To support - // this logic, writer options are built from the copy of `self.options` - // with updated values of these special fields. - let has_header = self - .options() - .has_header - .unwrap_or(state.config_options().catalog.has_header); - let newlines_in_values = self - .options() - .newlines_in_values - .unwrap_or(state.config_options().catalog.newlines_in_values); - - let options = self - .options() - .clone() - .with_has_header(has_header) - .with_newlines_in_values(newlines_in_values); - - let writer_options = CsvWriterOptions::try_from(&options)?; - - let sink = Arc::new(CsvSink::new(conf, writer_options)); - - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) - } - - fn file_source(&self) -> Arc { - Arc::new(CsvSource::default()) - } -} - -impl CsvFormat { - /// Return the inferred schema reading up to records_to_read from a - /// stream of delimited chunks returning the inferred schema and the - /// number of lines that were read - async fn infer_schema_from_stream( - &self, - state: &dyn Session, - mut records_to_read: usize, - stream: impl Stream>, - ) -> Result<(Schema, usize)> { - let mut total_records_read = 0; - let mut column_names = vec![]; - let mut column_type_possibilities = vec![]; - let mut record_number = -1; - - pin_mut!(stream); - - while let Some(chunk) = stream.next().await.transpose()? { - record_number += 1; - let first_chunk = record_number == 0; - let mut format = arrow::csv::reader::Format::default() - .with_header( - first_chunk - && self - .options - .has_header - .unwrap_or(state.config_options().catalog.has_header), - ) - .with_delimiter(self.options.delimiter) - .with_quote(self.options.quote); - - if let Some(null_regex) = &self.options.null_regex { - let regex = Regex::new(null_regex.as_str()) - .expect("Unable to parse CSV null regex."); - format = format.with_null_regex(regex); - } - - if let Some(escape) = self.options.escape { - format = format.with_escape(escape); - } - - if let Some(comment) = self.options.comment { - format = format.with_comment(comment); - } - - let (Schema { fields, .. }, records_read) = - format.infer_schema(chunk.reader(), Some(records_to_read))?; - - records_to_read -= records_read; - total_records_read += records_read; - - if first_chunk { - // set up initial structures for recording inferred schema across chunks - (column_names, column_type_possibilities) = fields - .into_iter() - .map(|field| { - let mut possibilities = HashSet::new(); - if records_read > 0 { - // at least 1 data row read, record the inferred datatype - possibilities.insert(field.data_type().clone()); - } - (field.name().clone(), possibilities) - }) - .unzip(); - } else { - if fields.len() != column_type_possibilities.len() { - return exec_err!( - "Encountered unequal lengths between records on CSV file whilst inferring schema. \ - Expected {} fields, found {} fields at record {}", - column_type_possibilities.len(), - fields.len(), - record_number + 1 - ); - } - - column_type_possibilities.iter_mut().zip(&fields).for_each( - |(possibilities, field)| { - possibilities.insert(field.data_type().clone()); - }, - ); - } - - if records_to_read == 0 { - break; - } + async fn put_multipart_opts( + &self, + _location: &Path, + _opts: PutMultipartOpts, + ) -> object_store::Result> { + unimplemented!() } - let schema = build_schema_helper(column_names, &column_type_possibilities); - Ok((schema, total_records_read)) - } -} - -fn build_schema_helper(names: Vec, types: &[HashSet]) -> Schema { - let fields = names - .into_iter() - .zip(types) - .map(|(field_name, data_type_possibilities)| { - // ripped from arrow::csv::reader::infer_reader_schema_with_csv_options - // determine data type based on possible types - // if there are incompatible types, use DataType::Utf8 - match data_type_possibilities.len() { - 1 => Field::new( - field_name, - data_type_possibilities.iter().next().unwrap().clone(), - true, - ), - 2 => { - if data_type_possibilities.contains(&DataType::Int64) - && data_type_possibilities.contains(&DataType::Float64) - { - // we have an integer and double, fall down to double - Field::new(field_name, DataType::Float64, true) - } else { - // default to Utf8 for conflicting datatypes (e.g bool and int) - Field::new(field_name, DataType::Utf8, true) - } - } - _ => Field::new(field_name, DataType::Utf8, true), - } - }) - .collect::(); - Schema::new(fields) -} + async fn get(&self, location: &Path) -> object_store::Result { + let bytes = self.bytes_to_repeat.clone(); + let range = 0..bytes.len() * self.max_iterations; + let arc = self.iterations_detected.clone(); + let stream = futures::stream::repeat_with(move || { + let arc_inner = arc.clone(); + *arc_inner.lock().unwrap() += 1; + Ok(bytes.clone()) + }) + .take(self.max_iterations) + .boxed(); + + Ok(GetResult { + payload: GetResultPayload::Stream(stream), + meta: ObjectMeta { + location: location.clone(), + last_modified: Default::default(), + size: range.end, + e_tag: None, + version: None, + }, + range: Default::default(), + attributes: Attributes::default(), + }) + } -impl Default for CsvSerializer { - fn default() -> Self { - Self::new() - } -} + async fn get_opts( + &self, + _location: &Path, + _opts: GetOptions, + ) -> object_store::Result { + unimplemented!() + } -/// Define a struct for serializing CSV records to a stream -pub struct CsvSerializer { - // CSV writer builder - builder: WriterBuilder, - // Flag to indicate whether there will be a header - header: bool, -} + async fn get_ranges( + &self, + _location: &Path, + _ranges: &[Range], + ) -> object_store::Result> { + unimplemented!() + } -impl CsvSerializer { - /// Constructor for the CsvSerializer object - pub fn new() -> Self { - Self { - builder: WriterBuilder::new(), - header: true, + async fn head(&self, _location: &Path) -> object_store::Result { + unimplemented!() } - } - /// Method for setting the CSV writer builder - pub fn with_builder(mut self, builder: WriterBuilder) -> Self { - self.builder = builder; - self - } + async fn delete(&self, _location: &Path) -> object_store::Result<()> { + unimplemented!() + } - /// Method for setting the CSV writer header status - pub fn with_header(mut self, header: bool) -> Self { - self.header = header; - self - } -} + fn list( + &self, + _prefix: Option<&Path>, + ) -> BoxStream<'_, object_store::Result> { + unimplemented!() + } -impl BatchSerializer for CsvSerializer { - fn serialize(&self, batch: RecordBatch, initial: bool) -> Result { - let mut buffer = Vec::with_capacity(4096); - let builder = self.builder.clone(); - let header = self.header && initial; - let mut writer = builder.with_header(header).build(&mut buffer); - writer.write(&batch)?; - drop(writer); - Ok(Bytes::from(buffer)) - } -} + async fn list_with_delimiter( + &self, + _prefix: Option<&Path>, + ) -> object_store::Result { + unimplemented!() + } -/// Implements [`DataSink`] for writing to a CSV file. -pub struct CsvSink { - /// Config options for writing data - config: FileSinkConfig, - writer_options: CsvWriterOptions, -} + async fn copy(&self, _from: &Path, _to: &Path) -> object_store::Result<()> { + unimplemented!() + } -impl Debug for CsvSink { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("CsvSink").finish() + async fn copy_if_not_exists( + &self, + _from: &Path, + _to: &Path, + ) -> object_store::Result<()> { + unimplemented!() + } } -} -impl DisplayAs for CsvSink { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "CsvSink(file_groups=",)?; - FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; - write!(f, ")") + impl VariableStream { + pub fn new(bytes_to_repeat: Bytes, max_iterations: usize) -> Self { + Self { + bytes_to_repeat, + max_iterations, + iterations_detected: Arc::new(Mutex::new(0)), } } - } -} -impl CsvSink { - /// Create from config. - pub fn new(config: FileSinkConfig, writer_options: CsvWriterOptions) -> Self { - Self { - config, - writer_options, + pub fn get_iterations_detected(&self) -> usize { + *self.iterations_detected.lock().unwrap() } } - /// Retrieve the writer options - pub fn writer_options(&self) -> &CsvWriterOptions { - &self.writer_options - } -} - -#[async_trait] -impl FileSink for CsvSink { - fn config(&self) -> &FileSinkConfig { - &self.config - } - - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, - demux_task: SpawnedTask>, - file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, - ) -> Result { - let builder = self.writer_options.writer_options.clone(); - let header = builder.header(); - let serializer = Arc::new( - CsvSerializer::new() - .with_builder(builder) - .with_header(header), - ) as _; - spawn_writer_tasks_and_join( - context, - serializer, - self.writer_options.compression.into(), - object_store, - demux_task, - file_stream_rx, - ) - .await - } -} - -#[async_trait] -impl DataSink for CsvSink { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> &SchemaRef { - self.config.output_schema() - } - - async fn write_all( - &self, - data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await - } -} - -#[cfg(test)] -mod tests { - use super::super::test_util::scan_format; - use super::*; - use crate::assert_batches_eq; - use crate::datasource::file_format::file_compression_type::FileCompressionType; - use crate::datasource::file_format::test_util::VariableStream; - use crate::datasource::file_format::{ - BatchDeserializer, DecoderDeserializer, DeserializerOutput, - }; - use crate::datasource::listing::ListingOptions; - use crate::execution::session_state::SessionStateBuilder; - use crate::physical_plan::collect; - use crate::prelude::{CsvReadOptions, SessionConfig, SessionContext}; - use crate::test_util::arrow_test_data; - - use arrow::array::{BooleanArray, Float64Array, Int32Array, StringArray}; - use arrow::compute::concat_batches; - use arrow::csv::ReaderBuilder; - use arrow::util::pretty::pretty_format_batches; - use datafusion_common::cast::as_string_array; - use datafusion_common::internal_err; - use datafusion_common::stats::Precision; - use datafusion_expr::{col, lit}; - - use chrono::DateTime; - use object_store::local::LocalFileSystem; - use object_store::path::Path; - use regex::Regex; - use rstest::*; - #[tokio::test] async fn read_small_batches() -> Result<()> { let config = SessionConfig::new().with_batch_size(2); @@ -1056,6 +474,11 @@ mod tests { async fn query_compress_data( file_compression_type: FileCompressionType, ) -> Result<()> { + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::DataFusionError; + use datafusion_datasource::file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD; + use futures::TryStreamExt; + let mut cfg = SessionConfig::new(); cfg.options_mut().catalog.has_header = true; let session_state = SessionStateBuilder::new() @@ -1675,10 +1098,7 @@ mod tests { } } - fn csv_expected_batch( - schema: SchemaRef, - line_count: usize, - ) -> Result { + fn csv_expected_batch(schema: SchemaRef, line_count: usize) -> Result { let mut c1 = Vec::with_capacity(line_count); let mut c2 = Vec::with_capacity(line_count); let mut c3 = Vec::with_capacity(line_count); @@ -1721,7 +1141,7 @@ mod tests { (int_value, float_value, bool_value, char_value) } - fn csv_schema() -> Arc { + fn csv_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("c1", DataType::Int32, true), Field::new("c2", DataType::Float64, true), diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs index 1a2aaf3af8be..ef9dd08c1c26 100644 --- a/datafusion/core/src/datasource/file_format/json.rs +++ b/datafusion/core/src/datasource/file_format/json.rs @@ -15,424 +15,27 @@ // specific language governing permissions and limitations // under the License. -//! [`JsonFormat`]: Line delimited JSON [`FileFormat`] abstractions - -use std::any::Any; -use std::collections::HashMap; -use std::fmt; -use std::fmt::Debug; -use std::io::BufReader; -use std::sync::Arc; - -use super::write::orchestration::spawn_writer_tasks_and_join; -use super::{ - Decoder, DecoderDeserializer, FileFormat, FileFormatFactory, - DEFAULT_SCHEMA_INFER_MAX_RECORD, -}; -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::write::demux::DemuxedStreamReceiver; -use crate::datasource::file_format::write::BatchSerializer; -use crate::datasource::physical_plan::{FileSink, FileSinkConfig, JsonSource}; -use crate::error::Result; -use crate::execution::SessionState; -use crate::physical_plan::insert::{DataSink, DataSinkExec}; -use crate::physical_plan::{ - DisplayAs, DisplayFormatType, SendableRecordBatchStream, Statistics, -}; - -use arrow::array::RecordBatch; -use arrow::datatypes::{Schema, SchemaRef}; -use arrow::error::ArrowError; -use arrow::json; -use arrow::json::reader::{infer_json_schema_from_iterator, ValueIter}; -use datafusion_catalog::Session; -use datafusion_common::config::{ConfigField, ConfigFileType, JsonOptions}; -use datafusion_common::file_options::json_writer::JsonWriterOptions; -use datafusion_common::{not_impl_err, GetExt, DEFAULT_JSON_EXTENSION}; -use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::display::FileGroupDisplay; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_execution::TaskContext; -use datafusion_expr::dml::InsertOp; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_plan::ExecutionPlan; - -use async_trait::async_trait; -use bytes::{Buf, Bytes}; -use datafusion_physical_expr_common::sort_expr::LexRequirement; -use object_store::{GetResultPayload, ObjectMeta, ObjectStore}; - -#[derive(Default)] -/// Factory struct used to create [JsonFormat] -pub struct JsonFormatFactory { - /// the options carried by format factory - pub options: Option, -} - -impl JsonFormatFactory { - /// Creates an instance of [JsonFormatFactory] - pub fn new() -> Self { - Self { options: None } - } - - /// Creates an instance of [JsonFormatFactory] with customized default options - pub fn new_with_options(options: JsonOptions) -> Self { - Self { - options: Some(options), - } - } -} - -impl FileFormatFactory for JsonFormatFactory { - fn create( - &self, - state: &dyn Session, - format_options: &HashMap, - ) -> Result> { - let state = state.as_any().downcast_ref::().unwrap(); - let json_options = match &self.options { - None => { - let mut table_options = state.default_table_options(); - table_options.set_config_format(ConfigFileType::JSON); - table_options.alter_with_string_hash_map(format_options)?; - table_options.json - } - Some(json_options) => { - let mut json_options = json_options.clone(); - for (k, v) in format_options { - json_options.set(k, v)?; - } - json_options - } - }; - - Ok(Arc::new(JsonFormat::default().with_options(json_options))) - } - - fn default(&self) -> Arc { - Arc::new(JsonFormat::default()) - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -impl GetExt for JsonFormatFactory { - fn get_ext(&self) -> String { - // Removes the dot, i.e. ".parquet" -> "parquet" - DEFAULT_JSON_EXTENSION[1..].to_string() - } -} - -impl Debug for JsonFormatFactory { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("JsonFormatFactory") - .field("options", &self.options) - .finish() - } -} - -/// New line delimited JSON `FileFormat` implementation. -#[derive(Debug, Default)] -pub struct JsonFormat { - options: JsonOptions, -} - -impl JsonFormat { - /// Set JSON options - pub fn with_options(mut self, options: JsonOptions) -> Self { - self.options = options; - self - } - - /// Retrieve JSON options - pub fn options(&self) -> &JsonOptions { - &self.options - } - - /// Set a limit in terms of records to scan to infer the schema - /// - defaults to `DEFAULT_SCHEMA_INFER_MAX_RECORD` - pub fn with_schema_infer_max_rec(mut self, max_rec: usize) -> Self { - self.options.schema_infer_max_rec = Some(max_rec); - self - } - - /// Set a `FileCompressionType` of JSON - /// - defaults to `FileCompressionType::UNCOMPRESSED` - pub fn with_file_compression_type( - mut self, - file_compression_type: FileCompressionType, - ) -> Self { - self.options.compression = file_compression_type.into(); - self - } -} - -#[async_trait] -impl FileFormat for JsonFormat { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_ext(&self) -> String { - JsonFormatFactory::new().get_ext() - } - - fn get_ext_with_compression( - &self, - file_compression_type: &FileCompressionType, - ) -> Result { - let ext = self.get_ext(); - Ok(format!("{}{}", ext, file_compression_type.get_ext())) - } - - async fn infer_schema( - &self, - _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = Vec::new(); - let mut records_to_read = self - .options - .schema_infer_max_rec - .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); - let file_compression_type = FileCompressionType::from(self.options.compression); - for object in objects { - let mut take_while = || { - let should_take = records_to_read > 0; - if should_take { - records_to_read -= 1; - } - should_take - }; - - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - GetResultPayload::File(file, _) => { - let decoder = file_compression_type.convert_read(file)?; - let mut reader = BufReader::new(decoder); - let iter = ValueIter::new(&mut reader, None); - infer_json_schema_from_iterator(iter.take_while(|_| take_while()))? - } - GetResultPayload::Stream(_) => { - let data = r.bytes().await?; - let decoder = file_compression_type.convert_read(data.reader())?; - let mut reader = BufReader::new(decoder); - let iter = ValueIter::new(&mut reader, None); - infer_json_schema_from_iterator(iter.take_while(|_| take_while()))? - } - }; - - schemas.push(schema); - if records_to_read == 0 { - break; - } - } - - let schema = Schema::try_merge(schemas)?; - Ok(Arc::new(schema)) - } - - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, - table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) - } - - async fn create_physical_plan( - &self, - _state: &dyn Session, - mut conf: FileScanConfig, - _filters: Option<&Arc>, - ) -> Result> { - let source = Arc::new(JsonSource::new()); - conf.file_compression_type = FileCompressionType::from(self.options.compression); - Ok(conf.with_source(source).build()) - } - - async fn create_writer_physical_plan( - &self, - input: Arc, - _state: &dyn Session, - conf: FileSinkConfig, - order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Json"); - } - - let writer_options = JsonWriterOptions::try_from(&self.options)?; - - let sink = Arc::new(JsonSink::new(conf, writer_options)); - - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) - } - - fn file_source(&self) -> Arc { - Arc::new(JsonSource::default()) - } -} - -impl Default for JsonSerializer { - fn default() -> Self { - Self::new() - } -} - -/// Define a struct for serializing Json records to a stream -pub struct JsonSerializer {} - -impl JsonSerializer { - /// Constructor for the JsonSerializer object - pub fn new() -> Self { - Self {} - } -} - -impl BatchSerializer for JsonSerializer { - fn serialize(&self, batch: RecordBatch, _initial: bool) -> Result { - let mut buffer = Vec::with_capacity(4096); - let mut writer = json::LineDelimitedWriter::new(&mut buffer); - writer.write(&batch)?; - Ok(Bytes::from(buffer)) - } -} - -/// Implements [`DataSink`] for writing to a Json file. -pub struct JsonSink { - /// Config options for writing data - config: FileSinkConfig, - /// Writer options for underlying Json writer - writer_options: JsonWriterOptions, -} - -impl Debug for JsonSink { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("JsonSink").finish() - } -} - -impl DisplayAs for JsonSink { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "JsonSink(file_groups=",)?; - FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; - write!(f, ")") - } - } - } -} - -impl JsonSink { - /// Create from config. - pub fn new(config: FileSinkConfig, writer_options: JsonWriterOptions) -> Self { - Self { - config, - writer_options, - } - } - - /// Retrieve the writer options - pub fn writer_options(&self) -> &JsonWriterOptions { - &self.writer_options - } -} - -#[async_trait] -impl FileSink for JsonSink { - fn config(&self) -> &FileSinkConfig { - &self.config - } - - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, - demux_task: SpawnedTask>, - file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, - ) -> Result { - let serializer = Arc::new(JsonSerializer::new()) as _; - spawn_writer_tasks_and_join( - context, - serializer, - self.writer_options.compression.into(), - object_store, - demux_task, - file_stream_rx, - ) - .await - } -} - -#[async_trait] -impl DataSink for JsonSink { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> &SchemaRef { - self.config.output_schema() - } - - async fn write_all( - &self, - data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await - } -} - -#[derive(Debug)] -pub(crate) struct JsonDecoder { - inner: json::reader::Decoder, -} - -impl JsonDecoder { - pub(crate) fn new(decoder: json::reader::Decoder) -> Self { - Self { inner: decoder } - } -} - -impl Decoder for JsonDecoder { - fn decode(&mut self, buf: &[u8]) -> Result { - self.inner.decode(buf) - } - - fn flush(&mut self) -> Result, ArrowError> { - self.inner.flush() - } - - fn can_flush_early(&self) -> bool { - false - } -} - -impl From for DecoderDeserializer { - fn from(decoder: json::reader::Decoder) -> Self { - DecoderDeserializer::new(JsonDecoder::new(decoder)) - } -} +//! Re-exports the [`datafusion_datasource_json::file_format`] module, and contains tests for it. +pub use datafusion_datasource_json::file_format::*; #[cfg(test)] mod tests { - use super::super::test_util::scan_format; + use std::sync::Arc; + use super::*; - use crate::datasource::file_format::{ + + use crate::datasource::file_format::test_util::scan_format; + use crate::prelude::{NdJsonReadOptions, SessionConfig, SessionContext}; + use crate::test::object_store::local_unpartitioned_file; + use arrow::array::RecordBatch; + use arrow_schema::Schema; + use bytes::Bytes; + use datafusion_catalog::Session; + use datafusion_datasource::decoder::{ BatchDeserializer, DecoderDeserializer, DeserializerOutput, }; - use crate::execution::options::NdJsonReadOptions; - use crate::physical_plan::collect; - use crate::prelude::{SessionConfig, SessionContext}; - use crate::test::object_store::local_unpartitioned_file; + use datafusion_datasource::file_format::FileFormat; + use datafusion_physical_plan::{collect, ExecutionPlan}; use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field}; @@ -442,6 +45,7 @@ mod tests { use datafusion_common::stats::Precision; use datafusion_common::{assert_batches_eq, internal_err}; + use datafusion_common::Result; use futures::StreamExt; use object_store::local::LocalFileSystem; use regex::Regex; diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs index 5dbf4957a4b5..df74e5d060e6 100644 --- a/datafusion/core/src/datasource/file_format/mod.rs +++ b/datafusion/core/src/datasource/file_format/mod.rs @@ -19,211 +19,33 @@ //! See write.rs for write related helper methods pub mod arrow; -pub mod avro; pub mod csv; pub mod json; -pub mod options; + +#[cfg(feature = "avro")] +pub mod avro; + #[cfg(feature = "parquet")] pub mod parquet; -use ::arrow::array::RecordBatch; -use arrow_schema::ArrowError; -use bytes::Buf; -use bytes::Bytes; -use datafusion_common::Result; +pub mod options; + pub use datafusion_datasource::file_compression_type; pub use datafusion_datasource::file_format::*; pub use datafusion_datasource::write; -use futures::stream::BoxStream; -use futures::StreamExt as _; -use futures::{ready, Stream}; -use std::collections::VecDeque; -use std::fmt; -use std::task::Poll; - -/// Possible outputs of a [`BatchDeserializer`]. -#[derive(Debug, PartialEq)] -pub enum DeserializerOutput { - /// A successfully deserialized [`RecordBatch`]. - RecordBatch(RecordBatch), - /// The deserializer requires more data to make progress. - RequiresMoreData, - /// The input data has been exhausted. - InputExhausted, -} - -/// Trait defining a scheme for deserializing byte streams into structured data. -/// Implementors of this trait are responsible for converting raw bytes into -/// `RecordBatch` objects. -pub trait BatchDeserializer: Send + fmt::Debug { - /// Feeds a message for deserialization, updating the internal state of - /// this `BatchDeserializer`. Note that one can call this function multiple - /// times before calling `next`, which will queue multiple messages for - /// deserialization. Returns the number of bytes consumed. - fn digest(&mut self, message: T) -> usize; - - /// Attempts to deserialize any pending messages and returns a - /// `DeserializerOutput` to indicate progress. - fn next(&mut self) -> Result; - - /// Informs the deserializer that no more messages will be provided for - /// deserialization. - fn finish(&mut self); -} - -/// A general interface for decoders such as [`arrow::json::reader::Decoder`] and -/// [`arrow::csv::reader::Decoder`]. Defines an interface similar to -/// [`Decoder::decode`] and [`Decoder::flush`] methods, but also includes -/// a method to check if the decoder can flush early. Intended to be used in -/// conjunction with [`DecoderDeserializer`]. -/// -/// [`arrow::json::reader::Decoder`]: ::arrow::json::reader::Decoder -/// [`arrow::csv::reader::Decoder`]: ::arrow::csv::reader::Decoder -/// [`Decoder::decode`]: ::arrow::json::reader::Decoder::decode -/// [`Decoder::flush`]: ::arrow::json::reader::Decoder::flush -pub(crate) trait Decoder: Send + fmt::Debug { - /// See [`arrow::json::reader::Decoder::decode`]. - /// - /// [`arrow::json::reader::Decoder::decode`]: ::arrow::json::reader::Decoder::decode - fn decode(&mut self, buf: &[u8]) -> Result; - - /// See [`arrow::json::reader::Decoder::flush`]. - /// - /// [`arrow::json::reader::Decoder::flush`]: ::arrow::json::reader::Decoder::flush - fn flush(&mut self) -> Result, ArrowError>; - - /// Whether the decoder can flush early in its current state. - fn can_flush_early(&self) -> bool; -} - -impl fmt::Debug for DecoderDeserializer { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Deserializer") - .field("buffered_queue", &self.buffered_queue) - .field("finalized", &self.finalized) - .finish() - } -} - -impl BatchDeserializer for DecoderDeserializer { - fn digest(&mut self, message: Bytes) -> usize { - if message.is_empty() { - return 0; - } - - let consumed = message.len(); - self.buffered_queue.push_back(message); - consumed - } - - fn next(&mut self) -> Result { - while let Some(buffered) = self.buffered_queue.front_mut() { - let decoded = self.decoder.decode(buffered)?; - buffered.advance(decoded); - - if buffered.is_empty() { - self.buffered_queue.pop_front(); - } - - // Flush when the stream ends or batch size is reached - // Certain implementations can flush early - if decoded == 0 || self.decoder.can_flush_early() { - return match self.decoder.flush() { - Ok(Some(batch)) => Ok(DeserializerOutput::RecordBatch(batch)), - Ok(None) => continue, - Err(e) => Err(e), - }; - } - } - if self.finalized { - Ok(DeserializerOutput::InputExhausted) - } else { - Ok(DeserializerOutput::RequiresMoreData) - } - } - - fn finish(&mut self) { - self.finalized = true; - // Ensure the decoder is flushed: - self.buffered_queue.push_back(Bytes::new()); - } -} - -/// A generic, decoder-based deserialization scheme for processing encoded data. -/// -/// This struct is responsible for converting a stream of bytes, which represent -/// encoded data, into a stream of `RecordBatch` objects, following the specified -/// schema and formatting options. It also handles any buffering necessary to satisfy -/// the `Decoder` interface. -pub(crate) struct DecoderDeserializer { - /// The underlying decoder used for deserialization - pub(crate) decoder: T, - /// The buffer used to store the remaining bytes to be decoded - pub(crate) buffered_queue: VecDeque, - /// Whether the input stream has been fully consumed - pub(crate) finalized: bool, -} - -impl DecoderDeserializer { - /// Creates a new `DecoderDeserializer` with the provided decoder. - pub fn new(decoder: T) -> Self { - DecoderDeserializer { - decoder, - buffered_queue: VecDeque::new(), - finalized: false, - } - } -} - -/// Deserializes a stream of bytes into a stream of [`RecordBatch`] objects using the -/// provided deserializer. -/// -/// Returns a boxed stream of `Result`. The stream yields [`RecordBatch`] -/// objects as they are produced by the deserializer, or an [`ArrowError`] if an error -/// occurs while polling the input or deserializing. -pub(crate) fn deserialize_stream<'a>( - mut input: impl Stream> + Unpin + Send + 'a, - mut deserializer: impl BatchDeserializer + 'a, -) -> BoxStream<'a, Result> { - futures::stream::poll_fn(move |cx| loop { - match ready!(input.poll_next_unpin(cx)).transpose()? { - Some(b) => _ = deserializer.digest(b), - None => deserializer.finish(), - }; - - return match deserializer.next()? { - DeserializerOutput::RecordBatch(rb) => Poll::Ready(Some(Ok(rb))), - DeserializerOutput::InputExhausted => Poll::Ready(None), - DeserializerOutput::RequiresMoreData => continue, - }; - }) - .boxed() -} #[cfg(test)] pub(crate) mod test_util { - use std::fmt::{self, Display}; - use std::ops::Range; - use std::sync::{Arc, Mutex}; + use std::sync::Arc; - use crate::datasource::listing::PartitionedFile; - use crate::datasource::object_store::ObjectStoreUrl; - use crate::test::object_store::local_unpartitioned_file; - use async_trait::async_trait; - use bytes::Bytes; use datafusion_catalog::Session; use datafusion_common::Result; - use datafusion_datasource::file_format::FileFormat; - use datafusion_datasource::file_scan_config::FileScanConfig; - use datafusion_physical_plan::ExecutionPlan; - use futures::stream::BoxStream; - use futures::StreamExt; - use object_store::local::LocalFileSystem; - use object_store::path::Path; - use object_store::{ - Attributes, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, - ObjectMeta, ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, + use datafusion_datasource::{ + file_format::FileFormat, file_scan_config::FileScanConfig, PartitionedFile, }; + use datafusion_execution::object_store::ObjectStoreUrl; + + use crate::test::object_store::local_unpartitioned_file; pub async fn scan_format( state: &dyn Session, @@ -232,8 +54,8 @@ pub(crate) mod test_util { file_name: &str, projection: Option>, limit: Option, - ) -> Result> { - let store = Arc::new(LocalFileSystem::new()) as _; + ) -> Result> { + let store = Arc::new(object_store::local::LocalFileSystem::new()) as _; let meta = local_unpartitioned_file(format!("{store_root}/{file_name}")); let file_schema = format @@ -270,129 +92,41 @@ pub(crate) mod test_util { .await?; Ok(exec) } +} - /// Mock ObjectStore to provide an variable stream of bytes on get - /// Able to keep track of how many iterations of the provided bytes were repeated - #[derive(Debug)] - pub struct VariableStream { - bytes_to_repeat: Bytes, - max_iterations: usize, - iterations_detected: Arc>, - } - - impl Display for VariableStream { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "VariableStream") - } - } - - #[async_trait] - impl ObjectStore for VariableStream { - async fn put_opts( - &self, - _location: &Path, - _payload: PutPayload, - _opts: PutOptions, - ) -> object_store::Result { - unimplemented!() - } - - async fn put_multipart_opts( - &self, - _location: &Path, - _opts: PutMultipartOpts, - ) -> object_store::Result> { - unimplemented!() - } - - async fn get(&self, location: &Path) -> object_store::Result { - let bytes = self.bytes_to_repeat.clone(); - let range = 0..bytes.len() * self.max_iterations; - let arc = self.iterations_detected.clone(); - let stream = futures::stream::repeat_with(move || { - let arc_inner = arc.clone(); - *arc_inner.lock().unwrap() += 1; - Ok(bytes.clone()) - }) - .take(self.max_iterations) - .boxed(); - - Ok(GetResult { - payload: GetResultPayload::Stream(stream), - meta: ObjectMeta { - location: location.clone(), - last_modified: Default::default(), - size: range.end, - e_tag: None, - version: None, - }, - range: Default::default(), - attributes: Attributes::default(), - }) - } - - async fn get_opts( - &self, - _location: &Path, - _opts: GetOptions, - ) -> object_store::Result { - unimplemented!() - } - - async fn get_ranges( - &self, - _location: &Path, - _ranges: &[Range], - ) -> object_store::Result> { - unimplemented!() - } - - async fn head(&self, _location: &Path) -> object_store::Result { - unimplemented!() - } - - async fn delete(&self, _location: &Path) -> object_store::Result<()> { - unimplemented!() - } - - fn list( - &self, - _prefix: Option<&Path>, - ) -> BoxStream<'_, object_store::Result> { - unimplemented!() - } - - async fn list_with_delimiter( - &self, - _prefix: Option<&Path>, - ) -> object_store::Result { - unimplemented!() - } - - async fn copy(&self, _from: &Path, _to: &Path) -> object_store::Result<()> { - unimplemented!() - } - - async fn copy_if_not_exists( - &self, - _from: &Path, - _to: &Path, - ) -> object_store::Result<()> { - unimplemented!() - } - } - - impl VariableStream { - pub fn new(bytes_to_repeat: Bytes, max_iterations: usize) -> Self { - Self { - bytes_to_repeat, - max_iterations, - iterations_detected: Arc::new(Mutex::new(0)), - } - } +#[cfg(test)] +mod tests { + #[cfg(feature = "parquet")] + #[tokio::test] + async fn write_parquet_results_error_handling() -> datafusion_common::Result<()> { + use std::sync::Arc; + + use object_store::local::LocalFileSystem; + use tempfile::TempDir; + use url::Url; + + use crate::{ + dataframe::DataFrameWriteOptions, + prelude::{CsvReadOptions, SessionContext}, + }; - pub fn get_iterations_detected(&self) -> usize { - *self.iterations_detected.lock().unwrap() - } + let ctx = SessionContext::new(); + // register a local file system object store for /tmp directory + let tmp_dir = TempDir::new()?; + let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); + let local_url = Url::parse("file://local").unwrap(); + ctx.register_object_store(&local_url, local); + + let options = CsvReadOptions::default() + .schema_infer_max_records(2) + .has_header(true); + let df = ctx.read_csv("tests/data/corrupt.csv", options).await?; + let out_dir_url = "file://local/out"; + let e = df + .write_parquet(out_dir_url, DataFrameWriteOptions::new(), None) + .await + .expect_err("should fail because input file does not match inferred schema"); + assert_eq!(e.strip_backtrace(), "Arrow error: Parser error: Error while parsing value d for column 0 at line 4"); + Ok(()) } } diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index 2f32479ed2b0..08e9a628dd61 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -19,16 +19,17 @@ use std::sync::Arc; -use crate::datasource::file_format::arrow::ArrowFormat; -use crate::datasource::file_format::file_compression_type::FileCompressionType; +#[cfg(feature = "avro")] +use crate::datasource::file_format::avro::AvroFormat; + #[cfg(feature = "parquet")] use crate::datasource::file_format::parquet::ParquetFormat; + +use crate::datasource::file_format::arrow::ArrowFormat; +use crate::datasource::file_format::file_compression_type::FileCompressionType; use crate::datasource::file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD; use crate::datasource::listing::ListingTableUrl; -use crate::datasource::{ - file_format::{avro::AvroFormat, csv::CsvFormat, json::JsonFormat}, - listing::ListingOptions, -}; +use crate::datasource::{file_format::csv::CsvFormat, listing::ListingOptions}; use crate::error::Result; use crate::execution::context::{SessionConfig, SessionState}; @@ -40,6 +41,7 @@ use datafusion_common::{ }; use async_trait::async_trait; +use datafusion_datasource_json::file_format::JsonFormat; use datafusion_expr::SortExpr; /// Options that control the reading of CSV files. @@ -629,6 +631,7 @@ impl ReadOptions<'_> for NdJsonReadOptions<'_> { } } +#[cfg(feature = "avro")] #[async_trait] impl ReadOptions<'_> for AvroReadOptions<'_> { fn to_listing_options( diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index 4a24871aeef7..e2c7d1ecafa3 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -15,1360 +15,17 @@ // specific language governing permissions and limitations // under the License. -//! [`ParquetFormat`]: Parquet [`FileFormat`] abstractions - -use std::any::Any; -use std::fmt; -use std::fmt::Debug; -use std::ops::Range; -use std::sync::Arc; - -use super::write::demux::DemuxedStreamReceiver; -use super::write::{create_writer, SharedBuffer}; -use super::{FileFormat, FileFormatFactory, FilePushdownSupport}; -use crate::arrow::array::RecordBatch; -use crate::arrow::datatypes::{Fields, Schema, SchemaRef}; -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::write::get_writer_schema; -use crate::datasource::physical_plan::parquet::can_expr_be_pushed_down_with_schemas; -use crate::datasource::physical_plan::parquet::source::ParquetSource; -use crate::datasource::physical_plan::{FileSink, FileSinkConfig}; -use crate::datasource::statistics::{create_max_min_accs, get_col_stats}; -use crate::error::Result; -use crate::execution::SessionState; -use crate::physical_plan::insert::{DataSink, DataSinkExec}; -use crate::physical_plan::{ - Accumulator, DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, - Statistics, -}; - -use arrow::compute::sum; -use arrow_schema::{DataType, Field, FieldRef}; -use datafusion_catalog::Session; -use datafusion_common::config::{ConfigField, ConfigFileType, TableParquetOptions}; -use datafusion_common::parsers::CompressionTypeVariant; -use datafusion_common::stats::Precision; -use datafusion_common::HashMap; -use datafusion_common::{ - internal_datafusion_err, internal_err, not_impl_err, DataFusionError, GetExt, - DEFAULT_PARQUET_EXTENSION, -}; -use datafusion_common_runtime::SpawnedTask; -use datafusion_datasource::display::FileGroupDisplay; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; -use datafusion_execution::TaskContext; -use datafusion_expr::dml::InsertOp; -use datafusion_expr::Expr; -use datafusion_functions_aggregate::min_max::{MaxAccumulator, MinAccumulator}; -use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr_common::sort_expr::LexRequirement; - -use async_trait::async_trait; -use bytes::Bytes; -use futures::future::BoxFuture; -use futures::{FutureExt, StreamExt, TryStreamExt}; -use log::debug; -use object_store::buffered::BufWriter; -use object_store::path::Path; -use object_store::{ObjectMeta, ObjectStore}; -use parquet::arrow::arrow_reader::statistics::StatisticsConverter; -use parquet::arrow::arrow_writer::{ - compute_leaves, get_column_writers, ArrowColumnChunk, ArrowColumnWriter, - ArrowLeafColumn, ArrowWriterOptions, -}; -use parquet::arrow::async_reader::MetadataFetch; -use parquet::arrow::{parquet_to_arrow_schema, ArrowSchemaConverter, AsyncArrowWriter}; -use parquet::errors::ParquetError; -use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData}; -use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; -use parquet::file::writer::SerializedFileWriter; -use parquet::format::FileMetaData; -use tokio::io::{AsyncWrite, AsyncWriteExt}; -use tokio::sync::mpsc::{self, Receiver, Sender}; -use tokio::task::JoinSet; - -/// Initial writing buffer size. Note this is just a size hint for efficiency. It -/// will grow beyond the set value if needed. -const INITIAL_BUFFER_BYTES: usize = 1048576; - -/// When writing parquet files in parallel, if the buffered Parquet data exceeds -/// this size, it is flushed to object store -const BUFFER_FLUSH_BYTES: usize = 1024000; - -#[derive(Default)] -/// Factory struct used to create [ParquetFormat] -pub struct ParquetFormatFactory { - /// inner options for parquet - pub options: Option, -} - -impl ParquetFormatFactory { - /// Creates an instance of [ParquetFormatFactory] - pub fn new() -> Self { - Self { options: None } - } - - /// Creates an instance of [ParquetFormatFactory] with customized default options - pub fn new_with_options(options: TableParquetOptions) -> Self { - Self { - options: Some(options), - } - } -} - -impl FileFormatFactory for ParquetFormatFactory { - fn create( - &self, - state: &dyn Session, - format_options: &std::collections::HashMap, - ) -> Result> { - let state = state.as_any().downcast_ref::().unwrap(); - let parquet_options = match &self.options { - None => { - let mut table_options = state.default_table_options(); - table_options.set_config_format(ConfigFileType::PARQUET); - table_options.alter_with_string_hash_map(format_options)?; - table_options.parquet - } - Some(parquet_options) => { - let mut parquet_options = parquet_options.clone(); - for (k, v) in format_options { - parquet_options.set(k, v)?; - } - parquet_options - } - }; - - Ok(Arc::new( - ParquetFormat::default().with_options(parquet_options), - )) - } - - fn default(&self) -> Arc { - Arc::new(ParquetFormat::default()) - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -impl GetExt for ParquetFormatFactory { - fn get_ext(&self) -> String { - // Removes the dot, i.e. ".parquet" -> "parquet" - DEFAULT_PARQUET_EXTENSION[1..].to_string() - } -} - -impl Debug for ParquetFormatFactory { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ParquetFormatFactory") - .field("ParquetFormatFactory", &self.options) - .finish() - } -} -/// The Apache Parquet `FileFormat` implementation -#[derive(Debug, Default)] -pub struct ParquetFormat { - options: TableParquetOptions, -} - -impl ParquetFormat { - /// Construct a new Format with no local overrides - pub fn new() -> Self { - Self::default() - } - - /// Activate statistics based row group level pruning - /// - If `None`, defaults to value on `config_options` - pub fn with_enable_pruning(mut self, enable: bool) -> Self { - self.options.global.pruning = enable; - self - } - - /// Return `true` if pruning is enabled - pub fn enable_pruning(&self) -> bool { - self.options.global.pruning - } - - /// Provide a hint to the size of the file metadata. If a hint is provided - /// the reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. - /// Without a hint, two read are required. One read to fetch the 8-byte parquet footer and then - /// another read to fetch the metadata length encoded in the footer. - /// - /// - If `None`, defaults to value on `config_options` - pub fn with_metadata_size_hint(mut self, size_hint: Option) -> Self { - self.options.global.metadata_size_hint = size_hint; - self - } - - /// Return the metadata size hint if set - pub fn metadata_size_hint(&self) -> Option { - self.options.global.metadata_size_hint - } - - /// Tell the parquet reader to skip any metadata that may be in - /// the file Schema. This can help avoid schema conflicts due to - /// metadata. - /// - /// - If `None`, defaults to value on `config_options` - pub fn with_skip_metadata(mut self, skip_metadata: bool) -> Self { - self.options.global.skip_metadata = skip_metadata; - self - } - - /// Returns `true` if schema metadata will be cleared prior to - /// schema merging. - pub fn skip_metadata(&self) -> bool { - self.options.global.skip_metadata - } - - /// Set Parquet options for the ParquetFormat - pub fn with_options(mut self, options: TableParquetOptions) -> Self { - self.options = options; - self - } - - /// Parquet options - pub fn options(&self) -> &TableParquetOptions { - &self.options - } - - /// Return `true` if should use view types. - /// - /// If this returns true, DataFusion will instruct the parquet reader - /// to read string / binary columns using view `StringView` or `BinaryView` - /// if the table schema specifies those types, regardless of any embedded metadata - /// that may specify an alternate Arrow type. The parquet reader is optimized - /// for reading `StringView` and `BinaryView` and such queries are significantly faster. - /// - /// If this returns false, the parquet reader will read the columns according to the - /// defaults or any embedded Arrow type information. This may result in reading - /// `StringArrays` and then casting to `StringViewArray` which is less efficient. - pub fn force_view_types(&self) -> bool { - self.options.global.schema_force_view_types - } - - /// If true, will use view types. See [`Self::force_view_types`] for details - pub fn with_force_view_types(mut self, use_views: bool) -> Self { - self.options.global.schema_force_view_types = use_views; - self - } - - /// Return `true` if binary types will be read as strings. - /// - /// If this returns true, DataFusion will instruct the parquet reader - /// to read binary columns such as `Binary` or `BinaryView` as the - /// corresponding string type such as `Utf8` or `LargeUtf8`. - /// The parquet reader has special optimizations for `Utf8` and `LargeUtf8` - /// validation, and such queries are significantly faster than reading - /// binary columns and then casting to string columns. - pub fn binary_as_string(&self) -> bool { - self.options.global.binary_as_string - } - - /// If true, will read binary types as strings. See [`Self::binary_as_string`] for details - pub fn with_binary_as_string(mut self, binary_as_string: bool) -> Self { - self.options.global.binary_as_string = binary_as_string; - self - } -} - -/// Clears all metadata (Schema level and field level) on an iterator -/// of Schemas -fn clear_metadata( - schemas: impl IntoIterator, -) -> impl Iterator { - schemas.into_iter().map(|schema| { - let fields = schema - .fields() - .iter() - .map(|field| { - field.as_ref().clone().with_metadata(Default::default()) // clear meta - }) - .collect::(); - Schema::new(fields) - }) -} - -async fn fetch_schema_with_location( - store: &dyn ObjectStore, - file: &ObjectMeta, - metadata_size_hint: Option, -) -> Result<(Path, Schema)> { - let loc_path = file.location.clone(); - let schema = fetch_schema(store, file, metadata_size_hint).await?; - Ok((loc_path, schema)) -} - -#[async_trait] -impl FileFormat for ParquetFormat { - fn as_any(&self) -> &dyn Any { - self - } - - fn get_ext(&self) -> String { - ParquetFormatFactory::new().get_ext() - } - - fn get_ext_with_compression( - &self, - file_compression_type: &FileCompressionType, - ) -> Result { - let ext = self.get_ext(); - match file_compression_type.get_variant() { - CompressionTypeVariant::UNCOMPRESSED => Ok(ext), - _ => internal_err!("Parquet FileFormat does not support compression."), - } - } - - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas: Vec<_> = futures::stream::iter(objects) - .map(|object| { - fetch_schema_with_location( - store.as_ref(), - object, - self.metadata_size_hint(), - ) - }) - .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 - .buffered(state.config_options().execution.meta_fetch_concurrency) - .try_collect() - .await?; - - // Schema inference adds fields based the order they are seen - // which depends on the order the files are processed. For some - // object stores (like local file systems) the order returned from list - // is not deterministic. Thus, to ensure deterministic schema inference - // sort the files first. - // https://github.com/apache/datafusion/pull/6629 - schemas.sort_by(|(location1, _), (location2, _)| location1.cmp(location2)); - - let schemas = schemas - .into_iter() - .map(|(_, schema)| schema) - .collect::>(); - - let schema = if self.skip_metadata() { - Schema::try_merge(clear_metadata(schemas)) - } else { - Schema::try_merge(schemas) - }?; - - let schema = if self.binary_as_string() { - transform_binary_to_string(&schema) - } else { - schema - }; - - let schema = if self.force_view_types() { - transform_schema_to_view(&schema) - } else { - schema - }; - - Ok(Arc::new(schema)) - } - - async fn infer_stats( - &self, - _state: &dyn Session, - store: &Arc, - table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result { - let stats = fetch_statistics( - store.as_ref(), - table_schema, - object, - self.metadata_size_hint(), - ) - .await?; - Ok(stats) - } - - async fn create_physical_plan( - &self, - _state: &dyn Session, - conf: FileScanConfig, - filters: Option<&Arc>, - ) -> Result> { - let mut predicate = None; - let mut metadata_size_hint = None; - - // If enable pruning then combine the filters to build the predicate. - // If disable pruning then set the predicate to None, thus readers - // will not prune data based on the statistics. - if self.enable_pruning() { - if let Some(pred) = filters.cloned() { - predicate = Some(pred); - } - } - if let Some(metadata) = self.metadata_size_hint() { - metadata_size_hint = Some(metadata); - } - - let mut source = ParquetSource::new(self.options.clone()); - - if let Some(predicate) = predicate { - source = source.with_predicate(Arc::clone(&conf.file_schema), predicate); - } - if let Some(metadata_size_hint) = metadata_size_hint { - source = source.with_metadata_size_hint(metadata_size_hint) - } - Ok(conf.with_source(Arc::new(source)).build()) - } - - async fn create_writer_physical_plan( - &self, - input: Arc, - _state: &dyn Session, - conf: FileSinkConfig, - order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Parquet"); - } - - let sink = Arc::new(ParquetSink::new(conf, self.options.clone())); - - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) - } - - fn supports_filters_pushdown( - &self, - file_schema: &Schema, - table_schema: &Schema, - filters: &[&Expr], - ) -> Result { - if !self.options().global.pushdown_filters { - return Ok(FilePushdownSupport::NoSupport); - } - - let all_supported = filters.iter().all(|filter| { - can_expr_be_pushed_down_with_schemas(filter, file_schema, table_schema) - }); - - Ok(if all_supported { - FilePushdownSupport::Supported - } else { - FilePushdownSupport::NotSupportedForFilter - }) - } - - fn file_source(&self) -> Arc { - Arc::new(ParquetSource::default()) - } -} - -/// Coerces the file schema if the table schema uses a view type. -#[cfg(not(target_arch = "wasm32"))] -pub fn coerce_file_schema_to_view_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| { - let dt = f.data_type(); - if dt.equals_datatype(&DataType::Utf8View) - || dt.equals_datatype(&DataType::BinaryView) - { - transform = true; - } - (f.name(), dt) - }) - .collect(); - - if !transform { - return None; - } - - let transformed_fields: Vec> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - (Some(DataType::Utf8View), DataType::Utf8 | DataType::LargeUtf8) => { - field_with_new_type(field, DataType::Utf8View) - } - ( - Some(DataType::BinaryView), - DataType::Binary | DataType::LargeBinary, - ) => field_with_new_type(field, DataType::BinaryView), - _ => Arc::clone(field), - }, - ) - .collect(); - - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) -} - -/// If the table schema uses a string type, coerce the file schema to use a string type. -/// -/// See [ParquetFormat::binary_as_string] for details -#[cfg(not(target_arch = "wasm32"))] -pub fn coerce_file_schema_to_string_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| (f.name(), f.data_type())) - .collect(); - let transformed_fields: Vec> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - // table schema uses string type, coerce the file schema to use string type - ( - Some(DataType::Utf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8) - } - // table schema uses large string type, coerce the file schema to use large string type - ( - Some(DataType::LargeUtf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::LargeUtf8) - } - // table schema uses string view type, coerce the file schema to use view type - ( - Some(DataType::Utf8View), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8View) - } - _ => Arc::clone(field), - }, - ) - .collect(); - - if !transform { - None - } else { - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) - } -} - -/// Create a new field with the specified data type, copying the other -/// properties from the input field -fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { - Arc::new(field.as_ref().clone().with_data_type(new_type)) -} - -/// Transform a schema to use view types for Utf8 and Binary -/// -/// See [ParquetFormat::force_view_types] for details -pub fn transform_schema_to_view(schema: &Schema) -> Schema { - let transformed_fields: Vec> = schema - .fields - .iter() - .map(|field| match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => { - field_with_new_type(field, DataType::Utf8View) - } - DataType::Binary | DataType::LargeBinary => { - field_with_new_type(field, DataType::BinaryView) - } - _ => Arc::clone(field), - }) - .collect(); - Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) -} - -/// Transform a schema so that any binary types are strings -pub fn transform_binary_to_string(schema: &Schema) -> Schema { - let transformed_fields: Vec> = schema - .fields - .iter() - .map(|field| match field.data_type() { - DataType::Binary => field_with_new_type(field, DataType::Utf8), - DataType::LargeBinary => field_with_new_type(field, DataType::LargeUtf8), - DataType::BinaryView => field_with_new_type(field, DataType::Utf8View), - _ => Arc::clone(field), - }) - .collect(); - Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) -} - -/// [`MetadataFetch`] adapter for reading bytes from an [`ObjectStore`] -struct ObjectStoreFetch<'a> { - store: &'a dyn ObjectStore, - meta: &'a ObjectMeta, -} - -impl<'a> ObjectStoreFetch<'a> { - fn new(store: &'a dyn ObjectStore, meta: &'a ObjectMeta) -> Self { - Self { store, meta } - } -} - -impl MetadataFetch for ObjectStoreFetch<'_> { - fn fetch( - &mut self, - range: Range, - ) -> BoxFuture<'_, Result> { - async { - self.store - .get_range(&self.meta.location, range) - .await - .map_err(ParquetError::from) - } - .boxed() - } -} - -/// Fetches parquet metadata from ObjectStore for given object -/// -/// This component is a subject to **change** in near future and is exposed for low level integrations -/// through [`ParquetFileReaderFactory`]. -/// -/// [`ParquetFileReaderFactory`]: crate::datasource::physical_plan::ParquetFileReaderFactory -pub async fn fetch_parquet_metadata( - store: &dyn ObjectStore, - meta: &ObjectMeta, - size_hint: Option, -) -> Result { - let file_size = meta.size; - let fetch = ObjectStoreFetch::new(store, meta); - - ParquetMetaDataReader::new() - .with_prefetch_hint(size_hint) - .load_and_finish(fetch, file_size) - .await - .map_err(DataFusionError::from) -} - -/// Read and parse the schema of the Parquet file at location `path` -async fn fetch_schema( - store: &dyn ObjectStore, - file: &ObjectMeta, - metadata_size_hint: Option, -) -> Result { - let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?; - let file_metadata = metadata.file_metadata(); - let schema = parquet_to_arrow_schema( - file_metadata.schema_descr(), - file_metadata.key_value_metadata(), - )?; - Ok(schema) -} - -/// Read and parse the statistics of the Parquet file at location `path` -/// -/// See [`statistics_from_parquet_meta_calc`] for more details -async fn fetch_statistics( - store: &dyn ObjectStore, - table_schema: SchemaRef, - file: &ObjectMeta, - metadata_size_hint: Option, -) -> Result { - let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?; - statistics_from_parquet_meta_calc(&metadata, table_schema) -} - -/// Convert statistics in [`ParquetMetaData`] into [`Statistics`] using ['StatisticsConverter`] -/// -/// The statistics are calculated for each column in the table schema -/// using the row group statistics in the parquet metadata. -pub fn statistics_from_parquet_meta_calc( - metadata: &ParquetMetaData, - table_schema: SchemaRef, -) -> Result { - let row_groups_metadata = metadata.row_groups(); - - let mut statistics = Statistics::new_unknown(&table_schema); - let mut has_statistics = false; - let mut num_rows = 0_usize; - let mut total_byte_size = 0_usize; - for row_group_meta in row_groups_metadata { - num_rows += row_group_meta.num_rows() as usize; - total_byte_size += row_group_meta.total_byte_size() as usize; - - if !has_statistics { - row_group_meta.columns().iter().for_each(|column| { - has_statistics = column.statistics().is_some(); - }); - } - } - statistics.num_rows = Precision::Exact(num_rows); - statistics.total_byte_size = Precision::Exact(total_byte_size); - - let file_metadata = metadata.file_metadata(); - let mut file_schema = parquet_to_arrow_schema( - file_metadata.schema_descr(), - file_metadata.key_value_metadata(), - )?; - if let Some(merged) = coerce_file_schema_to_string_type(&table_schema, &file_schema) { - file_schema = merged; - } - - if let Some(merged) = coerce_file_schema_to_view_type(&table_schema, &file_schema) { - file_schema = merged; - } - - statistics.column_statistics = if has_statistics { - let (mut max_accs, mut min_accs) = create_max_min_accs(&table_schema); - let mut null_counts_array = - vec![Precision::Exact(0); table_schema.fields().len()]; - - table_schema - .fields() - .iter() - .enumerate() - .for_each(|(idx, field)| { - match StatisticsConverter::try_new( - field.name(), - &file_schema, - file_metadata.schema_descr(), - ) { - Ok(stats_converter) => { - summarize_min_max_null_counts( - &mut min_accs, - &mut max_accs, - &mut null_counts_array, - idx, - num_rows, - &stats_converter, - row_groups_metadata, - ) - .ok(); - } - Err(e) => { - debug!("Failed to create statistics converter: {}", e); - null_counts_array[idx] = Precision::Exact(num_rows); - } - } - }); - - get_col_stats( - &table_schema, - null_counts_array, - &mut max_accs, - &mut min_accs, - ) - } else { - Statistics::unknown_column(&table_schema) - }; - - Ok(statistics) -} - -/// Deprecated -/// Use [`statistics_from_parquet_meta_calc`] instead. -/// This method was deprecated because it didn't need to be async so a new method was created -/// that exposes a synchronous API. -#[deprecated( - since = "40.0.0", - note = "please use `statistics_from_parquet_meta_calc` instead" -)] -pub async fn statistics_from_parquet_meta( - metadata: &ParquetMetaData, - table_schema: SchemaRef, -) -> Result { - statistics_from_parquet_meta_calc(metadata, table_schema) -} - -fn summarize_min_max_null_counts( - min_accs: &mut [Option], - max_accs: &mut [Option], - null_counts_array: &mut [Precision], - arrow_schema_index: usize, - num_rows: usize, - stats_converter: &StatisticsConverter, - row_groups_metadata: &[RowGroupMetaData], -) -> Result<()> { - let max_values = stats_converter.row_group_maxes(row_groups_metadata)?; - let min_values = stats_converter.row_group_mins(row_groups_metadata)?; - let null_counts = stats_converter.row_group_null_counts(row_groups_metadata)?; - - if let Some(max_acc) = &mut max_accs[arrow_schema_index] { - max_acc.update_batch(&[max_values])?; - } - - if let Some(min_acc) = &mut min_accs[arrow_schema_index] { - min_acc.update_batch(&[min_values])?; - } - - null_counts_array[arrow_schema_index] = Precision::Exact(match sum(&null_counts) { - Some(null_count) => null_count as usize, - None => num_rows, - }); - - Ok(()) -} - -/// Implements [`DataSink`] for writing to a parquet file. -pub struct ParquetSink { - /// Config options for writing data - config: FileSinkConfig, - /// Underlying parquet options - parquet_options: TableParquetOptions, - /// File metadata from successfully produced parquet files. The Mutex is only used - /// to allow inserting to HashMap from behind borrowed reference in DataSink::write_all. - written: Arc>>, -} - -impl Debug for ParquetSink { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ParquetSink").finish() - } -} - -impl DisplayAs for ParquetSink { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "ParquetSink(file_groups=",)?; - FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; - write!(f, ")") - } - } - } -} - -impl ParquetSink { - /// Create from config. - pub fn new(config: FileSinkConfig, parquet_options: TableParquetOptions) -> Self { - Self { - config, - parquet_options, - written: Default::default(), - } - } - - /// Retrieve the file metadata for the written files, keyed to the path - /// which may be partitioned (in the case of hive style partitioning). - pub fn written(&self) -> HashMap { - self.written.lock().clone() - } - - /// Create writer properties based upon configuration settings, - /// including partitioning and the inclusion of arrow schema metadata. - fn create_writer_props(&self) -> Result { - let schema = if self.parquet_options.global.allow_single_file_parallelism { - // If parallelizing writes, we may be also be doing hive style partitioning - // into multiple files which impacts the schema per file. - // Refer to `get_writer_schema()` - &get_writer_schema(&self.config) - } else { - self.config.output_schema() - }; - - // TODO: avoid this clone in follow up PR, where the writer properties & schema - // are calculated once on `ParquetSink::new` - let mut parquet_opts = self.parquet_options.clone(); - if !self.parquet_options.global.skip_arrow_metadata { - parquet_opts.arrow_schema(schema); - } - - Ok(WriterPropertiesBuilder::try_from(&parquet_opts)?.build()) - } - - /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore - /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized - async fn create_async_arrow_writer( - &self, - location: &Path, - object_store: Arc, - parquet_props: WriterProperties, - ) -> Result> { - let buf_writer = BufWriter::new(object_store, location.clone()); - let options = ArrowWriterOptions::new() - .with_properties(parquet_props) - .with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata); - - let writer = AsyncArrowWriter::try_new_with_options( - buf_writer, - get_writer_schema(&self.config), - options, - )?; - Ok(writer) - } - - /// Parquet options - pub fn parquet_options(&self) -> &TableParquetOptions { - &self.parquet_options - } -} - -#[async_trait] -impl FileSink for ParquetSink { - fn config(&self) -> &FileSinkConfig { - &self.config - } - - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, - demux_task: SpawnedTask>, - mut file_stream_rx: DemuxedStreamReceiver, - object_store: Arc, - ) -> Result { - let parquet_opts = &self.parquet_options; - let allow_single_file_parallelism = - parquet_opts.global.allow_single_file_parallelism; - - let mut file_write_tasks: JoinSet< - std::result::Result<(Path, FileMetaData), DataFusionError>, - > = JoinSet::new(); - - let parquet_props = self.create_writer_props()?; - let parallel_options = ParallelParquetWriterOptions { - max_parallel_row_groups: parquet_opts - .global - .maximum_parallel_row_group_writers, - max_buffered_record_batches_per_stream: parquet_opts - .global - .maximum_buffered_record_batches_per_stream, - }; - - while let Some((path, mut rx)) = file_stream_rx.recv().await { - if !allow_single_file_parallelism { - let mut writer = self - .create_async_arrow_writer( - &path, - Arc::clone(&object_store), - parquet_props.clone(), - ) - .await?; - let mut reservation = - MemoryConsumer::new(format!("ParquetSink[{}]", path)) - .register(context.memory_pool()); - file_write_tasks.spawn(async move { - while let Some(batch) = rx.recv().await { - writer.write(&batch).await?; - reservation.try_resize(writer.memory_size())?; - } - let file_metadata = writer - .close() - .await - .map_err(DataFusionError::ParquetError)?; - Ok((path, file_metadata)) - }); - } else { - let writer = create_writer( - // Parquet files as a whole are never compressed, since they - // manage compressed blocks themselves. - FileCompressionType::UNCOMPRESSED, - &path, - Arc::clone(&object_store), - ) - .await?; - let schema = get_writer_schema(&self.config); - let props = parquet_props.clone(); - let parallel_options_clone = parallel_options.clone(); - let pool = Arc::clone(context.memory_pool()); - file_write_tasks.spawn(async move { - let file_metadata = output_single_parquet_file_parallelized( - writer, - rx, - schema, - &props, - parallel_options_clone, - pool, - ) - .await?; - Ok((path, file_metadata)) - }); - } - } - - let mut row_count = 0; - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - let (path, file_metadata) = r?; - row_count += file_metadata.num_rows; - let mut written_files = self.written.lock(); - written_files - .try_insert(path.clone(), file_metadata) - .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; - drop(written_files); - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } - - demux_task - .join_unwind() - .await - .map_err(DataFusionError::ExecutionJoin)??; - - Ok(row_count as u64) - } -} - -#[async_trait] -impl DataSink for ParquetSink { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> &SchemaRef { - self.config.output_schema() - } - - async fn write_all( - &self, - data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await - } -} - -/// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] -/// Once the channel is exhausted, returns the ArrowColumnWriter. -async fn column_serializer_task( - mut rx: Receiver, - mut writer: ArrowColumnWriter, - mut reservation: MemoryReservation, -) -> Result<(ArrowColumnWriter, MemoryReservation)> { - while let Some(col) = rx.recv().await { - writer.write(&col)?; - reservation.try_resize(writer.memory_size())?; - } - Ok((writer, reservation)) -} - -type ColumnWriterTask = SpawnedTask>; -type ColSender = Sender; - -/// Spawns a parallel serialization task for each column -/// Returns join handles for each columns serialization task along with a send channel -/// to send arrow arrays to each serialization task. -fn spawn_column_parallel_row_group_writer( - schema: Arc, - parquet_props: Arc, - max_buffer_size: usize, - pool: &Arc, -) -> Result<(Vec, Vec)> { - let schema_desc = ArrowSchemaConverter::new().convert(&schema)?; - let col_writers = get_column_writers(&schema_desc, &parquet_props, &schema)?; - let num_columns = col_writers.len(); - - let mut col_writer_tasks = Vec::with_capacity(num_columns); - let mut col_array_channels = Vec::with_capacity(num_columns); - for writer in col_writers.into_iter() { - // Buffer size of this channel limits the number of arrays queued up for column level serialization - let (send_array, receive_array) = - mpsc::channel::(max_buffer_size); - col_array_channels.push(send_array); - - let reservation = - MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(pool); - let task = SpawnedTask::spawn(column_serializer_task( - receive_array, - writer, - reservation, - )); - col_writer_tasks.push(task); - } - - Ok((col_writer_tasks, col_array_channels)) -} - -/// Settings related to writing parquet files in parallel -#[derive(Clone)] -struct ParallelParquetWriterOptions { - max_parallel_row_groups: usize, - max_buffered_record_batches_per_stream: usize, -} - -/// This is the return type of calling [ArrowColumnWriter].close() on each column -/// i.e. the Vec of encoded columns which can be appended to a row group -type RBStreamSerializeResult = Result<(Vec, MemoryReservation, usize)>; - -/// Sends the ArrowArrays in passed [RecordBatch] through the channels to their respective -/// parallel column serializers. -async fn send_arrays_to_col_writers( - col_array_channels: &[ColSender], - rb: &RecordBatch, - schema: Arc, -) -> Result<()> { - // Each leaf column has its own channel, increment next_channel for each leaf column sent. - let mut next_channel = 0; - for (array, field) in rb.columns().iter().zip(schema.fields()) { - for c in compute_leaves(field, array)? { - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if col_array_channels[next_channel].send(c).await.is_err() { - return Ok(()); - } - - next_channel += 1; - } - } - - Ok(()) -} - -/// Spawns a tokio task which joins the parallel column writer tasks, -/// and finalizes the row group -fn spawn_rg_join_and_finalize_task( - column_writer_tasks: Vec, - rg_rows: usize, - pool: &Arc, -) -> SpawnedTask { - let mut rg_reservation = - MemoryConsumer::new("ParquetSink(SerializedRowGroupWriter)").register(pool); - - SpawnedTask::spawn(async move { - let num_cols = column_writer_tasks.len(); - let mut finalized_rg = Vec::with_capacity(num_cols); - for task in column_writer_tasks.into_iter() { - let (writer, _col_reservation) = task - .join_unwind() - .await - .map_err(DataFusionError::ExecutionJoin)??; - let encoded_size = writer.get_estimated_total_bytes(); - rg_reservation.grow(encoded_size); - finalized_rg.push(writer.close()?); - } - - Ok((finalized_rg, rg_reservation, rg_rows)) - }) -} - -/// This task coordinates the serialization of a parquet file in parallel. -/// As the query produces RecordBatches, these are written to a RowGroup -/// via parallel [ArrowColumnWriter] tasks. Once the desired max rows per -/// row group is reached, the parallel tasks are joined on another separate task -/// and sent to a concatenation task. This task immediately continues to work -/// on the next row group in parallel. So, parquet serialization is parallelized -/// across both columns and row_groups, with a theoretical max number of parallel tasks -/// given by n_columns * num_row_groups. -fn spawn_parquet_parallel_serialization_task( - mut data: Receiver, - serialize_tx: Sender>, - schema: Arc, - writer_props: Arc, - parallel_options: ParallelParquetWriterOptions, - pool: Arc, -) -> SpawnedTask> { - SpawnedTask::spawn(async move { - let max_buffer_rb = parallel_options.max_buffered_record_batches_per_stream; - let max_row_group_rows = writer_props.max_row_group_size(); - let (mut column_writer_handles, mut col_array_channels) = - spawn_column_parallel_row_group_writer( - Arc::clone(&schema), - Arc::clone(&writer_props), - max_buffer_rb, - &pool, - )?; - let mut current_rg_rows = 0; - - while let Some(mut rb) = data.recv().await { - // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case - // when max_row_group_rows < execution.batch_size as an alternative to a recursive async - // function. - loop { - if current_rg_rows + rb.num_rows() < max_row_group_rows { - send_arrays_to_col_writers( - &col_array_channels, - &rb, - Arc::clone(&schema), - ) - .await?; - current_rg_rows += rb.num_rows(); - break; - } else { - let rows_left = max_row_group_rows - current_rg_rows; - let a = rb.slice(0, rows_left); - send_arrays_to_col_writers( - &col_array_channels, - &a, - Arc::clone(&schema), - ) - .await?; - - // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup - // on a separate task, so that we can immediately start on the next RG before waiting - // for the current one to finish. - drop(col_array_channels); - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - max_row_group_rows, - &pool, - ); - - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { - return Ok(()); - } - - current_rg_rows = 0; - rb = rb.slice(rows_left, rb.num_rows() - rows_left); - - (column_writer_handles, col_array_channels) = - spawn_column_parallel_row_group_writer( - Arc::clone(&schema), - Arc::clone(&writer_props), - max_buffer_rb, - &pool, - )?; - } - } - } - - drop(col_array_channels); - // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows - if current_rg_rows > 0 { - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - current_rg_rows, - &pool, - ); +//! Re-exports the [`datafusion_datasource_parquet::file_format`] module, and contains tests for it. - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { - return Ok(()); - } - } - - Ok(()) - }) -} - -/// Consume RowGroups serialized by other parallel tasks and concatenate them in -/// to the final parquet file, while flushing finalized bytes to an [ObjectStore] -async fn concatenate_parallel_row_groups( - mut serialize_rx: Receiver>, - schema: Arc, - writer_props: Arc, - mut object_store_writer: Box, - pool: Arc, -) -> Result { - let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES); - - let mut file_reservation = - MemoryConsumer::new("ParquetSink(SerializedFileWriter)").register(&pool); - - let schema_desc = ArrowSchemaConverter::new().convert(schema.as_ref())?; - let mut parquet_writer = SerializedFileWriter::new( - merged_buff.clone(), - schema_desc.root_schema_ptr(), - writer_props, - )?; - - while let Some(task) = serialize_rx.recv().await { - let result = task.join_unwind().await; - let mut rg_out = parquet_writer.next_row_group()?; - let (serialized_columns, mut rg_reservation, _cnt) = - result.map_err(DataFusionError::ExecutionJoin)??; - for chunk in serialized_columns { - chunk.append_to_row_group(&mut rg_out)?; - rg_reservation.free(); - - let mut buff_to_flush = merged_buff.buffer.try_lock().unwrap(); - file_reservation.try_resize(buff_to_flush.len())?; - - if buff_to_flush.len() > BUFFER_FLUSH_BYTES { - object_store_writer - .write_all(buff_to_flush.as_slice()) - .await?; - buff_to_flush.clear(); - file_reservation.try_resize(buff_to_flush.len())?; // will set to zero - } - } - rg_out.close()?; - } - - let file_metadata = parquet_writer.close()?; - let final_buff = merged_buff.buffer.try_lock().unwrap(); - - object_store_writer.write_all(final_buff.as_slice()).await?; - object_store_writer.shutdown().await?; - file_reservation.free(); - - Ok(file_metadata) -} - -/// Parallelizes the serialization of a single parquet file, by first serializing N -/// independent RecordBatch streams in parallel to RowGroups in memory. Another -/// task then stitches these independent RowGroups together and streams this large -/// single parquet file to an ObjectStore in multiple parts. -async fn output_single_parquet_file_parallelized( - object_store_writer: Box, - data: Receiver, - output_schema: Arc, - parquet_props: &WriterProperties, - parallel_options: ParallelParquetWriterOptions, - pool: Arc, -) -> Result { - let max_rowgroups = parallel_options.max_parallel_row_groups; - // Buffer size of this channel limits maximum number of RowGroups being worked on in parallel - let (serialize_tx, serialize_rx) = - mpsc::channel::>(max_rowgroups); - - let arc_props = Arc::new(parquet_props.clone()); - let launch_serialization_task = spawn_parquet_parallel_serialization_task( - data, - serialize_tx, - Arc::clone(&output_schema), - Arc::clone(&arc_props), - parallel_options, - Arc::clone(&pool), - ); - let file_metadata = concatenate_parallel_row_groups( - serialize_rx, - Arc::clone(&output_schema), - Arc::clone(&arc_props), - object_store_writer, - pool, - ) - .await?; - - launch_serialization_task - .join_unwind() - .await - .map_err(DataFusionError::ExecutionJoin)??; - Ok(file_metadata) -} +pub use datafusion_datasource_parquet::file_format::*; #[cfg(test)] pub(crate) mod test_util { - use super::*; - use crate::test::object_store::local_unpartitioned_file; - - use parquet::arrow::ArrowWriter; - use tempfile::NamedTempFile; + use arrow::array::RecordBatch; + use datafusion_common::Result; + use object_store::ObjectMeta; - /// How many rows per page should be written - const ROWS_PER_PAGE: usize = 2; + use crate::test::object_store::local_unpartitioned_file; /// Writes `batches` to a temporary parquet file /// @@ -1378,12 +35,28 @@ pub(crate) mod test_util { pub async fn store_parquet( batches: Vec, multi_page: bool, - ) -> Result<(Vec, Vec)> { + ) -> Result<(Vec, Vec)> { + /// How many rows per page should be written + const ROWS_PER_PAGE: usize = 2; + /// write batches chunk_size rows at a time + fn write_in_chunks( + writer: &mut parquet::arrow::ArrowWriter, + batch: &RecordBatch, + chunk_size: usize, + ) { + let mut i = 0; + while i < batch.num_rows() { + let num = chunk_size.min(batch.num_rows() - i); + writer.write(&batch.slice(i, num)).unwrap(); + i += num; + } + } + // we need the tmp files to be sorted as some tests rely on the how the returning files are ordered // https://github.com/apache/datafusion/pull/6629 let tmp_files = { let mut tmp_files: Vec<_> = (0..batches.len()) - .map(|_| NamedTempFile::new().expect("creating temp file")) + .map(|_| tempfile::NamedTempFile::new().expect("creating temp file")) .collect(); tmp_files.sort_by(|a, b| a.path().cmp(b.path())); tmp_files @@ -1394,7 +67,7 @@ pub(crate) mod test_util { .into_iter() .zip(tmp_files.into_iter()) .map(|(batch, mut output)| { - let builder = WriterProperties::builder(); + let builder = parquet::file::properties::WriterProperties::builder(); let props = if multi_page { builder.set_data_page_row_count_limit(ROWS_PER_PAGE) } else { @@ -1402,9 +75,12 @@ pub(crate) mod test_util { } .build(); - let mut writer = - ArrowWriter::try_new(&mut output, batch.schema(), Some(props)) - .expect("creating writer"); + let mut writer = parquet::arrow::ArrowWriter::try_new( + &mut output, + batch.schema(), + Some(props), + ) + .expect("creating writer"); if multi_page { // write in smaller batches as the parquet writer @@ -1422,65 +98,69 @@ pub(crate) mod test_util { Ok((meta, files)) } - - /// write batches chunk_size rows at a time - fn write_in_chunks( - writer: &mut ArrowWriter, - batch: &RecordBatch, - chunk_size: usize, - ) { - let mut i = 0; - while i < batch.num_rows() { - let num = chunk_size.min(batch.num_rows() - i); - writer.write(&batch.slice(i, num)).unwrap(); - i += num; - } - } } #[cfg(test)] mod tests { - use super::super::test_util::scan_format; - use crate::datasource::listing::{ListingTableUrl, PartitionedFile}; - use crate::execution::SessionState; - use crate::physical_plan::collect; - use crate::test_util::bounded_stream; - use std::fmt::{Display, Formatter}; + + use std::fmt::{self, Display, Formatter}; + use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::task::{Context, Poll}; use std::time::Duration; - use super::*; - use crate::datasource::file_format::parquet::test_util::store_parquet; + use crate::datasource::file_format::test_util::scan_format; + use crate::execution::SessionState; use crate::physical_plan::metrics::MetricValue; use crate::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; - use arrow::array::{ - types::Int32Type, Array, ArrayRef, DictionaryArray, Int32Array, Int64Array, - StringArray, - }; - use arrow::datatypes::{DataType, Field}; - use async_trait::async_trait; + + use arrow::array::RecordBatch; + use arrow_schema::{Schema, SchemaRef}; + use datafusion_catalog::Session; use datafusion_common::cast::{ as_binary_array, as_binary_view_array, as_boolean_array, as_float32_array, as_float64_array, as_int32_array, as_timestamp_nanosecond_array, }; - use datafusion_common::config::ParquetOptions; + use datafusion_common::config::{ParquetOptions, TableParquetOptions}; + use datafusion_common::stats::Precision; use datafusion_common::ScalarValue::Utf8; - use datafusion_common::{assert_batches_eq, ScalarValue}; + use datafusion_common::{assert_batches_eq, Result, ScalarValue}; + use datafusion_datasource::file_format::FileFormat; + use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; + use datafusion_datasource::{ListingTableUrl, PartitionedFile}; + use datafusion_datasource_parquet::{ + fetch_parquet_metadata, fetch_statistics, statistics_from_parquet_meta_calc, + ParquetFormat, ParquetFormatFactory, ParquetSink, + }; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_execution::{RecordBatchStream, TaskContext}; + use datafusion_expr::dml::InsertOp; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use datafusion_physical_plan::{collect, ExecutionPlan}; + + use arrow::array::{ + types::Int32Type, Array, ArrayRef, DictionaryArray, Int32Array, Int64Array, + StringArray, + }; + use arrow::datatypes::{DataType, Field}; + use async_trait::async_trait; use futures::stream::BoxStream; + use futures::{Stream, StreamExt}; use log::error; use object_store::local::LocalFileSystem; + use object_store::ObjectMeta; use object_store::{ - GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOpts, PutOptions, - PutPayload, PutResult, + path::Path, GetOptions, GetResult, ListResult, MultipartUpload, ObjectStore, + PutMultipartOpts, PutOptions, PutPayload, PutResult, }; use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::arrow::ParquetRecordBatchStreamBuilder; use parquet::file::metadata::{KeyValue, ParquetColumnIndex, ParquetOffsetIndex}; use parquet::file::page_index::index::Index; + use parquet::format::FileMetaData; use tokio::fs::File; enum ForceViews { @@ -2302,7 +982,7 @@ mod tests { } #[tokio::test] async fn test_read_parquet_page_index() -> Result<()> { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let path = format!("{testdata}/alltypes_tiny_pages.parquet"); let file = File::open(path).await.unwrap(); let options = ArrowReaderOptions::new().with_page_index(true); @@ -2385,7 +1065,7 @@ mod tests { projection: Option>, limit: Option, ) -> Result> { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let state = state.as_any().downcast_ref::().unwrap(); let format = state .get_file_format_factory("parquet") @@ -2954,4 +1634,43 @@ mod tests { Ok(()) } + + /// Creates an bounded stream for testing purposes. + fn bounded_stream( + batch: RecordBatch, + limit: usize, + ) -> datafusion_execution::SendableRecordBatchStream { + Box::pin(BoundedStream { + count: 0, + limit, + batch, + }) + } + + struct BoundedStream { + limit: usize, + count: usize, + batch: RecordBatch, + } + + impl Stream for BoundedStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if self.count >= self.limit { + return Poll::Ready(None); + } + self.count += 1; + Poll::Ready(Some(Ok(self.batch.clone()))) + } + } + + impl RecordBatchStream for BoundedStream { + fn schema(&self) -> SchemaRef { + self.batch.schema() + } + } } diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 41e939d60b08..4d7762784d78 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -616,6 +616,7 @@ impl ListingOptions { /// using an [`ObjectStore`] instance, for example from local files or objects /// from AWS S3. /// +/// # Reading Directories /// For example, given the `table1` directory (or object store prefix) /// /// ```text @@ -651,13 +652,19 @@ impl ListingOptions { /// If the query has a predicate like `WHERE date = '2024-06-01'` /// only the corresponding directory will be read. /// -/// `ListingTable` also supports filter and projection pushdown for formats that +/// `ListingTable` also supports limit, filter and projection pushdown for formats that /// support it as such as Parquet. /// +/// # Implementation +/// +/// `ListingTable` Uses [`DataSourceExec`] to execute the data. See that struct +/// for more details. +/// +/// [`DataSourceExec`]: crate::datasource::source::DataSourceExec +/// /// # Example /// -/// Here is an example of reading a directory of parquet files using a -/// [`ListingTable`]: +/// To read a directory of parquet files using a [`ListingTable`]: /// /// ```no_run /// # use datafusion::prelude::SessionContext; @@ -1098,7 +1105,9 @@ impl ListingTable { ) })) .await?; - let file_list = stream::iter(file_list).flatten(); + let meta_fetch_concurrency = + ctx.config_options().execution.meta_fetch_concurrency; + let file_list = stream::iter(file_list).flatten_unordered(meta_fetch_concurrency); // collect the statistics if required by the config let files = file_list .map(|part_file| async { @@ -1115,7 +1124,7 @@ impl ListingTable { } }) .boxed() - .buffered(ctx.config_options().execution.meta_fetch_concurrency); + .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency); let (files, statistics) = get_statistics_with_limit( files, @@ -1173,7 +1182,6 @@ impl ListingTable { #[cfg(test)] mod tests { use super::*; - use crate::datasource::file_format::avro::AvroFormat; use crate::datasource::file_format::csv::CsvFormat; use crate::datasource::file_format::json::JsonFormat; #[cfg(feature = "parquet")] @@ -1185,7 +1193,6 @@ mod tests { assert_batches_eq, test::{columns, object_store::register_test_store}, }; - use datafusion_physical_plan::collect; use arrow::compute::SortOptions; use arrow::record_batch::RecordBatch; @@ -1193,9 +1200,12 @@ mod tests { use datafusion_common::{assert_contains, ScalarValue}; use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator}; use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_plan::collect; use datafusion_physical_plan::ExecutionPlanProperties; + use crate::test::object_store::{ensure_head_concurrency, make_test_store_and_state}; use tempfile::TempDir; + use url::Url; #[tokio::test] async fn read_single_file() -> Result<()> { @@ -1296,7 +1306,7 @@ mod tests { vec![vec![ col("int_col").add(lit(1)).sort(true, true), ]], - Err("Expected single column references in output_ordering, got int_col + Int32(1)"), + Err("Expected single column reference in sort_order[0][0], got int_col + Int32(1)"), ), // ok with one column ( @@ -1365,11 +1375,14 @@ mod tests { #[tokio::test] async fn read_empty_table() -> Result<()> { let ctx = SessionContext::new(); - let path = String::from("table/p1=v1/file.avro"); + let path = String::from("table/p1=v1/file.json"); register_test_store(&ctx, &[(&path, 100)]); - let opt = ListingOptions::new(Arc::new(AvroFormat {})) - .with_file_extension(AvroFormat.get_ext()) + let format = JsonFormat::default(); + let ext = format.get_ext(); + + let opt = ListingOptions::new(Arc::new(format)) + .with_file_extension(ext) .with_table_partition_cols(vec![(String::from("p1"), DataType::Utf8)]) .with_target_partitions(4); @@ -1479,9 +1492,9 @@ mod tests { // files that don't match the prefix or the default file extention assert_list_files_for_scan_grouping( &[ - "bucket/key-prefix/file0.avro", + "bucket/key-prefix/file0.json", "bucket/key-prefix/file1.parquet", - "bucket/other-prefix/roguefile.avro", + "bucket/other-prefix/roguefile.json", ], "test:///bucket/key-prefix/", 10, @@ -1569,11 +1582,11 @@ mod tests { // files that don't match the prefix or the default file ext assert_list_files_for_multi_paths( &[ - "bucket/key1/file0.avro", + "bucket/key1/file0.json", "bucket/key1/file1.csv", - "bucket/key1/file2.avro", + "bucket/key1/file2.json", "bucket/key2/file3.csv", - "bucket/key2/file4.avro", + "bucket/key2/file4.json", "bucket/key3/file5.csv", ], &["test:///bucket/key1/", "test:///bucket/key3/"], @@ -1585,6 +1598,81 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_assert_list_files_for_exact_paths() -> Result<()> { + // more expected partitions than files + assert_list_files_for_exact_paths( + &[ + "bucket/key1/file0", + "bucket/key1/file1", + "bucket/key1/file2", + "bucket/key2/file3", + "bucket/key2/file4", + ], + 12, + 5, + Some(""), + ) + .await?; + + // more files than meta_fetch_concurrency (32) + let files: Vec = + (0..64).map(|i| format!("bucket/key1/file{}", i)).collect(); + // Collect references to each string + let file_refs: Vec<&str> = files.iter().map(|s| s.as_str()).collect(); + assert_list_files_for_exact_paths(file_refs.as_slice(), 5, 5, Some("")).await?; + + // as many expected partitions as files + assert_list_files_for_exact_paths( + &[ + "bucket/key1/file0", + "bucket/key1/file1", + "bucket/key1/file2", + "bucket/key2/file3", + "bucket/key2/file4", + ], + 5, + 5, + Some(""), + ) + .await?; + + // more files as expected partitions + assert_list_files_for_exact_paths( + &[ + "bucket/key1/file0", + "bucket/key1/file1", + "bucket/key1/file2", + "bucket/key2/file3", + "bucket/key2/file4", + ], + 2, + 2, + Some(""), + ) + .await?; + + // no files => no groups + assert_list_files_for_exact_paths(&[], 2, 0, Some("")).await?; + + // files that don't match the default file ext + assert_list_files_for_exact_paths( + &[ + "bucket/key1/file0.json", + "bucket/key1/file1.csv", + "bucket/key1/file2.json", + "bucket/key2/file3.csv", + "bucket/key2/file4.json", + "bucket/key3/file5.csv", + ], + 2, + 2, + None, + ) + .await?; + Ok(()) + } + async fn load_table( ctx: &SessionContext, name: &str, @@ -1612,9 +1700,7 @@ mod tests { let ctx = SessionContext::new(); register_test_store(&ctx, &files.iter().map(|f| (*f, 10)).collect::>()); - let format = AvroFormat {}; - - let opt = ListingOptions::new(Arc::new(format)) + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) .with_file_extension_opt(file_ext) .with_target_partitions(target_partitions); @@ -1646,9 +1732,7 @@ mod tests { let ctx = SessionContext::new(); register_test_store(&ctx, &files.iter().map(|f| (*f, 10)).collect::>()); - let format = AvroFormat {}; - - let opt = ListingOptions::new(Arc::new(format)) + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) .with_file_extension_opt(file_ext) .with_target_partitions(target_partitions); @@ -1671,6 +1755,56 @@ mod tests { Ok(()) } + /// Check that the files listed by the table match the specified `output_partitioning` + /// when the object store contains `files`, and validate that file metadata is fetched + /// concurrently + async fn assert_list_files_for_exact_paths( + files: &[&str], + target_partitions: usize, + output_partitioning: usize, + file_ext: Option<&str>, + ) -> Result<()> { + let ctx = SessionContext::new(); + let (store, _) = make_test_store_and_state( + &files.iter().map(|f| (*f, 10)).collect::>(), + ); + + let meta_fetch_concurrency = ctx + .state() + .config_options() + .execution + .meta_fetch_concurrency; + let expected_concurrency = files.len().min(meta_fetch_concurrency); + let head_blocking_store = ensure_head_concurrency(store, expected_concurrency); + + let url = Url::parse("test://").unwrap(); + ctx.register_object_store(&url, head_blocking_store.clone()); + + let format = JsonFormat::default(); + + let opt = ListingOptions::new(Arc::new(format)) + .with_file_extension_opt(file_ext) + .with_target_partitions(target_partitions); + + let schema = Schema::new(vec![Field::new("a", DataType::Boolean, false)]); + + let table_paths = files + .iter() + .map(|t| ListingTableUrl::parse(format!("test:///{}", t)).unwrap()) + .collect(); + let config = ListingTableConfig::new_with_multi_paths(table_paths) + .with_listing_options(opt) + .with_schema(Arc::new(schema)); + + let table = ListingTable::try_new(config)?; + + let (file_list, _) = table.list_files_for_scan(&ctx.state(), &[], None).await?; + + assert_eq!(file_list.len(), output_partitioning); + + Ok(()) + } + #[tokio::test] async fn test_insert_into_append_new_json_files() -> Result<()> { let mut config_map: HashMap = HashMap::new(); @@ -1707,6 +1841,7 @@ mod tests { Ok(()) } + #[cfg(feature = "parquet")] #[tokio::test] async fn test_insert_into_append_2_new_parquet_files_defaults() -> Result<()> { let mut config_map: HashMap = HashMap::new(); @@ -1725,6 +1860,7 @@ mod tests { Ok(()) } + #[cfg(feature = "parquet")] #[tokio::test] async fn test_insert_into_append_1_new_parquet_files_defaults() -> Result<()> { let mut config_map: HashMap = HashMap::new(); @@ -2011,6 +2147,7 @@ mod tests { ) .await?; } + #[cfg(feature = "parquet")] "parquet" => { session_ctx .register_parquet( @@ -2020,6 +2157,7 @@ mod tests { ) .await?; } + #[cfg(feature = "avro")] "avro" => { session_ctx .register_avro( diff --git a/datafusion/core/src/datasource/memory.rs b/datafusion/core/src/datasource/memory.rs index b8bec410070c..d96944fa7a69 100644 --- a/datafusion/core/src/datasource/memory.rs +++ b/datafusion/core/src/datasource/memory.rs @@ -315,6 +315,10 @@ impl DisplayAs for MemSink { let partition_count = self.batches.len(); write!(f, "MemoryTable (partitions={partition_count})") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs index 2b7bb14b6f6c..18a1318dd40d 100644 --- a/datafusion/core/src/datasource/mod.rs +++ b/datafusion/core/src/datasource/mod.rs @@ -19,7 +19,6 @@ //! //! [`ListingTable`]: crate::datasource::listing::ListingTable -pub mod avro_to_arrow; pub mod cte_worktable; pub mod default_table_source; pub mod dynamic_file; @@ -30,11 +29,11 @@ pub mod listing_table_factory; pub mod memory; pub mod physical_plan; pub mod provider; -pub mod schema_adapter; mod statistics; pub mod stream; pub mod view; +pub use datafusion_datasource::schema_adapter; pub use datafusion_datasource::source; // backwards compatibility @@ -54,16 +53,55 @@ use datafusion_common::{plan_err, Result}; use datafusion_expr::{Expr, SortExpr}; use datafusion_physical_expr::{expressions, LexOrdering, PhysicalSortExpr}; -fn create_ordering( +/// Converts logical sort expressions to physical sort expressions +/// +/// This function transforms a collection of logical sort expressions into their physical +/// representation that can be used during query execution. +/// +/// # Arguments +/// +/// * `schema` - The schema containing column definitions +/// * `sort_order` - A collection of logical sort expressions grouped into lexicographic orderings +/// +/// # Returns +/// +/// A vector of lexicographic orderings for physical execution, or an error if the transformation fails +/// +/// # Examples +/// +/// ``` +/// // Create orderings from columns "id" and "name" +/// # use arrow::datatypes::{Schema, Field, DataType}; +/// # use datafusion::datasource::create_ordering; +/// # use datafusion_common::Column; +/// # use datafusion_expr::{Expr, SortExpr}; +/// # +/// // Create a schema with two fields +/// let schema = Schema::new(vec![ +/// Field::new("id", DataType::Int32, false), +/// Field::new("name", DataType::Utf8, false), +/// ]); +/// +/// let sort_exprs = vec![ +/// vec![ +/// SortExpr { expr: Expr::Column(Column::new(Some("t"), "id")), asc: true, nulls_first: false } +/// ], +/// vec![ +/// SortExpr { expr: Expr::Column(Column::new(Some("t"), "name")), asc: false, nulls_first: true } +/// ] +/// ]; +/// let result = create_ordering(&schema, &sort_exprs).unwrap(); +/// ``` +pub fn create_ordering( schema: &Schema, sort_order: &[Vec], ) -> Result> { let mut all_sort_orders = vec![]; - for exprs in sort_order { + for (group_idx, exprs) in sort_order.iter().enumerate() { // Construct PhysicalSortExpr objects from Expr objects: let mut sort_exprs = LexOrdering::default(); - for sort in exprs { + for (expr_idx, sort) in exprs.iter().enumerate() { match &sort.expr { Expr::Column(col) => match expressions::col(&col.name, schema) { Ok(expr) => { @@ -81,8 +119,11 @@ fn create_ordering( }, expr => { return plan_err!( - "Expected single column references in output_ordering, got {expr}" - ) + "Expected single column reference in sort_order[{}][{}], got {}", + group_idx, + expr_idx, + expr + ); } } } @@ -92,3 +133,226 @@ fn create_ordering( } Ok(all_sort_orders) } + +#[cfg(all(test, feature = "parquet"))] +mod tests { + + use crate::prelude::SessionContext; + + use std::fs; + use std::sync::Arc; + + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use arrow::record_batch::RecordBatch; + use datafusion_common::assert_batches_sorted_eq; + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::schema_adapter::{ + DefaultSchemaAdapterFactory, SchemaAdapter, SchemaAdapterFactory, SchemaMapper, + }; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource_parquet::source::ParquetSource; + + use datafusion_common::record_batch; + + use ::object_store::path::Path; + use ::object_store::ObjectMeta; + use datafusion_physical_plan::collect; + use tempfile::TempDir; + + #[tokio::test] + async fn can_override_schema_adapter() { + // Test shows that SchemaAdapter can add a column that doesn't existing in the + // record batches returned from parquet. This can be useful for schema evolution + // where older files may not have all columns. + + use datafusion_execution::object_store::ObjectStoreUrl; + let tmp_dir = TempDir::new().unwrap(); + let table_dir = tmp_dir.path().join("parquet_test"); + fs::DirBuilder::new().create(table_dir.as_path()).unwrap(); + let f1 = Field::new("id", DataType::Int32, true); + + let file_schema = Arc::new(Schema::new(vec![f1.clone()])); + let filename = "part.parquet".to_string(); + let path = table_dir.as_path().join(filename.clone()); + let file = fs::File::create(path.clone()).unwrap(); + let mut writer = + parquet::arrow::ArrowWriter::try_new(file, file_schema.clone(), None) + .unwrap(); + + let ids = Arc::new(Int32Array::from(vec![1i32])); + let rec_batch = RecordBatch::try_new(file_schema.clone(), vec![ids]).unwrap(); + + writer.write(&rec_batch).unwrap(); + writer.close().unwrap(); + + let location = Path::parse(path.to_str().unwrap()).unwrap(); + let metadata = fs::metadata(path.as_path()).expect("Local file metadata"); + let meta = ObjectMeta { + location, + last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), + size: metadata.len() as usize, + e_tag: None, + version: None, + }; + + let partitioned_file = PartitionedFile { + object_meta: meta, + partition_values: vec![], + range: None, + statistics: None, + extensions: None, + metadata_size_hint: None, + }; + + let f1 = Field::new("id", DataType::Int32, true); + let f2 = Field::new("extra_column", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![f1.clone(), f2.clone()])); + let source = Arc::new( + ParquetSource::default() + .with_schema_adapter_factory(Arc::new(TestSchemaAdapterFactory {})), + ); + let base_conf = + FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, source) + .with_file(partitioned_file); + + let parquet_exec = base_conf.build(); + + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + let read = collect(parquet_exec, task_ctx).await.unwrap(); + + let expected = [ + "+----+--------------+", + "| id | extra_column |", + "+----+--------------+", + "| 1 | foo |", + "+----+--------------+", + ]; + + assert_batches_sorted_eq!(expected, &read); + } + + #[test] + fn default_schema_adapter() { + let table_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]); + + // file has a subset of the table schema fields and different type + let file_schema = Schema::new(vec![ + Field::new("c", DataType::Float64, true), // not in table schema + Field::new("b", DataType::Float64, true), + ]); + + let adapter = DefaultSchemaAdapterFactory::from_schema(Arc::new(table_schema)); + let (mapper, indices) = adapter.map_schema(&file_schema).unwrap(); + assert_eq!(indices, vec![1]); + + let file_batch = record_batch!(("b", Float64, vec![1.0, 2.0])).unwrap(); + + let mapped_batch = mapper.map_batch(file_batch).unwrap(); + + // the mapped batch has the correct schema and the "b" column has been cast to Utf8 + let expected_batch = record_batch!( + ("a", Int32, vec![None, None]), // missing column filled with nulls + ("b", Utf8, vec!["1.0", "2.0"]) // b was cast to string and order was changed + ) + .unwrap(); + assert_eq!(mapped_batch, expected_batch); + } + + #[test] + fn default_schema_adapter_non_nullable_columns() { + let table_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), // "a"" is declared non nullable + Field::new("b", DataType::Utf8, true), + ]); + let file_schema = Schema::new(vec![ + // since file doesn't have "a" it will be filled with nulls + Field::new("b", DataType::Float64, true), + ]); + + let adapter = DefaultSchemaAdapterFactory::from_schema(Arc::new(table_schema)); + let (mapper, indices) = adapter.map_schema(&file_schema).unwrap(); + assert_eq!(indices, vec![0]); + + let file_batch = record_batch!(("b", Float64, vec![1.0, 2.0])).unwrap(); + + // Mapping fails because it tries to fill in a non-nullable column with nulls + let err = mapper.map_batch(file_batch).unwrap_err().to_string(); + assert!(err.contains("Invalid argument error: Column 'a' is declared as non-nullable but contains null values"), "{err}"); + } + + #[derive(Debug)] + struct TestSchemaAdapterFactory; + + impl SchemaAdapterFactory for TestSchemaAdapterFactory { + fn create( + &self, + projected_table_schema: SchemaRef, + _table_schema: SchemaRef, + ) -> Box { + Box::new(TestSchemaAdapter { + table_schema: projected_table_schema, + }) + } + } + + struct TestSchemaAdapter { + /// Schema for the table + table_schema: SchemaRef, + } + + impl SchemaAdapter for TestSchemaAdapter { + fn map_column_index(&self, index: usize, file_schema: &Schema) -> Option { + let field = self.table_schema.field(index); + Some(file_schema.fields.find(field.name())?.0) + } + + fn map_schema( + &self, + file_schema: &Schema, + ) -> datafusion_common::Result<(Arc, Vec)> { + let mut projection = Vec::with_capacity(file_schema.fields().len()); + + for (file_idx, file_field) in file_schema.fields.iter().enumerate() { + if self.table_schema.fields().find(file_field.name()).is_some() { + projection.push(file_idx); + } + } + + Ok((Arc::new(TestSchemaMapping {}), projection)) + } + } + + #[derive(Debug)] + struct TestSchemaMapping {} + + impl SchemaMapper for TestSchemaMapping { + fn map_batch( + &self, + batch: RecordBatch, + ) -> datafusion_common::Result { + let f1 = Field::new("id", DataType::Int32, true); + let f2 = Field::new("extra_column", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![f1, f2])); + + let extra_column = Arc::new(StringArray::from(vec!["foo"])); + let mut new_columns = batch.columns().to_vec(); + new_columns.push(extra_column); + + Ok(RecordBatch::try_new(schema, new_columns).unwrap()) + } + + fn map_partial_batch( + &self, + batch: RecordBatch, + ) -> datafusion_common::Result { + self.map_batch(batch) + } + } +} diff --git a/datafusion/core/src/datasource/physical_plan/arrow_file.rs b/datafusion/core/src/datasource/physical_plan/arrow_file.rs index a0e1135e2cac..59860e3594f5 100644 --- a/datafusion/core/src/datasource/physical_plan/arrow_file.rs +++ b/datafusion/core/src/datasource/physical_plan/arrow_file.rs @@ -21,9 +21,7 @@ use std::any::Any; use std::sync::Arc; use crate::datasource::listing::PartitionedFile; -use crate::datasource::physical_plan::{ - FileMeta, FileOpenFuture, FileOpener, JsonSource, -}; +use crate::datasource::physical_plan::{FileMeta, FileOpenFuture, FileOpener}; use crate::error::Result; use arrow::buffer::Buffer; @@ -34,6 +32,7 @@ use datafusion_common::{Constraints, Statistics}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::source::DataSourceExec; +use datafusion_datasource_json::source::JsonSource; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs index 08c22183302b..9fa2b3bc1482 100644 --- a/datafusion/core/src/datasource/physical_plan/avro.rs +++ b/datafusion/core/src/datasource/physical_plan/avro.rs @@ -15,350 +15,28 @@ // specific language governing permissions and limitations // under the License. -//! Execution plan for reading line-delimited Avro files - -use std::any::Any; -use std::fmt::Formatter; -use std::sync::Arc; - -use super::FileOpener; -#[cfg(feature = "avro")] -use crate::datasource::avro_to_arrow::Reader as AvroReader; - -use crate::error::Result; - -use arrow::datatypes::SchemaRef; -use datafusion_common::{Constraints, Statistics}; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::source::DataSourceExec; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, -}; - -use object_store::ObjectStore; - -/// Execution plan for scanning Avro data source -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] -pub struct AvroExec { - inner: DataSourceExec, - base_config: FileScanConfig, -} - -#[allow(unused, deprecated)] -impl AvroExec { - /// Create a new Avro reader execution plan provided base configurations - pub fn new(base_config: FileScanConfig) -> Self { - let ( - projected_schema, - projected_constraints, - projected_statistics, - projected_output_ordering, - ) = base_config.project(); - let cache = Self::compute_properties( - Arc::clone(&projected_schema), - &projected_output_ordering, - projected_constraints, - &base_config, - ); - let base_config = base_config.with_source(Arc::new(AvroSource::default())); - Self { - inner: DataSourceExec::new(Arc::new(base_config.clone())), - base_config, - } - } - - /// Ref to the base configs - pub fn base_config(&self) -> &FileScanConfig { - &self.base_config - } - - /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( - schema: SchemaRef, - orderings: &[LexOrdering], - constraints: Constraints, - file_scan_config: &FileScanConfig, - ) -> PlanProperties { - // Equivalence Properties - let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) - .with_constraints(constraints); - let n_partitions = file_scan_config.file_groups.len(); - - PlanProperties::new( - eq_properties, - Partitioning::UnknownPartitioning(n_partitions), // Output Partitioning - EmissionType::Incremental, - Boundedness::Bounded, - ) - } -} - -#[allow(unused, deprecated)] -impl DisplayAs for AvroExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - self.inner.fmt_as(t, f) - } -} - -#[allow(unused, deprecated)] -impl ExecutionPlan for AvroExec { - fn name(&self) -> &'static str { - "AvroExec" - } - - fn as_any(&self) -> &dyn Any { - self - } +//! Reexports the [`datafusion_datasource_json::source`] module, containing [Avro] based [`FileSource`]. +//! +//! [Avro]: https://avro.apache.org/ +//! [`FileSource`]: datafusion_datasource::file::FileSource - fn properties(&self) -> &PlanProperties { - self.inner.properties() - } - fn children(&self) -> Vec<&Arc> { - Vec::new() - } - fn with_new_children( - self: Arc, - _: Vec>, - ) -> Result> { - Ok(self) - } - #[cfg(not(feature = "avro"))] - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - Err(crate::error::DataFusionError::NotImplemented( - "Cannot execute avro plan without avro feature enabled".to_string(), - )) - } - #[cfg(feature = "avro")] - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.inner.execute(partition, context) - } - - fn statistics(&self) -> Result { - self.inner.statistics() - } - - fn metrics(&self) -> Option { - self.inner.metrics() - } - - fn fetch(&self) -> Option { - self.inner.fetch() - } - - fn with_fetch(&self, limit: Option) -> Option> { - self.inner.with_fetch(limit) - } -} - -/// AvroSource holds the extra configuration that is necessary for opening avro files -#[derive(Clone, Default)] -pub struct AvroSource { - schema: Option, - batch_size: Option, - projection: Option>, - metrics: ExecutionPlanMetricsSet, - projected_statistics: Option, -} - -impl AvroSource { - /// Initialize an AvroSource with default values - pub fn new() -> Self { - Self::default() - } - - #[cfg(feature = "avro")] - fn open(&self, reader: R) -> Result> { - AvroReader::try_new( - reader, - Arc::clone(self.schema.as_ref().expect("Schema must set before open")), - self.batch_size.expect("Batch size must set before open"), - self.projection.clone(), - ) - } -} - -impl FileSource for AvroSource { - #[cfg(feature = "avro")] - fn create_file_opener( - &self, - object_store: Arc, - _base_config: &FileScanConfig, - _partition: usize, - ) -> Arc { - Arc::new(private::AvroOpener { - config: Arc::new(self.clone()), - object_store, - }) - } - - #[cfg(not(feature = "avro"))] - fn create_file_opener( - &self, - _object_store: Arc, - _base_config: &FileScanConfig, - _partition: usize, - ) -> Arc { - panic!("Avro feature is not enabled in this build") - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - - fn with_schema(&self, schema: SchemaRef) -> Arc { - let mut conf = self.clone(); - conf.schema = Some(schema); - Arc::new(conf) - } - fn with_statistics(&self, statistics: Statistics) -> Arc { - let mut conf = self.clone(); - conf.projected_statistics = Some(statistics); - Arc::new(conf) - } - - fn with_projection(&self, config: &FileScanConfig) -> Arc { - let mut conf = self.clone(); - conf.projection = config.projected_file_column_names(); - Arc::new(conf) - } - - fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self.metrics - } - - fn statistics(&self) -> Result { - let statistics = &self.projected_statistics; - Ok(statistics - .clone() - .expect("projected_statistics must be set")) - } - - fn file_type(&self) -> &str { - "avro" - } - - fn repartitioned( - &self, - _target_partitions: usize, - _repartition_file_min_size: usize, - _output_ordering: Option, - _config: &FileScanConfig, - ) -> Result> { - Ok(None) - } -} - -#[cfg(feature = "avro")] -mod private { - use super::*; - use crate::datasource::physical_plan::FileMeta; - use crate::datasource::physical_plan::{FileOpenFuture, FileOpener}; - - use bytes::Buf; - use futures::StreamExt; - use object_store::{GetResultPayload, ObjectStore}; - - pub struct DeprecatedAvroConfig { - pub schema: SchemaRef, - pub batch_size: usize, - pub projection: Option>, - pub object_store: Arc, - } - - impl DeprecatedAvroConfig { - fn open(&self, reader: R) -> Result> { - AvroReader::try_new( - reader, - Arc::clone(&self.schema), - self.batch_size, - self.projection.clone(), - ) - } - } - - pub struct DeprecatedAvroOpener { - pub config: Arc, - } - impl FileOpener for DeprecatedAvroOpener { - fn open(&self, file_meta: FileMeta) -> Result { - let config = Arc::clone(&self.config); - Ok(Box::pin(async move { - let r = config.object_store.get(file_meta.location()).await?; - match r.payload { - GetResultPayload::File(file, _) => { - let reader = config.open(file)?; - Ok(futures::stream::iter(reader).boxed()) - } - GetResultPayload::Stream(_) => { - let bytes = r.bytes().await?; - let reader = config.open(bytes.reader())?; - Ok(futures::stream::iter(reader).boxed()) - } - } - })) - } - } - - pub struct AvroOpener { - pub config: Arc, - pub object_store: Arc, - } - - impl FileOpener for AvroOpener { - fn open(&self, file_meta: FileMeta) -> Result { - let config = Arc::clone(&self.config); - let object_store = Arc::clone(&self.object_store); - Ok(Box::pin(async move { - let r = object_store.get(file_meta.location()).await?; - match r.payload { - GetResultPayload::File(file, _) => { - let reader = config.open(file)?; - Ok(futures::stream::iter(reader).boxed()) - } - GetResultPayload::Stream(_) => { - let bytes = r.bytes().await?; - let reader = config.open(bytes.reader())?; - Ok(futures::stream::iter(reader).boxed()) - } - } - })) - } - } -} +pub use datafusion_datasource_avro::source::*; #[cfg(test)] -#[cfg(feature = "avro")] mod tests { - use super::*; - use crate::arrow::datatypes::{DataType, Field, SchemaBuilder}; - use crate::datasource::file_format::{avro::AvroFormat, FileFormat}; - use crate::datasource::listing::PartitionedFile; - use crate::datasource::object_store::ObjectStoreUrl; + + use std::sync::Arc; + use crate::prelude::SessionContext; - use crate::scalar::ScalarValue; use crate::test::object_store::local_unpartitioned_file; - + use arrow::datatypes::{DataType, Field, SchemaBuilder}; + use datafusion_common::{assert_batches_eq, test_util, Result, ScalarValue}; + use datafusion_datasource::file_format::FileFormat; + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource_avro::source::AvroSource; + use datafusion_datasource_avro::AvroFormat; + use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::ExecutionPlan; use futures::StreamExt; @@ -392,7 +70,7 @@ mod tests { let url = Url::parse("file://").unwrap(); session_ctx.register_object_store(&url, store.clone()); - let testdata = crate::test_util::arrow_test_data(); + let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); let meta = local_unpartitioned_file(filename); @@ -439,7 +117,7 @@ mod tests { "+----+----------+-------------+", ]; - crate::assert_batches_eq!(expected, &[batch]); + assert_batches_eq!(expected, &[batch]); let batch = results.next().await; assert!(batch.is_none()); @@ -458,7 +136,7 @@ mod tests { let session_ctx = SessionContext::new(); let state = session_ctx.state(); - let testdata = crate::test_util::arrow_test_data(); + let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); let object_store = Arc::new(LocalFileSystem::new()) as _; let object_store_url = ObjectStoreUrl::local_filesystem(); @@ -513,7 +191,7 @@ mod tests { "+----+----------+-------------+-------------+", ]; - crate::assert_batches_eq!(expected, &[batch]); + assert_batches_eq!(expected, &[batch]); let batch = results.next().await; assert!(batch.is_none()); @@ -532,7 +210,7 @@ mod tests { let session_ctx = SessionContext::new(); let state = session_ctx.state(); - let testdata = crate::test_util::arrow_test_data(); + let testdata = test_util::arrow_test_data(); let filename = format!("{testdata}/avro/alltypes_plain.avro"); let object_store = Arc::new(LocalFileSystem::new()) as _; let object_store_url = ObjectStoreUrl::local_filesystem(); @@ -587,7 +265,7 @@ mod tests { "| 1 | false | 2021-10-26 | 1 |", "+----+----------+------------+-------------+", ]; - crate::assert_batches_eq!(expected, &[batch]); + assert_batches_eq!(expected, &[batch]); let batch = results.next().await; assert!(batch.is_none()); diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index bc7d6779bbfd..e80d04fe4b2f 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -15,794 +15,68 @@ // specific language governing permissions and limitations // under the License. -//! Execution plan for reading CSV files - -use std::any::Any; -use std::fmt; -use std::io::{Read, Seek, SeekFrom}; -use std::sync::Arc; -use std::task::Poll; - -use super::{calculate_range, RangeCalculation}; - -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::{deserialize_stream, DecoderDeserializer}; -use crate::datasource::listing::{FileRange, ListingTableUrl, PartitionedFile}; -use crate::datasource::physical_plan::FileMeta; -use crate::datasource::physical_plan::{FileOpenFuture, FileOpener}; -use crate::error::{DataFusionError, Result}; -use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; - -use arrow::csv; -use arrow::datatypes::SchemaRef; -use datafusion_common::config::ConfigOptions; -use datafusion_common::{Constraints, Statistics}; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::source::DataSourceExec; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::projection::ProjectionExec; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; - -use futures::{StreamExt, TryStreamExt}; -use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; -use tokio::io::AsyncWriteExt; -use tokio::task::JoinSet; - -/// Old Csv source, deprecated with DataSourceExec implementation and CsvSource -/// -/// See examples on `CsvSource` -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] -pub struct CsvExec { - base_config: FileScanConfig, - inner: DataSourceExec, -} - -/// Builder for [`CsvExec`]. -/// -/// See example on [`CsvExec`]. -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "use FileScanConfig instead")] -pub struct CsvExecBuilder { - file_scan_config: FileScanConfig, - file_compression_type: FileCompressionType, - // TODO: it seems like these format options could be reused across all the various CSV config - has_header: bool, - delimiter: u8, - quote: u8, - terminator: Option, - escape: Option, - comment: Option, - newlines_in_values: bool, -} - -#[allow(unused, deprecated)] -impl CsvExecBuilder { - /// Create a new builder to read the provided file scan configuration. - pub fn new(file_scan_config: FileScanConfig) -> Self { - Self { - file_scan_config, - // TODO: these defaults are duplicated from `CsvOptions` - should they be computed? - has_header: false, - delimiter: b',', - quote: b'"', - terminator: None, - escape: None, - comment: None, - newlines_in_values: false, - file_compression_type: FileCompressionType::UNCOMPRESSED, - } - } - - /// Set whether the first row defines the column names. - /// - /// The default value is `false`. - pub fn with_has_header(mut self, has_header: bool) -> Self { - self.has_header = has_header; - self - } - - /// Set the column delimeter. - /// - /// The default is `,`. - pub fn with_delimeter(mut self, delimiter: u8) -> Self { - self.delimiter = delimiter; - self - } - - /// Set the quote character. - /// - /// The default is `"`. - pub fn with_quote(mut self, quote: u8) -> Self { - self.quote = quote; - self - } - - /// Set the line terminator. If not set, the default is CRLF. - /// - /// The default is None. - pub fn with_terminator(mut self, terminator: Option) -> Self { - self.terminator = terminator; - self - } - - /// Set the escape character. - /// - /// The default is `None` (i.e. quotes cannot be escaped). - pub fn with_escape(mut self, escape: Option) -> Self { - self.escape = escape; - self - } - - /// Set the comment character. - /// - /// The default is `None` (i.e. comments are not supported). - pub fn with_comment(mut self, comment: Option) -> Self { - self.comment = comment; - self - } - - /// Set whether newlines in (quoted) values are supported. - /// - /// Parsing newlines in quoted values may be affected by execution behaviour such as - /// parallel file scanning. Setting this to `true` ensures that newlines in values are - /// parsed successfully, which may reduce performance. - /// - /// The default value is `false`. - pub fn with_newlines_in_values(mut self, newlines_in_values: bool) -> Self { - self.newlines_in_values = newlines_in_values; - self - } - - /// Set the file compression type. - /// - /// The default is [`FileCompressionType::UNCOMPRESSED`]. - pub fn with_file_compression_type( - mut self, - file_compression_type: FileCompressionType, - ) -> Self { - self.file_compression_type = file_compression_type; - self - } +//! Reexports the [`datafusion_datasource_json::source`] module, containing CSV based [`FileSource`]. +//! +//! [`FileSource`]: datafusion_datasource::file::FileSource - /// Build a [`CsvExec`]. - #[must_use] - pub fn build(self) -> CsvExec { - let Self { - file_scan_config: base_config, - file_compression_type, - has_header, - delimiter, - quote, - terminator, - escape, - comment, - newlines_in_values, - } = self; - - let ( - projected_schema, - projected_constraints, - projected_statistics, - projected_output_ordering, - ) = base_config.project(); - let cache = CsvExec::compute_properties( - projected_schema, - &projected_output_ordering, - projected_constraints, - &base_config, - ); - let csv = CsvSource::new(has_header, delimiter, quote) - .with_comment(comment) - .with_escape(escape) - .with_terminator(terminator); - let base_config = base_config - .with_newlines_in_values(newlines_in_values) - .with_file_compression_type(file_compression_type) - .with_source(Arc::new(csv)); - - CsvExec { - inner: DataSourceExec::new(Arc::new(base_config.clone())), - base_config, - } - } -} - -#[allow(unused, deprecated)] -impl CsvExec { - /// Create a new CSV reader execution plan provided base and specific configurations - #[allow(clippy::too_many_arguments)] - pub fn new( - base_config: FileScanConfig, - has_header: bool, - delimiter: u8, - quote: u8, - terminator: Option, - escape: Option, - comment: Option, - newlines_in_values: bool, - file_compression_type: FileCompressionType, - ) -> Self { - CsvExecBuilder::new(base_config) - .with_has_header(has_header) - .with_delimeter(delimiter) - .with_quote(quote) - .with_terminator(terminator) - .with_escape(escape) - .with_comment(comment) - .with_newlines_in_values(newlines_in_values) - .with_file_compression_type(file_compression_type) - .build() - } - - /// Return a [`CsvExecBuilder`]. - /// - /// See example on [`CsvExec`] and [`CsvExecBuilder`] for specifying CSV table options. - pub fn builder(file_scan_config: FileScanConfig) -> CsvExecBuilder { - CsvExecBuilder::new(file_scan_config) - } - - /// Ref to the base configs - pub fn base_config(&self) -> &FileScanConfig { - &self.base_config - } - - fn file_scan_config(&self) -> FileScanConfig { - self.inner - .data_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - fn csv_source(&self) -> CsvSource { - let source = self.file_scan_config(); - source - .file_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - /// true if the first line of each file is a header - pub fn has_header(&self) -> bool { - self.csv_source().has_header() - } - - /// Specifies whether newlines in (quoted) values are supported. - /// - /// Parsing newlines in quoted values may be affected by execution behaviour such as - /// parallel file scanning. Setting this to `true` ensures that newlines in values are - /// parsed successfully, which may reduce performance. - /// - /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting. - pub fn newlines_in_values(&self) -> bool { - let source = self.file_scan_config(); - source.newlines_in_values() - } - - fn output_partitioning_helper(file_scan_config: &FileScanConfig) -> Partitioning { - Partitioning::UnknownPartitioning(file_scan_config.file_groups.len()) - } - - /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( - schema: SchemaRef, - orderings: &[LexOrdering], - constraints: Constraints, - file_scan_config: &FileScanConfig, - ) -> PlanProperties { - // Equivalence Properties - let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) - .with_constraints(constraints); - - PlanProperties::new( - eq_properties, - Self::output_partitioning_helper(file_scan_config), // Output Partitioning - EmissionType::Incremental, - Boundedness::Bounded, - ) - } - - fn with_file_groups(mut self, file_groups: Vec>) -> Self { - self.base_config.file_groups = file_groups.clone(); - let mut file_source = self.file_scan_config(); - file_source = file_source.with_file_groups(file_groups); - self.inner = self.inner.with_data_source(Arc::new(file_source)); - self - } -} - -#[allow(unused, deprecated)] -impl DisplayAs for CsvExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - self.inner.fmt_as(t, f) - } -} - -#[allow(unused, deprecated)] -impl ExecutionPlan for CsvExec { - fn name(&self) -> &'static str { - "CsvExec" - } - - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - // this is a leaf node and has no children - vec![] - } - - fn with_new_children( - self: Arc, - _: Vec>, - ) -> Result> { - Ok(self) - } - - /// Redistribute files across partitions according to their size - /// See comments on `FileGroupPartitioner` for more detail. - /// - /// Return `None` if can't get repartitioned (empty, compressed file, or `newlines_in_values` set). - fn repartitioned( - &self, - target_partitions: usize, - config: &ConfigOptions, - ) -> Result>> { - self.inner.repartitioned(target_partitions, config) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.inner.execute(partition, context) - } - - fn statistics(&self) -> Result { - self.inner.statistics() - } - - fn metrics(&self) -> Option { - self.inner.metrics() - } - - fn fetch(&self) -> Option { - self.inner.fetch() - } - - fn with_fetch(&self, limit: Option) -> Option> { - self.inner.with_fetch(limit) - } - - fn try_swapping_with_projection( - &self, - projection: &ProjectionExec, - ) -> Result>> { - self.inner.try_swapping_with_projection(projection) - } -} - -/// A Config for [`CsvOpener`] -/// -/// # Example: create a `DataSourceExec` for CSV -/// ``` -/// # use std::sync::Arc; -/// # use arrow::datatypes::Schema; -/// # use datafusion::datasource::{ -/// # physical_plan::FileScanConfig, -/// # listing::PartitionedFile, -/// # }; -/// # use datafusion::datasource::physical_plan::CsvSource; -/// # use datafusion_execution::object_store::ObjectStoreUrl; -/// # use datafusion::datasource::source::DataSourceExec; -/// -/// # let object_store_url = ObjectStoreUrl::local_filesystem(); -/// # let file_schema = Arc::new(Schema::empty()); -/// -/// let source = Arc::new(CsvSource::new( -/// true, -/// b',', -/// b'"', -/// ) -/// .with_terminator(Some(b'#') -/// )); -/// // Create a DataSourceExec for reading the first 100MB of `file1.csv` -/// let file_scan_config = FileScanConfig::new(object_store_url, file_schema, source) -/// .with_file(PartitionedFile::new("file1.csv", 100*1024*1024)) -/// .with_newlines_in_values(true); // The file contains newlines in values; -/// let exec = file_scan_config.build(); -/// ``` -#[derive(Debug, Clone, Default)] -pub struct CsvSource { - batch_size: Option, - file_schema: Option, - file_projection: Option>, - pub(crate) has_header: bool, - delimiter: u8, - quote: u8, - terminator: Option, - escape: Option, - comment: Option, - metrics: ExecutionPlanMetricsSet, - projected_statistics: Option, -} - -impl CsvSource { - /// Returns a [`CsvSource`] - pub fn new(has_header: bool, delimiter: u8, quote: u8) -> Self { - Self { - has_header, - delimiter, - quote, - ..Self::default() - } - } - - /// true if the first line of each file is a header - pub fn has_header(&self) -> bool { - self.has_header - } - /// A column delimiter - pub fn delimiter(&self) -> u8 { - self.delimiter - } - - /// The quote character - pub fn quote(&self) -> u8 { - self.quote - } - - /// The line terminator - pub fn terminator(&self) -> Option { - self.terminator - } - - /// Lines beginning with this byte are ignored. - pub fn comment(&self) -> Option { - self.comment - } - - /// The escape character - pub fn escape(&self) -> Option { - self.escape - } - - /// Initialize a CsvSource with escape - pub fn with_escape(&self, escape: Option) -> Self { - let mut conf = self.clone(); - conf.escape = escape; - conf - } - - /// Initialize a CsvSource with terminator - pub fn with_terminator(&self, terminator: Option) -> Self { - let mut conf = self.clone(); - conf.terminator = terminator; - conf - } - - /// Initialize a CsvSource with comment - pub fn with_comment(&self, comment: Option) -> Self { - let mut conf = self.clone(); - conf.comment = comment; - conf - } -} - -impl CsvSource { - fn open(&self, reader: R) -> Result> { - Ok(self.builder().build(reader)?) - } - - fn builder(&self) -> csv::ReaderBuilder { - let mut builder = csv::ReaderBuilder::new(Arc::clone( - self.file_schema - .as_ref() - .expect("Schema must be set before initializing builder"), - )) - .with_delimiter(self.delimiter) - .with_batch_size( - self.batch_size - .expect("Batch size must be set before initializing builder"), - ) - .with_header(self.has_header) - .with_quote(self.quote); - if let Some(terminator) = self.terminator { - builder = builder.with_terminator(terminator); - } - if let Some(proj) = &self.file_projection { - builder = builder.with_projection(proj.clone()); - } - if let Some(escape) = self.escape { - builder = builder.with_escape(escape) - } - if let Some(comment) = self.comment { - builder = builder.with_comment(comment); - } - - builder - } -} - -/// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`] -pub struct CsvOpener { - config: Arc, - file_compression_type: FileCompressionType, - object_store: Arc, -} - -impl CsvOpener { - /// Returns a [`CsvOpener`] - pub fn new( - config: Arc, - file_compression_type: FileCompressionType, - object_store: Arc, - ) -> Self { - Self { - config, - file_compression_type, - object_store, - } - } -} - -impl FileSource for CsvSource { - fn create_file_opener( - &self, - object_store: Arc, - base_config: &FileScanConfig, - _partition: usize, - ) -> Arc { - Arc::new(CsvOpener { - config: Arc::new(self.clone()), - file_compression_type: base_config.file_compression_type, - object_store, - }) - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - - fn with_schema(&self, schema: SchemaRef) -> Arc { - let mut conf = self.clone(); - conf.file_schema = Some(schema); - Arc::new(conf) - } - - fn with_statistics(&self, statistics: Statistics) -> Arc { - let mut conf = self.clone(); - conf.projected_statistics = Some(statistics); - Arc::new(conf) - } - - fn with_projection(&self, config: &FileScanConfig) -> Arc { - let mut conf = self.clone(); - conf.file_projection = config.file_column_projection_indices(); - Arc::new(conf) - } - - fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self.metrics - } - fn statistics(&self) -> Result { - let statistics = &self.projected_statistics; - Ok(statistics - .clone() - .expect("projected_statistics must be set")) - } - fn file_type(&self) -> &str { - "csv" - } - fn fmt_extra(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, ", has_header={}", self.has_header) - } -} - -impl FileOpener for CsvOpener { - /// Open a partitioned CSV file. - /// - /// If `file_meta.range` is `None`, the entire file is opened. - /// If `file_meta.range` is `Some(FileRange {start, end})`, this signifies that the partition - /// corresponds to the byte range [start, end) within the file. - /// - /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules - /// are applied to determine which lines to read: - /// 1. The first line of the partition is the line in which the index of the first character >= `start`. - /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides. - /// - /// Examples: - /// Consider the following partitions enclosed by braces `{}`: - /// - /// {A,1,2,3,4,5,6,7,8,9\n - /// A,1,2,3,4,5,6,7,8,9\n} - /// A,1,2,3,4,5,6,7,8,9\n - /// The lines read would be: [0, 1] - /// - /// A,{1,2,3,4,5,6,7,8,9\n - /// A,1,2,3,4,5,6,7,8,9\n - /// A},1,2,3,4,5,6,7,8,9\n - /// The lines read would be: [1, 2] - fn open(&self, file_meta: FileMeta) -> Result { - // `self.config.has_header` controls whether to skip reading the 1st line header - // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle - // partition, then don't skip first line - let mut csv_has_header = self.config.has_header; - if let Some(FileRange { start, .. }) = file_meta.range { - if start != 0 { - csv_has_header = false; - } - } - - let config = CsvSource { - has_header: csv_has_header, - ..(*self.config).clone() - }; - - let file_compression_type = self.file_compression_type.to_owned(); - - if file_meta.range.is_some() { - assert!( - !file_compression_type.is_compressed(), - "Reading compressed .csv in parallel is not supported" - ); - } - - let store = Arc::clone(&self.object_store); - let terminator = self.config.terminator; - - Ok(Box::pin(async move { - // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries) - - let calculated_range = - calculate_range(&file_meta, &store, terminator).await?; - - let range = match calculated_range { - RangeCalculation::Range(None) => None, - RangeCalculation::Range(Some(range)) => Some(range.into()), - RangeCalculation::TerminateEarly => { - return Ok( - futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed() - ) - } - }; - - let options = GetOptions { - range, - ..Default::default() - }; - - let result = store.get_opts(file_meta.location(), options).await?; - - match result.payload { - GetResultPayload::File(mut file, _) => { - let is_whole_file_scanned = file_meta.range.is_none(); - let decoder = if is_whole_file_scanned { - // Don't seek if no range as breaks FIFO files - file_compression_type.convert_read(file)? - } else { - file.seek(SeekFrom::Start(result.range.start as _))?; - file_compression_type.convert_read( - file.take((result.range.end - result.range.start) as u64), - )? - }; - - Ok(futures::stream::iter(config.open(decoder)?).boxed()) - } - GetResultPayload::Stream(s) => { - let decoder = config.builder().build_decoder(); - let s = s.map_err(DataFusionError::from); - let input = file_compression_type.convert_stream(s.boxed())?.fuse(); - - Ok(deserialize_stream( - input, - DecoderDeserializer::from(decoder), - )) - } - } - })) - } -} - -pub async fn plan_to_csv( - task_ctx: Arc, - plan: Arc, - path: impl AsRef, -) -> Result<()> { - let path = path.as_ref(); - let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; - let mut join_set = JoinSet::new(); - for i in 0..plan.output_partitioning().partition_count() { - let storeref = Arc::clone(&store); - let plan: Arc = Arc::clone(&plan); - let filename = format!("{}/part-{i}.csv", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; - - let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; - join_set.spawn(async move { - let mut buf_writer = BufWriter::new(storeref, file.clone()); - let mut buffer = Vec::with_capacity(1024); - //only write headers on first iteration - let mut write_headers = true; - while let Some(batch) = stream.next().await.transpose()? { - let mut writer = csv::WriterBuilder::new() - .with_header(write_headers) - .build(buffer); - writer.write(&batch)?; - buffer = writer.into_inner(); - buf_writer.write_all(&buffer).await?; - buffer.clear(); - //prevent writing headers more than once - write_headers = false; - } - buf_writer.shutdown().await.map_err(DataFusionError::from) - }); - } - - while let Some(result) = join_set.join_next().await { - match result { - Ok(res) => res?, // propagate DataFusion error - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } - - Ok(()) -} +pub use datafusion_datasource_csv::source::*; #[cfg(test)] mod tests { + + use std::collections::HashMap; use std::fs::{self, File}; use std::io::Write; + use std::sync::Arc; - use super::*; - use crate::dataframe::DataFrameWriteOptions; - use crate::datasource::file_format::csv::CsvFormat; - use crate::prelude::*; - use crate::test::{partitioned_csv_config, partitioned_file_groups}; - use crate::{scalar::ScalarValue, test_util::aggr_test_schema}; + use datafusion_datasource_csv::CsvFormat; + use object_store::ObjectStore; - use arrow::datatypes::*; - use bytes::Bytes; + use crate::prelude::CsvReadOptions; + use crate::prelude::SessionContext; + use crate::test::partitioned_file_groups; use datafusion_common::test_util::arrow_test_data; + use datafusion_common::{assert_batches_eq, Result}; + use datafusion_execution::config::SessionConfig; use datafusion_physical_plan::metrics::MetricsSet; + use datafusion_physical_plan::ExecutionPlan; + #[cfg(feature = "compression")] + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource_csv::partitioned_csv_config; + use datafusion_datasource_csv::source::CsvSource; + use futures::{StreamExt, TryStreamExt}; + + use arrow::datatypes::*; + use bytes::Bytes; use object_store::chunked::ChunkedStore; use object_store::local::LocalFileSystem; use rstest::*; use tempfile::TempDir; use url::Url; + fn aggr_test_schema() -> SchemaRef { + let mut f1 = Field::new("c1", DataType::Utf8, false); + f1.set_metadata(HashMap::from_iter(vec![("testing".into(), "test".into())])); + let schema = Schema::new(vec![ + f1, + Field::new("c2", DataType::UInt32, false), + Field::new("c3", DataType::Int8, false), + Field::new("c4", DataType::Int16, false), + Field::new("c5", DataType::Int32, false), + Field::new("c6", DataType::Int64, false), + Field::new("c7", DataType::UInt8, false), + Field::new("c8", DataType::UInt16, false), + Field::new("c9", DataType::UInt32, false), + Field::new("c10", DataType::UInt64, false), + Field::new("c11", DataType::Float32, false), + Field::new("c12", DataType::Float64, false), + Field::new("c13", DataType::Utf8, false), + ]); + + Arc::new(schema) + } + #[rstest( file_compression_type, case(FileCompressionType::UNCOMPRESSED), @@ -816,8 +90,6 @@ mod tests { async fn csv_exec_with_projection( file_compression_type: FileCompressionType, ) -> Result<()> { - use crate::datasource::file_format::csv::CsvFormat; - let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); let file_schema = aggr_test_schema(); @@ -863,7 +135,7 @@ mod tests { "+----+-----+------------+", ]; - crate::assert_batches_eq!(expected, &[batch.slice(0, 5)]); + assert_batches_eq!(expected, &[batch.slice(0, 5)]); Ok(()) } @@ -880,8 +152,6 @@ mod tests { async fn csv_exec_with_mixed_order_projection( file_compression_type: FileCompressionType, ) -> Result<()> { - use crate::datasource::file_format::csv::CsvFormat; - let cfg = SessionConfig::new().set_str("datafusion.catalog.has_header", "true"); let session_ctx = SessionContext::new_with_config(cfg); let task_ctx = session_ctx.task_ctx(); @@ -926,7 +196,7 @@ mod tests { "+------------+----+-----+", ]; - crate::assert_batches_eq!(expected, &[batch.slice(0, 5)]); + assert_batches_eq!(expected, &[batch.slice(0, 5)]); Ok(()) } @@ -943,7 +213,7 @@ mod tests { async fn csv_exec_with_limit( file_compression_type: FileCompressionType, ) -> Result<()> { - use crate::datasource::file_format::csv::CsvFormat; + use futures::StreamExt; let cfg = SessionConfig::new().set_str("datafusion.catalog.has_header", "true"); let session_ctx = SessionContext::new_with_config(cfg); @@ -986,7 +256,7 @@ mod tests { "| b | 5 | -82 | 22080 | 1824882165 | 7373730676428214987 | 208 | 34331 | 3342719438 | 3330177516592499461 | 0.82634634 | 0.40975383525297016 | Ig1QcuKsjHXkproePdERo2w0mYzIqd |", "+----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+"]; - crate::assert_batches_eq!(expected, &[batch]); + assert_batches_eq!(expected, &[batch]); Ok(()) } @@ -1004,8 +274,6 @@ mod tests { async fn csv_exec_with_missing_column( file_compression_type: FileCompressionType, ) -> Result<()> { - use crate::datasource::file_format::csv::CsvFormat; - let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); let file_schema = aggr_test_schema_with_missing_col(); @@ -1054,7 +322,7 @@ mod tests { async fn csv_exec_with_partition( file_compression_type: FileCompressionType, ) -> Result<()> { - use crate::datasource::file_format::csv::CsvFormat; + use datafusion_common::ScalarValue; let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); @@ -1109,7 +377,7 @@ mod tests { "| b | 2021-10-26 |", "+----+------------+", ]; - crate::assert_batches_eq!(expected, &[batch.slice(0, 5)]); + assert_batches_eq!(expected, &[batch.slice(0, 5)]); let metrics = csv.metrics().expect("doesn't found metrics"); let time_elapsed_processing = get_value(&metrics, "time_elapsed_processing"); @@ -1242,7 +510,7 @@ mod tests { "+---+---+", ]; - crate::assert_batches_eq!(expected, &result); + assert_batches_eq!(expected, &result); } #[tokio::test] @@ -1273,7 +541,7 @@ mod tests { "+---+---+", ]; - crate::assert_batches_eq!(expected, &result); + assert_batches_eq!(expected, &result); let e = session_ctx .read_csv("memory:///", CsvReadOptions::new().terminator(Some(b'\n'))) @@ -1313,7 +581,7 @@ mod tests { "| id3 | value3 |", "+------+--------+", ]; - crate::assert_batches_eq!(expected, &df); + assert_batches_eq!(expected, &df); Ok(()) } @@ -1342,7 +610,7 @@ mod tests { "| value | end |", "+-------+-----------------------------+", ]; - crate::assert_batches_eq!(expected, &df); + assert_batches_eq!(expected, &df); Ok(()) } @@ -1362,7 +630,11 @@ mod tests { let out_dir_url = "file://local/out"; let e = df - .write_csv(out_dir_url, DataFrameWriteOptions::new(), None) + .write_csv( + out_dir_url, + crate::dataframe::DataFrameWriteOptions::new(), + None, + ) .await .expect_err("should fail because input file does not match inferred schema"); assert_eq!(e.strip_backtrace(), "Arrow error: Parser error: Error while parsing value d for column 0 at line 4"); @@ -1400,8 +672,12 @@ mod tests { let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out/"; let out_dir_url = "file://local/out/"; let df = ctx.sql("SELECT c1, c2 FROM test").await?; - df.write_csv(out_dir_url, DataFrameWriteOptions::new(), None) - .await?; + df.write_csv( + out_dir_url, + crate::dataframe::DataFrameWriteOptions::new(), + None, + ) + .await?; // create a new context and verify that the results were saved to a partitioned csv file let ctx = SessionContext::new_with_config( diff --git a/datafusion/core/src/datasource/physical_plan/json.rs b/datafusion/core/src/datasource/physical_plan/json.rs index c9a22add2afc..9bab75fc88c3 100644 --- a/datafusion/core/src/datasource/physical_plan/json.rs +++ b/datafusion/core/src/datasource/physical_plan/json.rs @@ -15,453 +15,43 @@ // specific language governing permissions and limitations // under the License. -//! Execution plan for reading line-delimited JSON files - -use std::any::Any; -use std::io::{BufReader, Read, Seek, SeekFrom}; -use std::sync::Arc; -use std::task::Poll; - -use super::{calculate_range, RangeCalculation}; - -use crate::datasource::file_format::file_compression_type::FileCompressionType; -use crate::datasource::file_format::{deserialize_stream, DecoderDeserializer}; -use crate::datasource::listing::{ListingTableUrl, PartitionedFile}; -use crate::datasource::physical_plan::FileMeta; -use crate::datasource::physical_plan::{FileOpenFuture, FileOpener}; -use crate::error::{DataFusionError, Result}; -use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; - -use arrow::json::ReaderBuilder; -use arrow::{datatypes::SchemaRef, json}; -use datafusion_common::{Constraints, Statistics}; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::source::DataSourceExec; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; - -use futures::{StreamExt, TryStreamExt}; -use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; -use tokio::io::AsyncWriteExt; -use tokio::task::JoinSet; - -/// Execution plan for scanning NdJson data source -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] -pub struct NdJsonExec { - inner: DataSourceExec, - base_config: FileScanConfig, - file_compression_type: FileCompressionType, -} - -#[allow(unused, deprecated)] -impl NdJsonExec { - /// Create a new JSON reader execution plan provided base configurations - pub fn new( - base_config: FileScanConfig, - file_compression_type: FileCompressionType, - ) -> Self { - let ( - projected_schema, - projected_constraints, - projected_statistics, - projected_output_ordering, - ) = base_config.project(); - let cache = Self::compute_properties( - projected_schema, - &projected_output_ordering, - projected_constraints, - &base_config, - ); - - let json = JsonSource::default(); - let base_config = base_config - .with_file_compression_type(file_compression_type) - .with_source(Arc::new(json)); - - Self { - inner: DataSourceExec::new(Arc::new(base_config.clone())), - file_compression_type: base_config.file_compression_type, - base_config, - } - } - - /// Ref to the base configs - pub fn base_config(&self) -> &FileScanConfig { - &self.base_config - } - - /// Ref to file compression type - pub fn file_compression_type(&self) -> &FileCompressionType { - &self.file_compression_type - } - - fn file_scan_config(&self) -> FileScanConfig { - self.inner - .data_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - fn json_source(&self) -> JsonSource { - let source = self.file_scan_config(); - source - .file_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - fn output_partitioning_helper(file_scan_config: &FileScanConfig) -> Partitioning { - Partitioning::UnknownPartitioning(file_scan_config.file_groups.len()) - } - - /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( - schema: SchemaRef, - orderings: &[LexOrdering], - constraints: Constraints, - file_scan_config: &FileScanConfig, - ) -> PlanProperties { - // Equivalence Properties - let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) - .with_constraints(constraints); - - PlanProperties::new( - eq_properties, - Self::output_partitioning_helper(file_scan_config), // Output Partitioning - EmissionType::Incremental, - Boundedness::Bounded, - ) - } - - fn with_file_groups(mut self, file_groups: Vec>) -> Self { - self.base_config.file_groups = file_groups.clone(); - let mut file_source = self.file_scan_config(); - file_source = file_source.with_file_groups(file_groups); - self.inner = self.inner.with_data_source(Arc::new(file_source)); - self - } -} - -#[allow(unused, deprecated)] -impl DisplayAs for NdJsonExec { - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - self.inner.fmt_as(t, f) - } -} - -#[allow(unused, deprecated)] -impl ExecutionPlan for NdJsonExec { - fn name(&self) -> &'static str { - "NdJsonExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &PlanProperties { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - Vec::new() - } - - fn with_new_children( - self: Arc, - _: Vec>, - ) -> Result> { - Ok(self) - } - - fn repartitioned( - &self, - target_partitions: usize, - config: &datafusion_common::config::ConfigOptions, - ) -> Result>> { - self.inner.repartitioned(target_partitions, config) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.inner.execute(partition, context) - } - - fn statistics(&self) -> Result { - self.inner.statistics() - } - - fn metrics(&self) -> Option { - self.inner.metrics() - } - - fn fetch(&self) -> Option { - self.inner.fetch() - } - - fn with_fetch(&self, limit: Option) -> Option> { - self.inner.with_fetch(limit) - } -} - -/// A [`FileOpener`] that opens a JSON file and yields a [`FileOpenFuture`] -pub struct JsonOpener { - batch_size: usize, - projected_schema: SchemaRef, - file_compression_type: FileCompressionType, - object_store: Arc, -} - -impl JsonOpener { - /// Returns a [`JsonOpener`] - pub fn new( - batch_size: usize, - projected_schema: SchemaRef, - file_compression_type: FileCompressionType, - object_store: Arc, - ) -> Self { - Self { - batch_size, - projected_schema, - file_compression_type, - object_store, - } - } -} - -/// JsonSource holds the extra configuration that is necessary for [`JsonOpener`] -#[derive(Clone, Default)] -pub struct JsonSource { - batch_size: Option, - metrics: ExecutionPlanMetricsSet, - projected_statistics: Option, -} - -impl JsonSource { - /// Initialize a JsonSource with default values - pub fn new() -> Self { - Self::default() - } -} - -impl FileSource for JsonSource { - fn create_file_opener( - &self, - object_store: Arc, - base_config: &FileScanConfig, - _partition: usize, - ) -> Arc { - Arc::new(JsonOpener { - batch_size: self - .batch_size - .expect("Batch size must set before creating opener"), - projected_schema: base_config.projected_file_schema(), - file_compression_type: base_config.file_compression_type, - object_store, - }) - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut conf = self.clone(); - conf.batch_size = Some(batch_size); - Arc::new(conf) - } - - fn with_schema(&self, _schema: SchemaRef) -> Arc { - Arc::new(Self { ..self.clone() }) - } - fn with_statistics(&self, statistics: Statistics) -> Arc { - let mut conf = self.clone(); - conf.projected_statistics = Some(statistics); - Arc::new(conf) - } - - fn with_projection(&self, _config: &FileScanConfig) -> Arc { - Arc::new(Self { ..self.clone() }) - } - - fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self.metrics - } - - fn statistics(&self) -> Result { - let statistics = &self.projected_statistics; - Ok(statistics - .clone() - .expect("projected_statistics must be set to call")) - } - - fn file_type(&self) -> &str { - "json" - } -} - -impl FileOpener for JsonOpener { - /// Open a partitioned NDJSON file. - /// - /// If `file_meta.range` is `None`, the entire file is opened. - /// Else `file_meta.range` is `Some(FileRange{start, end})`, which corresponds to the byte range [start, end) within the file. - /// - /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules - /// are applied to determine which lines to read: - /// 1. The first line of the partition is the line in which the index of the first character >= `start`. - /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides. - /// - /// See [`CsvOpener`](super::CsvOpener) for an example. - fn open(&self, file_meta: FileMeta) -> Result { - let store = Arc::clone(&self.object_store); - let schema = Arc::clone(&self.projected_schema); - let batch_size = self.batch_size; - let file_compression_type = self.file_compression_type.to_owned(); - - Ok(Box::pin(async move { - let calculated_range = calculate_range(&file_meta, &store, None).await?; - - let range = match calculated_range { - RangeCalculation::Range(None) => None, - RangeCalculation::Range(Some(range)) => Some(range.into()), - RangeCalculation::TerminateEarly => { - return Ok( - futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed() - ) - } - }; +//! Reexports the [`datafusion_datasource_json::source`] module, containing JSON based [`FileSource`]. +//! +//! [`FileSource`]: datafusion_datasource::file::FileSource - let options = GetOptions { - range, - ..Default::default() - }; - - let result = store.get_opts(file_meta.location(), options).await?; - - match result.payload { - GetResultPayload::File(mut file, _) => { - let bytes = match file_meta.range { - None => file_compression_type.convert_read(file)?, - Some(_) => { - file.seek(SeekFrom::Start(result.range.start as _))?; - let limit = result.range.end - result.range.start; - file_compression_type.convert_read(file.take(limit as u64))? - } - }; - - let reader = ReaderBuilder::new(schema) - .with_batch_size(batch_size) - .build(BufReader::new(bytes))?; - - Ok(futures::stream::iter(reader).boxed()) - } - GetResultPayload::Stream(s) => { - let s = s.map_err(DataFusionError::from); - - let decoder = ReaderBuilder::new(schema) - .with_batch_size(batch_size) - .build_decoder()?; - let input = file_compression_type.convert_stream(s.boxed())?.fuse(); - - Ok(deserialize_stream( - input, - DecoderDeserializer::from(decoder), - )) - } - } - })) - } -} - -pub async fn plan_to_json( - task_ctx: Arc, - plan: Arc, - path: impl AsRef, -) -> Result<()> { - let path = path.as_ref(); - let parsed = ListingTableUrl::parse(path)?; - let object_store_url = parsed.object_store(); - let store = task_ctx.runtime_env().object_store(&object_store_url)?; - let mut join_set = JoinSet::new(); - for i in 0..plan.output_partitioning().partition_count() { - let storeref = Arc::clone(&store); - let plan: Arc = Arc::clone(&plan); - let filename = format!("{}/part-{i}.json", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; - - let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; - join_set.spawn(async move { - let mut buf_writer = BufWriter::new(storeref, file.clone()); - - let mut buffer = Vec::with_capacity(1024); - while let Some(batch) = stream.next().await.transpose()? { - let mut writer = json::LineDelimitedWriter::new(buffer); - writer.write(&batch)?; - buffer = writer.into_inner(); - buf_writer.write_all(&buffer).await?; - buffer.clear(); - } - - buf_writer.shutdown().await.map_err(DataFusionError::from) - }); - } - - while let Some(result) = join_set.join_next().await { - match result { - Ok(res) => res?, // propagate DataFusion error - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } - - Ok(()) -} +#[allow(deprecated)] +pub use datafusion_datasource_json::source::*; #[cfg(test)] mod tests { + + use super::*; + use std::fs; use std::path::Path; + use std::sync::Arc; - use super::*; use crate::dataframe::DataFrameWriteOptions; - use crate::datasource::file_format::{json::JsonFormat, FileFormat}; - use crate::datasource::listing::PartitionedFile; - use crate::datasource::object_store::ObjectStoreUrl; - use crate::execution::context::SessionState; - use crate::prelude::{ - CsvReadOptions, NdJsonReadOptions, SessionConfig, SessionContext, - }; + use crate::execution::SessionState; + use crate::prelude::{CsvReadOptions, NdJsonReadOptions, SessionContext}; use crate::test::partitioned_file_groups; - use crate::{assert_batches_eq, assert_batches_sorted_eq}; + use datafusion_common::cast::{as_int32_array, as_int64_array, as_string_array}; + use datafusion_common::Result; + use datafusion_common::{assert_batches_eq, assert_batches_sorted_eq}; + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource::file_format::FileFormat; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource_json::JsonFormat; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_plan::ExecutionPlan; use arrow::array::Array; + use arrow::datatypes::SchemaRef; use arrow::datatypes::{Field, SchemaBuilder}; - use datafusion_common::cast::{as_int32_array, as_int64_array, as_string_array}; - use object_store::chunked::ChunkedStore; use object_store::local::LocalFileSystem; + use object_store::ObjectStore; use rstest::*; use tempfile::TempDir; use url::Url; @@ -577,6 +167,8 @@ mod tests { let state = session_ctx.state(); let task_ctx = session_ctx.task_ctx(); use arrow::datatypes::DataType; + use datafusion_datasource::file_scan_config::FileScanConfig; + use futures::StreamExt; let tmp_dir = TempDir::new()?; let (object_store_url, file_groups, file_schema) = @@ -638,10 +230,13 @@ mod tests { async fn nd_json_exec_file_with_missing_column( file_compression_type: FileCompressionType, ) -> Result<()> { + use arrow::datatypes::DataType; + use datafusion_datasource::file_scan_config::FileScanConfig; + use futures::StreamExt; + let session_ctx = SessionContext::new(); let state = session_ctx.state(); let task_ctx = session_ctx.task_ctx(); - use arrow::datatypes::DataType; let tmp_dir = TempDir::new()?; let (object_store_url, file_groups, actual_schema) = @@ -686,6 +281,9 @@ mod tests { async fn nd_json_exec_file_projection( file_compression_type: FileCompressionType, ) -> Result<()> { + use datafusion_datasource::file_scan_config::FileScanConfig; + use futures::StreamExt; + let session_ctx = SessionContext::new(); let state = session_ctx.state(); let task_ctx = session_ctx.task_ctx(); @@ -731,6 +329,9 @@ mod tests { async fn nd_json_exec_file_mixed_order_projection( file_compression_type: FileCompressionType, ) -> Result<()> { + use datafusion_datasource::file_scan_config::FileScanConfig; + use futures::StreamExt; + let session_ctx = SessionContext::new(); let state = session_ctx.state(); let task_ctx = session_ctx.task_ctx(); @@ -936,6 +537,8 @@ mod tests { async fn test_json_with_repartitioning( file_compression_type: FileCompressionType, ) -> Result<()> { + use datafusion_execution::config::SessionConfig; + let config = SessionConfig::new() .with_repartition_file_scans(true) .with_repartition_file_min_size(0) diff --git a/datafusion/core/src/datasource/physical_plan/mod.rs b/datafusion/core/src/datasource/physical_plan/mod.rs index 42f6912afec0..cae04e5ee6b8 100644 --- a/datafusion/core/src/datasource/physical_plan/mod.rs +++ b/datafusion/core/src/datasource/physical_plan/mod.rs @@ -18,32 +18,39 @@ //! Execution plans that read file formats mod arrow_file; -mod avro; -mod csv; -mod json; +pub mod csv; +pub mod json; + #[cfg(feature = "parquet")] pub mod parquet; -pub(crate) use self::csv::plan_to_csv; -pub(crate) use self::json::plan_to_json; +#[cfg(feature = "avro")] +pub mod avro; + +#[allow(deprecated)] +#[cfg(feature = "avro")] +pub use avro::{AvroExec, AvroSource}; + #[cfg(feature = "parquet")] -pub use self::parquet::source::ParquetSource; +pub use datafusion_datasource_parquet::source::ParquetSource; #[cfg(feature = "parquet")] #[allow(deprecated)] -pub use self::parquet::{ +pub use datafusion_datasource_parquet::{ ParquetExec, ParquetExecBuilder, ParquetFileMetrics, ParquetFileReaderFactory, }; -use crate::datasource::listing::FileRange; -use crate::error::Result; -use crate::physical_plan::DisplayAs; + #[allow(deprecated)] pub use arrow_file::ArrowExec; pub use arrow_file::ArrowSource; + #[allow(deprecated)] -pub use avro::AvroExec; -pub use avro::AvroSource; +pub use json::NdJsonExec; + +pub use json::{JsonOpener, JsonSource}; + #[allow(deprecated)] pub use csv::{CsvExec, CsvExecBuilder}; + pub use csv::{CsvOpener, CsvSource}; pub use datafusion_datasource::file::FileSource; pub use datafusion_datasource::file_groups::FileGroupPartitioner; @@ -56,121 +63,10 @@ pub use datafusion_datasource::file_sink_config::*; pub use datafusion_datasource::file_stream::{ FileOpenFuture, FileOpener, FileStream, OnError, }; -use futures::StreamExt; -#[allow(deprecated)] -pub use json::NdJsonExec; -pub use json::{JsonOpener, JsonSource}; - -use object_store::{path::Path, GetOptions, GetRange, ObjectStore}; -use std::{ops::Range, sync::Arc}; - -/// Represents the possible outcomes of a range calculation. -/// -/// This enum is used to encapsulate the result of calculating the range of -/// bytes to read from an object (like a file) in an object store. -/// -/// Variants: -/// - `Range(Option>)`: -/// Represents a range of bytes to be read. It contains an `Option` wrapping a -/// `Range`. `None` signifies that the entire object should be read, -/// while `Some(range)` specifies the exact byte range to read. -/// - `TerminateEarly`: -/// Indicates that the range calculation determined no further action is -/// necessary, possibly because the calculated range is empty or invalid. -enum RangeCalculation { - Range(Option>), - TerminateEarly, -} - -/// Calculates an appropriate byte range for reading from an object based on the -/// provided metadata. -/// -/// This asynchronous function examines the `FileMeta` of an object in an object store -/// and determines the range of bytes to be read. The range calculation may adjust -/// the start and end points to align with meaningful data boundaries (like newlines). -/// -/// Returns a `Result` wrapping a `RangeCalculation`, which is either a calculated byte range or an indication to terminate early. -/// -/// Returns an `Error` if any part of the range calculation fails, such as issues in reading from the object store or invalid range boundaries. -async fn calculate_range( - file_meta: &FileMeta, - store: &Arc, - terminator: Option, -) -> Result { - let location = file_meta.location(); - let file_size = file_meta.object_meta.size; - let newline = terminator.unwrap_or(b'\n'); - - match file_meta.range { - None => Ok(RangeCalculation::Range(None)), - Some(FileRange { start, end }) => { - let (start, end) = (start as usize, end as usize); - - let start_delta = if start != 0 { - find_first_newline(store, location, start - 1, file_size, newline).await? - } else { - 0 - }; - - let end_delta = if end != file_size { - find_first_newline(store, location, end - 1, file_size, newline).await? - } else { - 0 - }; - - let range = start + start_delta..end + end_delta; - - if range.start == range.end { - return Ok(RangeCalculation::TerminateEarly); - } - - Ok(RangeCalculation::Range(Some(range))) - } - } -} - -/// Asynchronously finds the position of the first newline character in a specified byte range -/// within an object, such as a file, in an object store. -/// -/// This function scans the contents of the object starting from the specified `start` position -/// up to the `end` position, looking for the first occurrence of a newline character. -/// It returns the position of the first newline relative to the start of the range. -/// -/// Returns a `Result` wrapping a `usize` that represents the position of the first newline character found within the specified range. If no newline is found, it returns the length of the scanned data, effectively indicating the end of the range. -/// -/// The function returns an `Error` if any issues arise while reading from the object store or processing the data stream. -/// -async fn find_first_newline( - object_store: &Arc, - location: &Path, - start: usize, - end: usize, - newline: u8, -) -> Result { - let options = GetOptions { - range: Some(GetRange::Bounded(start..end)), - ..Default::default() - }; - - let result = object_store.get_opts(location, options).await?; - let mut result_stream = result.into_stream(); - - let mut index = 0; - - while let Some(chunk) = result_stream.next().await.transpose()? { - if let Some(position) = chunk.iter().position(|&byte| byte == newline) { - return Ok(index + position); - } - - index += chunk.len(); - } - - Ok(index) -} #[cfg(test)] mod tests { - use super::*; + use std::sync::Arc; use arrow::array::{ cast::AsArray, diff --git a/datafusion/core/src/datasource/physical_plan/parquet/mod.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs similarity index 72% rename from datafusion/core/src/datasource/physical_plan/parquet/mod.rs rename to datafusion/core/src/datasource/physical_plan/parquet.rs index f677c73cc881..888f3ad9e3b9 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/mod.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet.rs @@ -15,576 +15,57 @@ // specific language governing permissions and limitations // under the License. -//! [`ParquetExec`] FileSource for reading Parquet files - -mod access_plan; -mod metrics; -mod opener; -mod page_filter; -mod reader; -mod row_filter; -mod row_group_filter; -pub mod source; -mod writer; - -use std::any::Any; -use std::fmt::Formatter; -use std::sync::Arc; - -use crate::datasource::listing::PartitionedFile; -use crate::datasource::physical_plan::{parquet::source::ParquetSource, DisplayAs}; -use crate::datasource::schema_adapter::SchemaAdapterFactory; -use crate::{ - config::TableParquetOptions, - error::Result, - execution::context::TaskContext, - physical_plan::{ - metrics::MetricsSet, DisplayFormatType, ExecutionPlan, Partitioning, - PlanProperties, SendableRecordBatchStream, Statistics, - }, -}; - -pub use access_plan::{ParquetAccessPlan, RowGroupAccess}; -use arrow::datatypes::SchemaRef; -use datafusion_common::config::ConfigOptions; -use datafusion_common::Constraints; -use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::source::DataSourceExec; -use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; -use datafusion_physical_optimizer::pruning::PruningPredicate; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -pub use metrics::ParquetFileMetrics; -pub use page_filter::PagePruningAccessPlanFilter; -pub use reader::{DefaultParquetFileReaderFactory, ParquetFileReaderFactory}; -pub use row_filter::build_row_filter; -pub use row_filter::can_expr_be_pushed_down_with_schemas; -pub use row_group_filter::RowGroupAccessPlanFilter; -pub use writer::plan_to_parquet; - -use log::debug; - -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] -/// Deprecated Execution plan replaced with DataSourceExec -pub struct ParquetExec { - inner: DataSourceExec, - base_config: FileScanConfig, - table_parquet_options: TableParquetOptions, - /// Optional predicate for row filtering during parquet scan - predicate: Option>, - /// Optional predicate for pruning row groups (derived from `predicate`) - pruning_predicate: Option>, - /// Optional user defined parquet file reader factory - parquet_file_reader_factory: Option>, - /// Optional user defined schema adapter - schema_adapter_factory: Option>, -} - -#[allow(unused, deprecated)] -impl From for ParquetExecBuilder { - fn from(exec: ParquetExec) -> Self { - exec.into_builder() - } -} - -/// [`ParquetExecBuilder`], deprecated builder for [`ParquetExec`]. -/// -/// ParquetExec is replaced with `DataSourceExec` and it includes `ParquetSource` -/// -/// See example on [`ParquetSource`]. -#[deprecated( - since = "46.0.0", - note = "use DataSourceExec with ParquetSource instead" -)] -#[allow(unused, deprecated)] -pub struct ParquetExecBuilder { - file_scan_config: FileScanConfig, - predicate: Option>, - metadata_size_hint: Option, - table_parquet_options: TableParquetOptions, - parquet_file_reader_factory: Option>, - schema_adapter_factory: Option>, -} - -#[allow(unused, deprecated)] -impl ParquetExecBuilder { - /// Create a new builder to read the provided file scan configuration - pub fn new(file_scan_config: FileScanConfig) -> Self { - Self::new_with_options(file_scan_config, TableParquetOptions::default()) - } - - /// Create a new builder to read the data specified in the file scan - /// configuration with the provided `TableParquetOptions`. - pub fn new_with_options( - file_scan_config: FileScanConfig, - table_parquet_options: TableParquetOptions, - ) -> Self { - Self { - file_scan_config, - predicate: None, - metadata_size_hint: None, - table_parquet_options, - parquet_file_reader_factory: None, - schema_adapter_factory: None, - } - } - - /// Update the list of files groups to read - pub fn with_file_groups(mut self, file_groups: Vec>) -> Self { - self.file_scan_config.file_groups = file_groups; - self - } - - /// Set the filter predicate when reading. - /// - /// See the "Predicate Pushdown" section of the [`ParquetExec`] documentation - /// for more details. - pub fn with_predicate(mut self, predicate: Arc) -> Self { - self.predicate = Some(predicate); - self - } - - /// Set the metadata size hint - /// - /// This value determines how many bytes at the end of the file the default - /// [`ParquetFileReaderFactory`] will request in the initial IO. If this is - /// too small, the ParquetExec will need to make additional IO requests to - /// read the footer. - pub fn with_metadata_size_hint(mut self, metadata_size_hint: usize) -> Self { - self.metadata_size_hint = Some(metadata_size_hint); - self - } - - /// Set the options for controlling how the ParquetExec reads parquet files. - /// - /// See also [`Self::new_with_options`] - pub fn with_table_parquet_options( - mut self, - table_parquet_options: TableParquetOptions, - ) -> Self { - self.table_parquet_options = table_parquet_options; - self - } - - /// Set optional user defined parquet file reader factory. - /// - /// You can use [`ParquetFileReaderFactory`] to more precisely control how - /// data is read from parquet files (e.g. skip re-reading metadata, coalesce - /// I/O operations, etc). - /// - /// The default reader factory reads directly from an [`ObjectStore`] - /// instance using individual I/O operations for the footer and each page. - /// - /// If a custom `ParquetFileReaderFactory` is provided, then data access - /// operations will be routed to this factory instead of [`ObjectStore`]. - /// - /// [`ObjectStore`]: object_store::ObjectStore - pub fn with_parquet_file_reader_factory( - mut self, - parquet_file_reader_factory: Arc, - ) -> Self { - self.parquet_file_reader_factory = Some(parquet_file_reader_factory); - self - } - - /// Set optional schema adapter factory. - /// - /// [`SchemaAdapterFactory`] allows user to specify how fields from the - /// parquet file get mapped to that of the table schema. The default schema - /// adapter uses arrow's cast library to map the parquet fields to the table - /// schema. - pub fn with_schema_adapter_factory( - mut self, - schema_adapter_factory: Arc, - ) -> Self { - self.schema_adapter_factory = Some(schema_adapter_factory); - self - } +//! Reexports the [`datafusion_datasource_parquet`] crate, containing Parquet based [`FileSource`]. +//! +//! [`FileSource`]: datafusion_datasource::file::FileSource - /// Convenience: build an `Arc`d `ParquetExec` from this builder - pub fn build_arc(self) -> Arc { - Arc::new(self.build()) - } - - /// Build a [`ParquetExec`] - #[must_use] - pub fn build(self) -> ParquetExec { - let Self { - file_scan_config, - predicate, - metadata_size_hint, - table_parquet_options, - parquet_file_reader_factory, - schema_adapter_factory, - } = self; - let mut parquet = ParquetSource::new(table_parquet_options); - if let Some(predicate) = predicate.clone() { - parquet = parquet - .with_predicate(Arc::clone(&file_scan_config.file_schema), predicate); - } - if let Some(metadata_size_hint) = metadata_size_hint { - parquet = parquet.with_metadata_size_hint(metadata_size_hint) - } - if let Some(parquet_reader_factory) = parquet_file_reader_factory { - parquet = parquet.with_parquet_file_reader_factory(parquet_reader_factory) - } - if let Some(schema_factory) = schema_adapter_factory { - parquet = parquet.with_schema_adapter_factory(schema_factory); - } - - let base_config = file_scan_config.with_source(Arc::new(parquet.clone())); - debug!("Creating ParquetExec, files: {:?}, projection {:?}, predicate: {:?}, limit: {:?}", - base_config.file_groups, base_config.projection, predicate, base_config.limit); - - ParquetExec { - inner: DataSourceExec::new(Arc::new(base_config.clone())), - base_config, - predicate, - pruning_predicate: parquet.pruning_predicate, - schema_adapter_factory: parquet.schema_adapter_factory, - parquet_file_reader_factory: parquet.parquet_file_reader_factory, - table_parquet_options: parquet.table_parquet_options, - } - } -} - -#[allow(unused, deprecated)] -impl ParquetExec { - /// Create a new Parquet reader execution plan provided file list and schema. - pub fn new( - base_config: FileScanConfig, - predicate: Option>, - metadata_size_hint: Option, - table_parquet_options: TableParquetOptions, - ) -> Self { - let mut builder = - ParquetExecBuilder::new_with_options(base_config, table_parquet_options); - if let Some(predicate) = predicate { - builder = builder.with_predicate(predicate); - } - if let Some(metadata_size_hint) = metadata_size_hint { - builder = builder.with_metadata_size_hint(metadata_size_hint); - } - builder.build() - } - /// Return a [`ParquetExecBuilder`]. - /// - /// See example on [`ParquetExec`] and [`ParquetExecBuilder`] for specifying - /// parquet table options. - pub fn builder(file_scan_config: FileScanConfig) -> ParquetExecBuilder { - ParquetExecBuilder::new(file_scan_config) - } - - /// Convert this `ParquetExec` into a builder for modification - pub fn into_builder(self) -> ParquetExecBuilder { - // list out fields so it is clear what is being dropped - // (note the fields which are dropped are re-created as part of calling - // `build` on the builder) - let file_scan_config = self.file_scan_config(); - let parquet = self.parquet_source(); - - ParquetExecBuilder { - file_scan_config, - predicate: parquet.predicate, - metadata_size_hint: parquet.metadata_size_hint, - table_parquet_options: parquet.table_parquet_options, - parquet_file_reader_factory: parquet.parquet_file_reader_factory, - schema_adapter_factory: parquet.schema_adapter_factory, - } - } - fn file_scan_config(&self) -> FileScanConfig { - self.inner - .data_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - fn parquet_source(&self) -> ParquetSource { - self.file_scan_config() - .file_source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - /// [`FileScanConfig`] that controls this scan (such as which files to read) - pub fn base_config(&self) -> &FileScanConfig { - &self.base_config - } - /// Options passed to the parquet reader for this scan - pub fn table_parquet_options(&self) -> &TableParquetOptions { - &self.table_parquet_options - } - /// Optional predicate. - pub fn predicate(&self) -> Option<&Arc> { - self.predicate.as_ref() - } - /// Optional reference to this parquet scan's pruning predicate - pub fn pruning_predicate(&self) -> Option<&Arc> { - self.pruning_predicate.as_ref() - } - /// return the optional file reader factory - pub fn parquet_file_reader_factory( - &self, - ) -> Option<&Arc> { - self.parquet_file_reader_factory.as_ref() - } - /// Optional user defined parquet file reader factory. - pub fn with_parquet_file_reader_factory( - mut self, - parquet_file_reader_factory: Arc, - ) -> Self { - let mut parquet = self.parquet_source(); - parquet.parquet_file_reader_factory = - Some(Arc::clone(&parquet_file_reader_factory)); - let file_source = self.file_scan_config(); - self.inner = self - .inner - .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); - self.parquet_file_reader_factory = Some(parquet_file_reader_factory); - self - } - /// return the optional schema adapter factory - pub fn schema_adapter_factory(&self) -> Option<&Arc> { - self.schema_adapter_factory.as_ref() - } - /// Set optional schema adapter factory. - /// - /// [`SchemaAdapterFactory`] allows user to specify how fields from the - /// parquet file get mapped to that of the table schema. The default schema - /// adapter uses arrow's cast library to map the parquet fields to the table - /// schema. - pub fn with_schema_adapter_factory( - mut self, - schema_adapter_factory: Arc, - ) -> Self { - let mut parquet = self.parquet_source(); - parquet.schema_adapter_factory = Some(Arc::clone(&schema_adapter_factory)); - let file_source = self.file_scan_config(); - self.inner = self - .inner - .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); - self.schema_adapter_factory = Some(schema_adapter_factory); - self - } - /// If true, the predicate will be used during the parquet scan. - /// Defaults to false - /// - /// [`Expr`]: datafusion_expr::Expr - pub fn with_pushdown_filters(mut self, pushdown_filters: bool) -> Self { - let mut parquet = self.parquet_source(); - parquet.table_parquet_options.global.pushdown_filters = pushdown_filters; - let file_source = self.file_scan_config(); - self.inner = self - .inner - .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); - self.table_parquet_options.global.pushdown_filters = pushdown_filters; - self - } - - /// Return the value described in [`Self::with_pushdown_filters`] - fn pushdown_filters(&self) -> bool { - self.parquet_source() - .table_parquet_options - .global - .pushdown_filters - } - /// If true, the `RowFilter` made by `pushdown_filters` may try to - /// minimize the cost of filter evaluation by reordering the - /// predicate [`Expr`]s. If false, the predicates are applied in - /// the same order as specified in the query. Defaults to false. - /// - /// [`Expr`]: datafusion_expr::Expr - pub fn with_reorder_filters(mut self, reorder_filters: bool) -> Self { - let mut parquet = self.parquet_source(); - parquet.table_parquet_options.global.reorder_filters = reorder_filters; - let file_source = self.file_scan_config(); - self.inner = self - .inner - .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); - self.table_parquet_options.global.reorder_filters = reorder_filters; - self - } - /// Return the value described in [`Self::with_reorder_filters`] - fn reorder_filters(&self) -> bool { - self.parquet_source() - .table_parquet_options - .global - .reorder_filters - } - /// If enabled, the reader will read the page index - /// This is used to optimize filter pushdown - /// via `RowSelector` and `RowFilter` by - /// eliminating unnecessary IO and decoding - fn bloom_filter_on_read(&self) -> bool { - self.parquet_source() - .table_parquet_options - .global - .bloom_filter_on_read - } - /// Return the value described in [`ParquetSource::with_enable_page_index`] - fn enable_page_index(&self) -> bool { - self.parquet_source() - .table_parquet_options - .global - .enable_page_index - } - - fn output_partitioning_helper(file_config: &FileScanConfig) -> Partitioning { - Partitioning::UnknownPartitioning(file_config.file_groups.len()) - } - - /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( - schema: SchemaRef, - orderings: &[LexOrdering], - constraints: Constraints, - file_config: &FileScanConfig, - ) -> PlanProperties { - PlanProperties::new( - EquivalenceProperties::new_with_orderings(schema, orderings) - .with_constraints(constraints), - Self::output_partitioning_helper(file_config), // Output Partitioning - EmissionType::Incremental, - Boundedness::Bounded, - ) - } - - /// Updates the file groups to read and recalculates the output partitioning - /// - /// Note this function does not update statistics or other properties - /// that depend on the file groups. - fn with_file_groups_and_update_partitioning( - mut self, - file_groups: Vec>, - ) -> Self { - let mut config = self.file_scan_config(); - config.file_groups = file_groups; - self.inner = self.inner.with_data_source(Arc::new(config)); - self - } -} - -#[allow(unused, deprecated)] -impl DisplayAs for ParquetExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - self.inner.fmt_as(t, f) - } -} - -#[allow(unused, deprecated)] -impl ExecutionPlan for ParquetExec { - fn name(&self) -> &'static str { - "ParquetExec" - } - - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - // this is a leaf node and has no children - vec![] - } - - fn with_new_children( - self: Arc, - _: Vec>, - ) -> Result> { - Ok(self) - } - - /// Redistribute files across partitions according to their size - /// See comments on `FileGroupPartitioner` for more detail. - fn repartitioned( - &self, - target_partitions: usize, - config: &ConfigOptions, - ) -> Result>> { - self.inner.repartitioned(target_partitions, config) - } - - fn execute( - &self, - partition_index: usize, - ctx: Arc, - ) -> Result { - self.inner.execute(partition_index, ctx) - } - fn metrics(&self) -> Option { - self.inner.metrics() - } - fn statistics(&self) -> Result { - self.inner.statistics() - } - fn fetch(&self) -> Option { - self.inner.fetch() - } - - fn with_fetch(&self, limit: Option) -> Option> { - self.inner.with_fetch(limit) - } -} - -fn should_enable_page_index( - enable_page_index: bool, - page_pruning_predicate: &Option>, -) -> bool { - enable_page_index - && page_pruning_predicate.is_some() - && page_pruning_predicate - .as_ref() - .map(|p| p.filter_number() > 0) - .unwrap_or(false) -} +pub use datafusion_datasource_parquet::*; #[cfg(test)] mod tests { // See also `parquet_exec` integration test use std::fs::{self, File}; use std::io::Write; + use std::sync::Arc; use std::sync::Mutex; - use super::*; use crate::dataframe::DataFrameWriteOptions; use crate::datasource::file_format::options::CsvReadOptions; use crate::datasource::file_format::parquet::test_util::store_parquet; use crate::datasource::file_format::test_util::scan_format; - use crate::datasource::listing::{FileRange, ListingOptions, PartitionedFile}; - use crate::datasource::object_store::ObjectStoreUrl; + use crate::datasource::listing::ListingOptions; use crate::execution::context::SessionState; - use crate::physical_plan::displayable; use crate::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use crate::test::object_store::local_unpartitioned_file; - use crate::{ - assert_batches_sorted_eq, - datasource::file_format::{parquet::ParquetFormat, FileFormat}, - physical_plan::collect, - }; - use arrow::array::{ ArrayRef, Date64Array, Int32Array, Int64Array, Int8Array, StringArray, StructArray, }; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaBuilder}; use arrow::record_batch::RecordBatch; + use arrow_schema::SchemaRef; use bytes::{BufMut, BytesMut}; - use datafusion_common::{assert_contains, ScalarValue}; + use datafusion_common::config::TableParquetOptions; + use datafusion_common::{ + assert_batches_eq, assert_batches_sorted_eq, assert_contains, Result, ScalarValue, + }; + use datafusion_datasource::file_format::FileFormat; + use datafusion_datasource::file_meta::FileMeta; + use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::source::DataSourceExec; + + use datafusion_datasource::{FileRange, PartitionedFile}; + use datafusion_datasource_parquet::source::ParquetSource; + use datafusion_datasource_parquet::{ + DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat, + }; + use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{col, lit, when, Expr}; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; + use datafusion_physical_plan::{collect, displayable}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; - use crate::datasource::physical_plan::parquet::source::ParquetSource; use chrono::{TimeZone, Utc}; use futures::StreamExt; use object_store::local::LocalFileSystem; @@ -743,28 +224,6 @@ mod tests { ) } - #[tokio::test] - async fn write_parquet_results_error_handling() -> Result<()> { - let ctx = SessionContext::new(); - // register a local file system object store for /tmp directory - let tmp_dir = TempDir::new()?; - let local = Arc::new(LocalFileSystem::new_with_prefix(&tmp_dir)?); - let local_url = Url::parse("file://local").unwrap(); - ctx.register_object_store(&local_url, local); - - let options = CsvReadOptions::default() - .schema_infer_max_records(2) - .has_header(true); - let df = ctx.read_csv("tests/data/corrupt.csv", options).await?; - let out_dir_url = "file://local/out"; - let e = df - .write_parquet(out_dir_url, DataFrameWriteOptions::new(), None) - .await - .expect_err("should fail because input file does not match inferred schema"); - assert_eq!(e.strip_backtrace(), "Arrow error: Parser error: Error while parsing value d for column 0 at line 4"); - Ok(()) - } - #[tokio::test] async fn evolved_schema() { let c1: ArrayRef = @@ -1288,7 +747,7 @@ mod tests { #[tokio::test] async fn parquet_exec_with_projection() -> Result<()> { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let filename = "alltypes_plain.parquet"; let session_ctx = SessionContext::new(); let state = session_ctx.state(); @@ -1376,7 +835,7 @@ mod tests { let session_ctx = SessionContext::new(); let state = session_ctx.state(); - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let filename = format!("{testdata}/alltypes_plain.parquet"); let meta = local_unpartitioned_file(filename); @@ -1409,7 +868,7 @@ mod tests { let object_store_url = ObjectStoreUrl::local_filesystem(); let store = state.runtime_env().object_store(&object_store_url).unwrap(); - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let filename = format!("{testdata}/alltypes_plain.parquet"); let meta = local_unpartitioned_file(filename); @@ -1492,7 +951,7 @@ mod tests { "| 1 | false | 1 | 10 | 26 |", "+----+----------+-------------+-------+-----+", ]; - crate::assert_batches_eq!(expected, &[batch]); + assert_batches_eq!(expected, &[batch]); let batch = results.next().await; assert!(batch.is_none()); @@ -1675,7 +1134,7 @@ mod tests { .await; // should have a pruning predicate - let pruning_predicate = &rt.parquet_source.pruning_predicate; + let pruning_predicate = rt.parquet_source.pruning_predicate(); assert!(pruning_predicate.is_some()); // convert to explain plan form @@ -1716,7 +1175,7 @@ mod tests { .round_trip(vec![batches.clone()]) .await; - let pruning_predicate = &rt0.parquet_source.pruning_predicate; + let pruning_predicate = rt0.parquet_source.pruning_predicate(); assert!(pruning_predicate.is_some()); let display0 = displayable(rt0.parquet_exec.as_ref()) @@ -1758,9 +1217,9 @@ mod tests { .await; // should have a pruning predicate - let pruning_predicate = &rt1.parquet_source.pruning_predicate; + let pruning_predicate = rt1.parquet_source.pruning_predicate(); assert!(pruning_predicate.is_some()); - let pruning_predicate = &rt2.parquet_source.predicate; + let pruning_predicate = rt2.parquet_source.predicate(); assert!(pruning_predicate.is_some()); // convert to explain plan form @@ -1801,14 +1260,14 @@ mod tests { .await; // Should not contain a pruning predicate (since nothing can be pruned) - let pruning_predicate = &rt.parquet_source.pruning_predicate; + let pruning_predicate = rt.parquet_source.pruning_predicate(); assert!( pruning_predicate.is_none(), "Still had pruning predicate: {pruning_predicate:?}" ); // but does still has a pushdown down predicate - let predicate = rt.parquet_source.predicate.as_ref(); + let predicate = rt.parquet_source.predicate(); let filter_phys = logical2physical(&filter, rt.parquet_exec.schema().as_ref()); assert_eq!(predicate.unwrap().to_string(), filter_phys.to_string()); } @@ -1836,7 +1295,7 @@ mod tests { .await; // Should have a pruning predicate - let pruning_predicate = &rt.parquet_source.pruning_predicate; + let pruning_predicate = rt.parquet_source.pruning_predicate(); assert!(pruning_predicate.is_some()); } @@ -1996,7 +1455,7 @@ mod tests { "| {id: 4, name: aaa2} | 2 | test02 |", "+---------------------+----+--------+", ]; - crate::assert_batches_eq!(expected, &batch); + assert_batches_eq!(expected, &batch); Ok(()) } @@ -2026,7 +1485,7 @@ mod tests { "| {id: 4, name: aaa2} | 2 | test02 |", "+---------------------+----+--------+", ]; - crate::assert_batches_eq!(expected, &batch); + assert_batches_eq!(expected, &batch); Ok(()) } @@ -2108,7 +1567,7 @@ mod tests { fn create_reader( &self, partition_index: usize, - file_meta: crate::datasource::physical_plan::FileMeta, + file_meta: FileMeta, metadata_size_hint: Option, metrics: &ExecutionPlanMetricsSet, ) -> Result> diff --git a/datafusion/core/src/datasource/statistics.rs b/datafusion/core/src/datasource/statistics.rs index f02927619a7d..acc16e0687c9 100644 --- a/datafusion/core/src/datasource/statistics.rs +++ b/datafusion/core/src/datasource/statistics.rs @@ -27,13 +27,6 @@ use crate::arrow::datatypes::SchemaRef; use crate::error::Result; use crate::physical_plan::{ColumnStatistics, Statistics}; -#[cfg(feature = "parquet")] -use crate::{ - arrow::datatypes::Schema, - functions_aggregate::min_max::{MaxAccumulator, MinAccumulator}, - physical_plan::Accumulator, -}; - use super::listing::PartitionedFile; /// Get all files as well as the file level summary statistics (no statistic for partition columns). @@ -152,28 +145,6 @@ pub async fn get_statistics_with_limit( Ok((result_files, statistics)) } -// only adding this cfg b/c this is the only feature it's used with currently -#[cfg(feature = "parquet")] -pub(crate) fn create_max_min_accs( - schema: &Schema, -) -> (Vec>, Vec>) { - let max_values: Vec> = schema - .fields() - .iter() - .map(|field| { - MaxAccumulator::try_new(min_max_aggregate_data_type(field.data_type())).ok() - }) - .collect(); - let min_values: Vec> = schema - .fields() - .iter() - .map(|field| { - MinAccumulator::try_new(min_max_aggregate_data_type(field.data_type())).ok() - }) - .collect(); - (max_values, min_values) -} - fn add_row_stats( file_num_rows: Precision, num_rows: Precision, @@ -185,52 +156,6 @@ fn add_row_stats( } } -// only adding this cfg b/c this is the only feature it's used with currently -#[cfg(feature = "parquet")] -pub(crate) fn get_col_stats( - schema: &Schema, - null_counts: Vec>, - max_values: &mut [Option], - min_values: &mut [Option], -) -> Vec { - (0..schema.fields().len()) - .map(|i| { - let max_value = match max_values.get_mut(i).unwrap() { - Some(max_value) => max_value.evaluate().ok(), - None => None, - }; - let min_value = match min_values.get_mut(i).unwrap() { - Some(min_value) => min_value.evaluate().ok(), - None => None, - }; - ColumnStatistics { - null_count: null_counts[i], - max_value: max_value.map(Precision::Exact).unwrap_or(Precision::Absent), - min_value: min_value.map(Precision::Exact).unwrap_or(Precision::Absent), - sum_value: Precision::Absent, - distinct_count: Precision::Absent, - } - }) - .collect() -} - -// Min/max aggregation can take Dictionary encode input but always produces unpacked -// (aka non Dictionary) output. We need to adjust the output data type to reflect this. -// The reason min/max aggregate produces unpacked output because there is only one -// min/max value per group; there is no needs to keep them Dictionary encode -// -// only adding this cfg b/c this is the only feature it's used with currently -#[cfg(feature = "parquet")] -fn min_max_aggregate_data_type( - input_type: &arrow_schema::DataType, -) -> &arrow_schema::DataType { - if let arrow_schema::DataType::Dictionary(_, value_type) = input_type { - value_type.as_ref() - } else { - input_type - } -} - /// If the given value is numerically greater than the original maximum value, /// return the new maximum value with appropriate exactness information. fn set_max_if_greater( diff --git a/datafusion/core/src/datasource/view.rs b/datafusion/core/src/datasource/view.rs index 91e9b6789fda..e4f57b0d9798 100644 --- a/datafusion/core/src/datasource/view.rs +++ b/datafusion/core/src/datasource/view.rs @@ -30,7 +30,6 @@ use datafusion_catalog::Session; use datafusion_common::config::ConfigOptions; use datafusion_common::Column; use datafusion_expr::{LogicalPlanBuilder, TableProviderFilterPushDown}; -use datafusion_optimizer::analyzer::expand_wildcard_rule::ExpandWildcardRule; use datafusion_optimizer::analyzer::type_coercion::TypeCoercion; use datafusion_optimizer::Analyzer; @@ -68,11 +67,11 @@ impl ViewTable { fn apply_required_rule(logical_plan: LogicalPlan) -> Result { let options = ConfigOptions::default(); - Analyzer::with_rules(vec![ - Arc::new(ExpandWildcardRule::new()), - Arc::new(TypeCoercion::new()), - ]) - .execute_and_check(logical_plan, &options, |_, _| {}) + Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]).execute_and_check( + logical_plan, + &options, + |_, _| {}, + ) } /// Get definition ref diff --git a/datafusion/core/src/execution/context/csv.rs b/datafusion/core/src/execution/context/csv.rs index 9b4c0e3b2964..3e7db1caa20f 100644 --- a/datafusion/core/src/execution/context/csv.rs +++ b/datafusion/core/src/execution/context/csv.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::datasource::physical_plan::plan_to_csv; use datafusion_common::TableReference; +use datafusion_datasource_csv::source::plan_to_csv; use std::sync::Arc; use super::super::options::{CsvReadOptions, ReadOptions}; diff --git a/datafusion/core/src/execution/context/json.rs b/datafusion/core/src/execution/context/json.rs index 013c47d046fc..e9d799400863 100644 --- a/datafusion/core/src/execution/context/json.rs +++ b/datafusion/core/src/execution/context/json.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::datasource::physical_plan::plan_to_json; use datafusion_common::TableReference; +use datafusion_datasource_json::source::plan_to_json; use std::sync::Arc; use super::super::options::{NdJsonReadOptions, ReadOptions}; diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index c27d1e4fd46b..ad0993ed43ca 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -83,12 +83,14 @@ use object_store::ObjectStore; use parking_lot::RwLock; use url::Url; -mod avro; mod csv; mod json; #[cfg(feature = "parquet")] mod parquet; +#[cfg(feature = "avro")] +mod avro; + /// DataFilePaths adds a method to convert strings and vector of strings to vector of [`ListingTableUrl`] URLs. /// This allows methods such [`SessionContext::read_csv`] and [`SessionContext::read_avro`] /// to take either a single file or multiple files. diff --git a/datafusion/core/src/execution/context/parquet.rs b/datafusion/core/src/execution/context/parquet.rs index 67ccacaea666..6ec9796fe90d 100644 --- a/datafusion/core/src/execution/context/parquet.rs +++ b/datafusion/core/src/execution/context/parquet.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use super::super::options::{ParquetReadOptions, ReadOptions}; use super::{DataFilePaths, DataFrame, ExecutionPlan, Result, SessionContext}; -use crate::datasource::physical_plan::parquet::plan_to_parquet; +use datafusion_datasource_parquet::plan_to_parquet; use datafusion_common::TableReference; use parquet::file::properties::WriterProperties; diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index bdaae4f6985b..f4b0fd0c125f 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -68,7 +68,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_optimizer::optimizer::PhysicalOptimizer; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::ExecutionPlan; -use datafusion_sql::parser::{DFParser, Statement}; +use datafusion_sql::parser::{DFParserBuilder, Statement}; use datafusion_sql::planner::{ContextProvider, ParserOptions, PlannerContext, SqlToRel}; use async_trait::async_trait; @@ -258,6 +258,14 @@ impl Session for SessionState { fn as_any(&self) -> &dyn Any { self } + + fn table_options(&self) -> &TableOptions { + self.table_options() + } + + fn table_options_mut(&mut self) -> &mut TableOptions { + self.table_options_mut() + } } impl SessionState { @@ -272,22 +280,6 @@ impl SessionState { .build() } - /// Returns new [`SessionState`] using the provided - /// [`SessionConfig`], [`RuntimeEnv`], and [`CatalogProviderList`] - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - pub fn new_with_config_rt_and_catalog_list( - config: SessionConfig, - runtime: Arc, - catalog_list: Arc, - ) -> Self { - SessionStateBuilder::new() - .with_config(config) - .with_runtime_env(runtime) - .with_catalog_list(catalog_list) - .with_default_features() - .build() - } - pub(crate) fn resolve_table_ref( &self, table_ref: impl Into, @@ -326,53 +318,6 @@ impl SessionState { }) } - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Replace the random session id. - pub fn with_session_id(mut self, session_id: String) -> Self { - self.session_id = session_id; - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// override default query planner with `query_planner` - pub fn with_query_planner( - mut self, - query_planner: Arc, - ) -> Self { - self.query_planner = query_planner; - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Override the [`AnalyzerRule`]s optimizer plan rules. - pub fn with_analyzer_rules( - mut self, - rules: Vec>, - ) -> Self { - self.analyzer = Analyzer::with_rules(rules); - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Replace the entire list of [`OptimizerRule`]s used to optimize plans - pub fn with_optimizer_rules( - mut self, - rules: Vec>, - ) -> Self { - self.optimizer = Optimizer::with_rules(rules); - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Replace the entire list of [`PhysicalOptimizerRule`]s used to optimize plans - pub fn with_physical_optimizer_rules( - mut self, - physical_optimizers: Vec>, - ) -> Self { - self.physical_optimizers = PhysicalOptimizer::with_rules(physical_optimizers); - self - } - /// Add `analyzer_rule` to the end of the list of /// [`AnalyzerRule`]s used to rewrite queries. pub fn add_analyzer_rule( @@ -383,17 +328,6 @@ impl SessionState { self } - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Add `optimizer_rule` to the end of the list of - /// [`OptimizerRule`]s used to rewrite queries. - pub fn add_optimizer_rule( - mut self, - optimizer_rule: Arc, - ) -> Self { - self.optimizer.rules.push(optimizer_rule); - self - } - // the add_optimizer_rule takes an owned reference // it should probably be renamed to `with_optimizer_rule` to follow builder style // and `add_optimizer_rule` that takes &mut self added instead of this @@ -404,52 +338,11 @@ impl SessionState { self.optimizer.rules.push(optimizer_rule); } - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Add `physical_optimizer_rule` to the end of the list of - /// [`PhysicalOptimizerRule`]s used to rewrite queries. - pub fn add_physical_optimizer_rule( - mut self, - physical_optimizer_rule: Arc, - ) -> Self { - self.physical_optimizers.rules.push(physical_optimizer_rule); - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Adds a new [`ConfigExtension`] to TableOptions - pub fn add_table_options_extension( - mut self, - extension: T, - ) -> Self { - self.table_options.extensions.insert(extension); - self - } - - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Registers a [`FunctionFactory`] to handle `CREATE FUNCTION` statements - pub fn with_function_factory( - mut self, - function_factory: Arc, - ) -> Self { - self.function_factory = Some(function_factory); - self - } - /// Registers a [`FunctionFactory`] to handle `CREATE FUNCTION` statements pub fn set_function_factory(&mut self, function_factory: Arc) { self.function_factory = Some(function_factory); } - #[deprecated(since = "40.0.0", note = "Use SessionStateBuilder")] - /// Replace the extension [`SerializerRegistry`] - pub fn with_serializer_registry( - mut self, - registry: Arc, - ) -> Self { - self.serializer_registry = registry; - self - } - /// Get the function factory pub fn function_factory(&self) -> Option<&Arc> { self.function_factory.as_ref() @@ -483,12 +376,21 @@ impl SessionState { MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks." ) })?; - let mut statements = DFParser::parse_sql_with_dialect(sql, dialect.as_ref())?; + + let recursion_limit = self.config.options().sql_parser.recursion_limit; + + let mut statements = DFParserBuilder::new(sql) + .with_dialect(dialect.as_ref()) + .with_recursion_limit(recursion_limit) + .build()? + .parse_statements()?; + if statements.len() > 1 { return not_impl_err!( "The context currently only supports a single SQL statement" ); } + let statement = statements.pop_front().ok_or_else(|| { plan_datafusion_err!("No SQL statements were provided in the query string") })?; @@ -522,7 +424,12 @@ impl SessionState { ) })?; - let expr = DFParser::parse_sql_into_expr_with_dialect(sql, dialect.as_ref())?; + let recursion_limit = self.config.options().sql_parser.recursion_limit; + let expr = DFParserBuilder::new(sql) + .with_dialect(dialect.as_ref()) + .with_recursion_limit(recursion_limit) + .build()? + .parse_expr()?; Ok(expr) } @@ -582,6 +489,7 @@ impl SessionState { enable_options_value_normalization: sql_parser_options .enable_options_value_normalization, support_varchar_with_length: sql_parser_options.support_varchar_with_length, + map_varchar_to_utf8view: sql_parser_options.map_varchar_to_utf8view, collect_spans: sql_parser_options.collect_spans, } } @@ -818,18 +726,17 @@ impl SessionState { self.config.options() } - /// return the TableOptions options with its extensions - pub fn default_table_options(&self) -> TableOptions { - self.table_options - .combine_with_session_config(self.config_options()) - } - /// Return the table options pub fn table_options(&self) -> &TableOptions { &self.table_options } - /// Return mutable table options + /// return the TableOptions options with its extensions + pub fn default_table_options(&self) -> TableOptions { + Session::default_table_options(self) + } + + /// Returns a mutable reference to [`TableOptions`] pub fn table_options_mut(&mut self) -> &mut TableOptions { &mut self.table_options } @@ -998,7 +905,13 @@ pub struct SessionStateBuilder { } impl SessionStateBuilder { - /// Returns a new [`SessionStateBuilder`] with no options set. + /// Returns a new empty [`SessionStateBuilder`]. + /// + /// See [`Self::with_default_features`] to install the default set of functions, + /// catalogs, etc. + /// + /// To create a `SessionStateBuilder` with default features such as functions, + /// please see [`Self::new_with_default_features`]. pub fn new() -> Self { Self { session_id: None, @@ -1028,9 +941,10 @@ impl SessionStateBuilder { } } - /// Returns a new [SessionStateBuilder] based on an existing [SessionState] + /// Returns a new [SessionStateBuilder] based on an existing [SessionState]. + /// /// The session id for the new builder will be unset; all other fields will - /// be cloned from what is set in the provided session state. If the default + /// be cloned from `existing`. If the default /// catalog exists in existing session state, the new session state will not /// create default catalog and schema. pub fn new_from_existing(existing: SessionState) -> Self { @@ -1079,16 +993,60 @@ impl SessionStateBuilder { } } - /// Create default builder with defaults for table_factories, file formats, expr_planners and builtin + /// Adds defaults for table_factories, file formats, expr_planners and builtin /// scalar, aggregate and windows functions. - pub fn with_default_features(self) -> Self { - self.with_table_factories(SessionStateDefaults::default_table_factories()) - .with_file_formats(SessionStateDefaults::default_file_formats()) - .with_expr_planners(SessionStateDefaults::default_expr_planners()) - .with_scalar_functions(SessionStateDefaults::default_scalar_functions()) - .with_aggregate_functions(SessionStateDefaults::default_aggregate_functions()) - .with_window_functions(SessionStateDefaults::default_window_functions()) - .with_table_function_list(SessionStateDefaults::default_table_functions()) + /// + /// Note overwrites any previously registered items with the same name. + pub fn with_default_features(mut self) -> Self { + self.table_factories + .get_or_insert_with(HashMap::new) + .extend(SessionStateDefaults::default_table_factories()); + + self.file_formats + .get_or_insert_with(Vec::new) + .extend(SessionStateDefaults::default_file_formats()); + + self.expr_planners + .get_or_insert_with(Vec::new) + .extend(SessionStateDefaults::default_expr_planners()); + + self.scalar_functions + .get_or_insert_with(Vec::new) + .extend(SessionStateDefaults::default_scalar_functions()); + + self.aggregate_functions + .get_or_insert_with(Vec::new) + .extend(SessionStateDefaults::default_aggregate_functions()); + + self.window_functions + .get_or_insert_with(Vec::new) + .extend(SessionStateDefaults::default_window_functions()); + + self.table_functions + .get_or_insert_with(HashMap::new) + .extend( + SessionStateDefaults::default_table_functions() + .into_iter() + .map(|f| (f.name().to_string(), f)), + ); + + self + } + + /// Returns a new [`SessionStateBuilder`] with default features. + /// + /// This is equivalent to calling [`Self::new()`] followed by [`Self::with_default_features()`]. + /// + /// ``` + /// use datafusion::execution::session_state::SessionStateBuilder; + /// + /// // Create a new SessionState with default features + /// let session_state = SessionStateBuilder::new_with_default_features() + /// .with_session_id("my_session".to_string()) + /// .build(); + /// ``` + pub fn new_with_default_features() -> Self { + Self::new().with_default_features() } /// Set the session id. @@ -2141,4 +2099,19 @@ mod tests { assert!(table_factories.contains_key("employee")); Ok(()) } + + #[test] + fn test_with_default_features_not_override() -> Result<()> { + use crate::test_util::TestTableFactory; + + // Test whether the table_factory has been overridden. + let table_factory = Arc::new(TestTableFactory {}); + let session_state = SessionStateBuilder::new() + .with_table_factory("test".to_string(), table_factory) + .with_default_features() + .build(); + assert!(session_state.table_factories().get("test").is_some()); + + Ok(()) + } } diff --git a/datafusion/core/src/execution/session_state_defaults.rs b/datafusion/core/src/execution/session_state_defaults.rs index 33bf01cf35cd..b48ef90f2bd5 100644 --- a/datafusion/core/src/execution/session_state_defaults.rs +++ b/datafusion/core/src/execution/session_state_defaults.rs @@ -18,6 +18,7 @@ use crate::catalog::{CatalogProvider, TableProviderFactory}; use crate::catalog_common::listing_schema::ListingSchemaProvider; use crate::datasource::file_format::arrow::ArrowFormatFactory; +#[cfg(feature = "avro")] use crate::datasource::file_format::avro::AvroFormatFactory; use crate::datasource::file_format::csv::CsvFormatFactory; use crate::datasource::file_format::json::JsonFormatFactory; @@ -135,6 +136,7 @@ impl SessionStateDefaults { Arc::new(JsonFormatFactory::new()), Arc::new(CsvFormatFactory::new()), Arc::new(ArrowFormatFactory::new()), + #[cfg(feature = "avro")] Arc::new(AvroFormatFactory::new()), ]; diff --git a/datafusion/core/src/lib.rs b/datafusion/core/src/lib.rs index 9a0d0157c1ae..b4d5f9740baa 100644 --- a/datafusion/core/src/lib.rs +++ b/datafusion/core/src/lib.rs @@ -298,10 +298,10 @@ //! (built in or user provided) ExecutionPlan //! ``` //! -//! DataFusion includes several built in data sources for common use -//! cases, and can be extended by implementing the [`TableProvider`] -//! trait. A [`TableProvider`] provides information for planning and -//! an [`ExecutionPlan`]s for execution. +//! A [`TableProvider`] provides information for planning and +//! an [`ExecutionPlan`]s for execution. DataFusion includes [`ListingTable`] +//! which supports reading several common file formats, and you can support any +//! new file format by implementing the [`TableProvider`] trait. See also: //! //! 1. [`ListingTable`]: Reads data from Parquet, JSON, CSV, or AVRO //! files. Supports single files or multiple files with HIVE style @@ -314,7 +314,7 @@ //! //! [`ListingTable`]: crate::datasource::listing::ListingTable //! [`MemTable`]: crate::datasource::memory::MemTable -//! [`StreamingTable`]: datafusion_catalog::streaming::StreamingTable +//! [`StreamingTable`]: crate::catalog::streaming::StreamingTable //! //! ## Plan Representations //! @@ -1026,12 +1026,6 @@ doc_comment::doctest!( library_user_guide_adding_udfs ); -#[cfg(doctest)] -doc_comment::doctest!( - "../../../docs/source/library-user-guide/api-health.md", - library_user_guide_api_health -); - #[cfg(doctest)] doc_comment::doctest!( "../../../docs/source/library-user-guide/building-logical-plans.md", @@ -1097,3 +1091,15 @@ doc_comment::doctest!( "../../../docs/source/library-user-guide/working-with-exprs.md", library_user_guide_working_with_exprs ); + +#[cfg(doctest)] +doc_comment::doctest!( + "../../../docs/source/library-user-guide/upgrading.md", + library_user_guide_upgrading +); + +#[cfg(doctest)] +doc_comment::doctest!( + "../../../docs/source/contributor-guide/api-health.md", + contributor_guide_api_health +); diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index a74cdcc5920b..6aff9280ffad 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -19,6 +19,7 @@ use std::borrow::Cow; use std::collections::HashMap; +use std::str::FromStr; use std::sync::Arc; use crate::datasource::file_format::file_type_to_format; @@ -77,8 +78,9 @@ use datafusion_expr::expr::{ use datafusion_expr::expr_rewriter::unnormalize_cols; use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary; use datafusion_expr::{ - DescribeTable, DmlStatement, Extension, FetchType, Filter, JoinType, RecursiveQuery, - SkipType, SortExpr, StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, + Analyze, DescribeTable, DmlStatement, Explain, Extension, FetchType, Filter, + JoinType, RecursiveQuery, SkipType, SortExpr, StringifiedPlan, WindowFrame, + WindowFrameBound, WriteOp, }; use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; use datafusion_physical_expr::expressions::Literal; @@ -87,6 +89,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::execution_plan::InvariantLevel; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::unnest::ListUnnest; +use datafusion_physical_plan::DisplayFormatType; use crate::schema_equivalence::schema_satisfied_by; use async_trait::async_trait; @@ -175,16 +178,17 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { logical_plan: &LogicalPlan, session_state: &SessionState, ) -> Result> { - match self.handle_explain(logical_plan, session_state).await? { - Some(plan) => Ok(plan), - None => { - let plan = self - .create_initial_plan(logical_plan, session_state) - .await?; - - self.optimize_physical_plan(plan, session_state, |_, _| {}) - } + if let Some(plan) = self + .handle_explain_or_analyze(logical_plan, session_state) + .await? + { + return Ok(plan); } + let plan = self + .create_initial_plan(logical_plan, session_state) + .await?; + + self.optimize_physical_plan(plan, session_state, |_, _| {}) } /// Create a physical expression from a logical expression @@ -1713,147 +1717,179 @@ impl DefaultPhysicalPlanner { /// Returns /// Some(plan) if optimized, and None if logical_plan was not an /// explain (and thus needs to be optimized as normal) - async fn handle_explain( + async fn handle_explain_or_analyze( &self, logical_plan: &LogicalPlan, session_state: &SessionState, ) -> Result>> { - if let LogicalPlan::Explain(e) = logical_plan { - use PlanType::*; - let mut stringified_plans = vec![]; + let execution_plan = match logical_plan { + LogicalPlan::Explain(e) => self.handle_explain(e, session_state).await?, + LogicalPlan::Analyze(a) => self.handle_analyze(a, session_state).await?, + _ => return Ok(None), + }; + Ok(Some(execution_plan)) + } - let config = &session_state.config_options().explain; + /// Planner for `LogicalPlan::Explain` + async fn handle_explain( + &self, + e: &Explain, + session_state: &SessionState, + ) -> Result> { + use PlanType::*; + let mut stringified_plans = vec![]; - if !config.physical_plan_only { - stringified_plans.clone_from(&e.stringified_plans); - if e.logical_optimization_succeeded { - stringified_plans.push(e.plan.to_stringified(FinalLogicalPlan)); - } + let config = &session_state.config_options().explain; + let explain_format = DisplayFormatType::from_str(&config.format)?; + + let skip_logical_plan = + config.physical_plan_only || explain_format == DisplayFormatType::TreeRender; + + if !skip_logical_plan { + stringified_plans.clone_from(&e.stringified_plans); + if e.logical_optimization_succeeded { + stringified_plans.push(e.plan.to_stringified(FinalLogicalPlan)); } + } - if !config.logical_plan_only && e.logical_optimization_succeeded { - match self - .create_initial_plan(e.plan.as_ref(), session_state) - .await - { - Ok(input) => { - // Include statistics / schema if enabled - stringified_plans.push( - displayable(input.as_ref()) - .set_show_statistics(config.show_statistics) - .set_show_schema(config.show_schema) - .to_stringified(e.verbose, InitialPhysicalPlan), - ); + if !config.logical_plan_only && e.logical_optimization_succeeded { + match self + .create_initial_plan(e.plan.as_ref(), session_state) + .await + { + Ok(input) => { + // Include statistics / schema if enabled + stringified_plans.push( + displayable(input.as_ref()) + .set_show_statistics(config.show_statistics) + .set_show_schema(config.show_schema) + .to_stringified( + e.verbose, + InitialPhysicalPlan, + explain_format, + ), + ); - // Show statistics + schema in verbose output even if not - // explicitly requested - if e.verbose { - if !config.show_statistics { - stringified_plans.push( - displayable(input.as_ref()) - .set_show_statistics(true) - .to_stringified( - e.verbose, - InitialPhysicalPlanWithStats, - ), - ); - } - if !config.show_schema { - stringified_plans.push( - displayable(input.as_ref()) - .set_show_schema(true) - .to_stringified( - e.verbose, - InitialPhysicalPlanWithSchema, - ), - ); - } + // Show statistics + schema in verbose output even if not + // explicitly requested + if e.verbose { + if !config.show_statistics { + stringified_plans.push( + displayable(input.as_ref()) + .set_show_statistics(true) + .to_stringified( + e.verbose, + InitialPhysicalPlanWithStats, + explain_format, + ), + ); } + if !config.show_schema { + stringified_plans.push( + displayable(input.as_ref()) + .set_show_schema(true) + .to_stringified( + e.verbose, + InitialPhysicalPlanWithSchema, + explain_format, + ), + ); + } + } - let optimized_plan = self.optimize_physical_plan( - input, - session_state, - |plan, optimizer| { - let optimizer_name = optimizer.name().to_string(); - let plan_type = OptimizedPhysicalPlan { optimizer_name }; - stringified_plans.push( - displayable(plan) - .set_show_statistics(config.show_statistics) - .set_show_schema(config.show_schema) - .to_stringified(e.verbose, plan_type), - ); - }, - ); - match optimized_plan { - Ok(input) => { - // This plan will includes statistics if show_statistics is on - stringified_plans.push( - displayable(input.as_ref()) - .set_show_statistics(config.show_statistics) - .set_show_schema(config.show_schema) - .to_stringified(e.verbose, FinalPhysicalPlan), - ); - - // Show statistics + schema in verbose output even if not - // explicitly requested - if e.verbose { - if !config.show_statistics { - stringified_plans.push( - displayable(input.as_ref()) - .set_show_statistics(true) - .to_stringified( - e.verbose, - FinalPhysicalPlanWithStats, - ), - ); - } - if !config.show_schema { - stringified_plans.push( - displayable(input.as_ref()) - .set_show_schema(true) - .to_stringified( - e.verbose, - FinalPhysicalPlanWithSchema, - ), - ); - } + let optimized_plan = self.optimize_physical_plan( + input, + session_state, + |plan, optimizer| { + let optimizer_name = optimizer.name().to_string(); + let plan_type = OptimizedPhysicalPlan { optimizer_name }; + stringified_plans.push( + displayable(plan) + .set_show_statistics(config.show_statistics) + .set_show_schema(config.show_schema) + .to_stringified(e.verbose, plan_type, explain_format), + ); + }, + ); + match optimized_plan { + Ok(input) => { + // This plan will includes statistics if show_statistics is on + stringified_plans.push( + displayable(input.as_ref()) + .set_show_statistics(config.show_statistics) + .set_show_schema(config.show_schema) + .to_stringified( + e.verbose, + FinalPhysicalPlan, + explain_format, + ), + ); + + // Show statistics + schema in verbose output even if not + // explicitly requested + if e.verbose { + if !config.show_statistics { + stringified_plans.push( + displayable(input.as_ref()) + .set_show_statistics(true) + .to_stringified( + e.verbose, + FinalPhysicalPlanWithStats, + explain_format, + ), + ); + } + if !config.show_schema { + stringified_plans.push( + displayable(input.as_ref()) + .set_show_schema(true) + .to_stringified( + e.verbose, + FinalPhysicalPlanWithSchema, + explain_format, + ), + ); } } - Err(DataFusionError::Context(optimizer_name, e)) => { - let plan_type = OptimizedPhysicalPlan { optimizer_name }; - stringified_plans - .push(StringifiedPlan::new(plan_type, e.to_string())) - } - Err(e) => return Err(e), } + Err(DataFusionError::Context(optimizer_name, e)) => { + let plan_type = OptimizedPhysicalPlan { optimizer_name }; + stringified_plans + .push(StringifiedPlan::new(plan_type, e.to_string())) + } + Err(e) => return Err(e), } - Err(err) => { - stringified_plans.push(StringifiedPlan::new( - PhysicalPlanError, - err.strip_backtrace(), - )); - } + } + Err(err) => { + stringified_plans.push(StringifiedPlan::new( + PhysicalPlanError, + err.strip_backtrace(), + )); } } - - Ok(Some(Arc::new(ExplainExec::new( - SchemaRef::new(e.schema.as_ref().to_owned().into()), - stringified_plans, - e.verbose, - )))) - } else if let LogicalPlan::Analyze(a) = logical_plan { - let input = self.create_physical_plan(&a.input, session_state).await?; - let schema = SchemaRef::new((*a.schema).clone().into()); - let show_statistics = session_state.config_options().explain.show_statistics; - Ok(Some(Arc::new(AnalyzeExec::new( - a.verbose, - show_statistics, - input, - schema, - )))) - } else { - Ok(None) } + + Ok(Arc::new(ExplainExec::new( + Arc::clone(e.schema.inner()), + stringified_plans, + e.verbose, + ))) + } + + async fn handle_analyze( + &self, + a: &Analyze, + session_state: &SessionState, + ) -> Result> { + let input = self.create_physical_plan(&a.input, session_state).await?; + let schema = SchemaRef::new((*a.schema).clone().into()); + let show_statistics = session_state.config_options().explain.show_statistics; + Ok(Arc::new(AnalyzeExec::new( + a.verbose, + show_statistics, + input, + schema, + ))) } /// Optimize a physical plan by applying each physical optimizer, @@ -2720,6 +2756,10 @@ mod tests { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "NoOpExecutionPlan") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/src/test/mod.rs b/datafusion/core/src/test/mod.rs index 23885c41c61a..be707f7e19d0 100644 --- a/datafusion/core/src/test/mod.rs +++ b/datafusion/core/src/test/mod.rs @@ -29,30 +29,29 @@ use crate::datasource::file_format::csv::CsvFormat; use crate::datasource::file_format::file_compression_type::FileCompressionType; use crate::datasource::file_format::FileFormat; use crate::datasource::listing::PartitionedFile; -use crate::datasource::object_store::ObjectStoreUrl; + use crate::datasource::physical_plan::CsvSource; use crate::datasource::{MemTable, TableProvider}; use crate::error::Result; use crate::logical_expr::LogicalPlan; -use crate::test::object_store::local_unpartitioned_file; use crate::test_util::{aggr_test_schema, arrow_test_data}; use arrow::array::{self, Array, ArrayRef, Decimal128Builder, Int32Array}; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::DataFusionError; -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::source::DataSourceExec; #[cfg(feature = "compression")] use bzip2::write::BzEncoder; #[cfg(feature = "compression")] use bzip2::Compression as BzCompression; +use datafusion_datasource_csv::partitioned_csv_config; #[cfg(feature = "compression")] use flate2::write::GzEncoder; #[cfg(feature = "compression")] use flate2::Compression as GzCompression; +use object_store::local_unpartitioned_file; #[cfg(feature = "compression")] use xz2::write::XzEncoder; #[cfg(feature = "compression")] @@ -186,16 +185,6 @@ pub fn partitioned_file_groups( .collect::>()) } -/// Returns a [`FileScanConfig`] for given `file_groups` -pub fn partitioned_csv_config( - schema: SchemaRef, - file_groups: Vec>, - file_source: Arc, -) -> FileScanConfig { - FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, file_source) - .with_file_groups(file_groups) -} - pub fn assert_fields_eq(plan: &LogicalPlan, expected: Vec<&str>) { let actual: Vec = plan .schema() diff --git a/datafusion/core/src/test/object_store.rs b/datafusion/core/src/test/object_store.rs index cac430c5b49d..e1328770cabd 100644 --- a/datafusion/core/src/test/object_store.rs +++ b/datafusion/core/src/test/object_store.rs @@ -20,12 +20,22 @@ use crate::execution::context::SessionState; use crate::execution::session_state::SessionStateBuilder; use crate::prelude::SessionContext; +use futures::stream::BoxStream; use futures::FutureExt; -use object_store::{memory::InMemory, path::Path, ObjectMeta, ObjectStore}; +use object_store::{ + memory::InMemory, path::Path, Error, GetOptions, GetResult, ListResult, + MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOpts, PutOptions, PutPayload, + PutResult, +}; +use std::fmt::{Debug, Display, Formatter}; use std::sync::Arc; +use tokio::{ + sync::Barrier, + time::{timeout, Duration}, +}; use url::Url; -/// Returns a test object store with the provided `ctx` +/// Registers a test object store with the provided `ctx` pub fn register_test_store(ctx: &SessionContext, files: &[(&str, u64)]) { let url = Url::parse("test://").unwrap(); ctx.register_object_store(&url, make_test_store_and_state(files).0); @@ -61,3 +71,121 @@ pub fn local_unpartitioned_file(path: impl AsRef) -> ObjectMeta version: None, } } + +/// Blocks the object_store `head` call until `concurrency` number of calls are pending. +pub fn ensure_head_concurrency( + object_store: Arc, + concurrency: usize, +) -> Arc { + Arc::new(BlockingObjectStore::new(object_store, concurrency)) +} + +/// An object store that “blocks” in its `head` call until an expected number of concurrent calls are reached. +#[derive(Debug)] +struct BlockingObjectStore { + inner: Arc, + barrier: Arc, +} + +impl BlockingObjectStore { + const NAME: &'static str = "BlockingObjectStore"; + fn new(inner: Arc, expected_concurrency: usize) -> Self { + Self { + inner, + barrier: Arc::new(Barrier::new(expected_concurrency)), + } + } +} + +impl Display for BlockingObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(&self.inner, f) + } +} + +/// All trait methods are forwarded to the inner object store, except for +/// the `head` method which waits until the expected number of concurrent calls is reached. +#[async_trait::async_trait] +impl ObjectStore for BlockingObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + self.inner.put_opts(location, payload, opts).await + } + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOpts, + ) -> object_store::Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + self.inner.get_opts(location, options).await + } + + async fn head(&self, location: &Path) -> object_store::Result { + println!( + "{} received head call for {location}", + BlockingObjectStore::NAME + ); + // Wait until the expected number of concurrent calls is reached, but timeout after 1 second to avoid hanging failing tests. + let wait_result = timeout(Duration::from_secs(1), self.barrier.wait()).await; + match wait_result { + Ok(_) => println!( + "{} barrier reached for {location}", + BlockingObjectStore::NAME + ), + Err(_) => { + let error_message = format!( + "{} barrier wait timed out for {location}", + BlockingObjectStore::NAME + ); + log::error!("{}", error_message); + return Err(Error::Generic { + store: BlockingObjectStore::NAME, + source: error_message.into(), + }); + } + } + // Forward the call to the inner object store. + self.inner.head(location).await + } + + async fn delete(&self, location: &Path) -> object_store::Result<()> { + self.inner.delete(location).await + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'_, object_store::Result> { + self.inner.list(prefix) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> object_store::Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy(&self, from: &Path, to: &Path) -> object_store::Result<()> { + self.inner.copy(from, to).await + } + + async fn copy_if_not_exists( + &self, + from: &Path, + to: &Path, + ) -> object_store::Result<()> { + self.inner.copy_if_not_exists(from, to).await + } +} diff --git a/datafusion/core/src/test_util/mod.rs b/datafusion/core/src/test_util/mod.rs index 50e33b27e1bb..d6865ca3d532 100644 --- a/datafusion/core/src/test_util/mod.rs +++ b/datafusion/core/src/test_util/mod.rs @@ -27,9 +27,7 @@ use std::collections::HashMap; use std::fs::File; use std::io::Write; use std::path::Path; -use std::pin::Pin; use std::sync::Arc; -use std::task::{Context, Poll}; use crate::catalog::{TableProvider, TableProviderFactory}; use crate::dataframe::DataFrame; @@ -37,7 +35,7 @@ use crate::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use crate::datasource::{empty::EmptyTable, provider_as_source}; use crate::error::Result; use crate::logical_expr::{LogicalPlanBuilder, UNNAMED_TABLE}; -use crate::physical_plan::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; +use crate::physical_plan::ExecutionPlan; use crate::prelude::{CsvReadOptions, SessionContext}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -47,7 +45,7 @@ use datafusion_common::TableReference; use datafusion_expr::{CreateExternalTable, Expr, SortExpr, TableType}; use async_trait::async_trait; -use futures::Stream; + use tempfile::TempDir; // backwards compatibility #[cfg(feature = "parquet")] @@ -236,39 +234,3 @@ pub fn register_unbounded_file_with_ordering( ctx.register_table(table_name, Arc::new(StreamTable::new(Arc::new(config))))?; Ok(()) } - -struct BoundedStream { - limit: usize, - count: usize, - batch: RecordBatch, -} - -impl Stream for BoundedStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - if self.count >= self.limit { - return Poll::Ready(None); - } - self.count += 1; - Poll::Ready(Some(Ok(self.batch.clone()))) - } -} - -impl RecordBatchStream for BoundedStream { - fn schema(&self) -> SchemaRef { - self.batch.schema() - } -} - -/// Creates an bounded stream for testing purposes. -pub fn bounded_stream(batch: RecordBatch, limit: usize) -> SendableRecordBatchStream { - Box::pin(BoundedStream { - count: 0, - limit, - batch, - }) -} diff --git a/datafusion/core/tests/core_integration.rs b/datafusion/core/tests/core_integration.rs index 66b4103160e7..9bcb9e41f86a 100644 --- a/datafusion/core/tests/core_integration.rs +++ b/datafusion/core/tests/core_integration.rs @@ -42,8 +42,13 @@ mod custom_sources_cases; /// Run all tests that are found in the `optimizer` directory mod optimizer; +/// Run all tests that are found in the `physical_optimizer` directory mod physical_optimizer; +/// Run all tests that are found in the `serde` directory +mod serde; + +/// Run all tests that are found in the `catalog` directory mod catalog; #[cfg(test)] diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index aafefac04e32..eb930b9a60bc 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -138,6 +138,11 @@ impl DisplayAs for CustomExecutionPlan { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "CustomExecutionPlan: projection={:#?}", self.projection) } + + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index af0506a50558..f68bcfaf1550 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -92,6 +92,10 @@ impl DisplayAs for CustomPlan { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "CustomPlan: batch_size={}", self.batches.len(),) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index 1fd6dfec79fb..66c886510e96 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -141,6 +141,10 @@ impl DisplayAs for StatisticsValidation { self.stats.num_rows, ) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/tests/dataframe/dataframe_functions.rs b/datafusion/core/tests/dataframe/dataframe_functions.rs index 28c0740ca76b..fec3ab786fce 100644 --- a/datafusion/core/tests/dataframe/dataframe_functions.rs +++ b/datafusion/core/tests/dataframe/dataframe_functions.rs @@ -1145,9 +1145,9 @@ async fn test_count_wildcard() -> Result<()> { .build() .unwrap(); - let expected = "Sort: count(Int64(1)) ASC NULLS LAST [count(Int64(1)):Int64]\ - \n Projection: count(Int64(1)) [count(Int64(1)):Int64]\ - \n Aggregate: groupBy=[[test.b]], aggr=[[count(Int64(1))]] [b:UInt32, count(Int64(1)):Int64]\ + let expected = "Sort: count(*) ASC NULLS LAST [count(*):Int64]\ + \n Projection: count(*) [count(*):Int64]\ + \n Aggregate: groupBy=[[test.b]], aggr=[[count(Int64(1)) AS count(*)]] [b:UInt32, count(*):Int64]\ \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; let formatted_plan = plan.display_indent_schema().to_string(); diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index b134ec54b13d..a902cf8ae65b 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -32,8 +32,7 @@ use arrow::datatypes::{ }; use arrow::error::ArrowError; use arrow::util::pretty::pretty_format_batches; -use datafusion_expr::utils::COUNT_STAR_EXPANSION; -use datafusion_functions_aggregate::count::{count_all, count_udaf}; +use datafusion_functions_aggregate::count::{count_all, count_all_window}; use datafusion_functions_aggregate::expr_fn::{ array_agg, avg, count, count_distinct, max, median, min, sum, }; @@ -54,7 +53,7 @@ use datafusion::execution::context::SessionContext; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::logical_expr::{ColumnarValue, Volatility}; use datafusion::prelude::{ - AvroReadOptions, CsvReadOptions, JoinType, NdJsonReadOptions, ParquetReadOptions, + CsvReadOptions, JoinType, NdJsonReadOptions, ParquetReadOptions, }; use datafusion::test_util::{ parquet_test_data, populate_csv_partitions, register_aggregate_csv, test_table, @@ -2455,7 +2454,7 @@ async fn test_count_wildcard_on_sort() -> Result<()> { let ctx = create_join_context()?; let sql_results = ctx - .sql("select b,count(1) from t1 group by b order by count(1)") + .sql("select b, count(*) from t1 group by b order by count(*)") .await? .explain(false, false)? .collect() @@ -2469,9 +2468,52 @@ async fn test_count_wildcard_on_sort() -> Result<()> { .explain(false, false)? .collect() .await?; - //make sure sql plan same with df plan + + let expected_sql_result = "+---------------+------------------------------------------------------------------------------------------------------------+\ + \n| plan_type | plan |\ + \n+---------------+------------------------------------------------------------------------------------------------------------+\ + \n| logical_plan | Projection: t1.b, count(*) |\ + \n| | Sort: count(Int64(1)) AS count(*) AS count(*) ASC NULLS LAST |\ + \n| | Projection: t1.b, count(Int64(1)) AS count(*), count(Int64(1)) |\ + \n| | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1))]] |\ + \n| | TableScan: t1 projection=[b] |\ + \n| physical_plan | ProjectionExec: expr=[b@0 as b, count(*)@1 as count(*)] |\ + \n| | SortPreservingMergeExec: [count(Int64(1))@2 ASC NULLS LAST] |\ + \n| | SortExec: expr=[count(Int64(1))@2 ASC NULLS LAST], preserve_partitioning=[true] |\ + \n| | ProjectionExec: expr=[b@0 as b, count(Int64(1))@1 as count(*), count(Int64(1))@1 as count(Int64(1))] |\ + \n| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(Int64(1))] |\ + \n| | CoalesceBatchesExec: target_batch_size=8192 |\ + \n| | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=4 |\ + \n| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 |\ + \n| | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(Int64(1))] |\ + \n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ + \n| | |\ + \n+---------------+------------------------------------------------------------------------------------------------------------+"; + + assert_eq!( + expected_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + + let expected_df_result = "+---------------+--------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+--------------------------------------------------------------------------------+\ +\n| logical_plan | Sort: count(*) ASC NULLS LAST |\ +\n| | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1)) AS count(*)]] |\ +\n| | TableScan: t1 projection=[b] |\ +\n| physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST] |\ +\n| | SortExec: expr=[count(*)@1 ASC NULLS LAST], preserve_partitioning=[true] |\ +\n| | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(*)] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=4 |\ +\n| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 |\ +\n| | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(*)] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+--------------------------------------------------------------------------------+"; + assert_eq!( - pretty_format_batches(&sql_results)?.to_string(), + expected_df_result, pretty_format_batches(&df_results)?.to_string() ); Ok(()) @@ -2481,12 +2523,35 @@ async fn test_count_wildcard_on_sort() -> Result<()> { async fn test_count_wildcard_on_where_in() -> Result<()> { let ctx = create_join_context()?; let sql_results = ctx - .sql("SELECT a,b FROM t1 WHERE a in (SELECT count(1) FROM t2)") + .sql("SELECT a, b FROM t1 WHERE a in (SELECT count(*) FROM t2)") .await? .explain(false, false)? .collect() .await?; + let expected_sql_result = "+---------------+------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | LeftSemi Join: CAST(t1.a AS Int64) = __correlated_sq_1.count(*) |\ +\n| | TableScan: t1 projection=[a, b] |\ +\n| | SubqueryAlias: __correlated_sq_1 |\ +\n| | Projection: count(Int64(1)) AS count(*) |\ +\n| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] |\ +\n| | TableScan: t2 projection=[] |\ +\n| physical_plan | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(count(*)@0, CAST(t1.a AS Int64)@2)], projection=[a@0, b@1] |\ +\n| | ProjectionExec: expr=[4 as count(*)] |\ +\n| | PlaceholderRowExec |\ +\n| | ProjectionExec: expr=[a@0 as a, b@1 as b, CAST(a@0 AS Int64) as CAST(t1.a AS Int64)] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+------------------------------------------------------------------------------------------------------------------------+"; + + assert_eq!( + expected_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + // In the same SessionContext, AliasGenerator will increase subquery_alias id by 1 // https://github.com/apache/datafusion/blame/cf45eb9020092943b96653d70fafb143cc362e19/datafusion/optimizer/src/alias.rs#L40-L43 // for compare difference between sql and df logical plan, we need to create a new SessionContext here @@ -2509,9 +2574,26 @@ async fn test_count_wildcard_on_where_in() -> Result<()> { .collect() .await?; + let actual_df_result= "+---------------+------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | LeftSemi Join: CAST(t1.a AS Int64) = __correlated_sq_1.count(*) |\ +\n| | TableScan: t1 projection=[a, b] |\ +\n| | SubqueryAlias: __correlated_sq_1 |\ +\n| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] |\ +\n| | TableScan: t2 projection=[] |\ +\n| physical_plan | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(count(*)@0, CAST(t1.a AS Int64)@2)], projection=[a@0, b@1] |\ +\n| | ProjectionExec: expr=[4 as count(*)] |\ +\n| | PlaceholderRowExec |\ +\n| | ProjectionExec: expr=[a@0 as a, b@1 as b, CAST(a@0 AS Int64) as CAST(t1.a AS Int64)] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+------------------------------------------------------------------------------------------------------------------------+"; + // make sure sql plan same with df plan assert_eq!( - pretty_format_batches(&sql_results)?.to_string(), + actual_df_result, pretty_format_batches(&df_results)?.to_string() ); @@ -2522,11 +2604,34 @@ async fn test_count_wildcard_on_where_in() -> Result<()> { async fn test_count_wildcard_on_where_exist() -> Result<()> { let ctx = create_join_context()?; let sql_results = ctx - .sql("SELECT a, b FROM t1 WHERE EXISTS (SELECT count(1) FROM t2)") + .sql("SELECT a, b FROM t1 WHERE EXISTS (SELECT count(*) FROM t2)") .await? .explain(false, false)? .collect() .await?; + + let actual_sql_result = + "+---------------+---------------------------------------------------------+\ + \n| plan_type | plan |\ + \n+---------------+---------------------------------------------------------+\ + \n| logical_plan | LeftSemi Join: |\ + \n| | TableScan: t1 projection=[a, b] |\ + \n| | SubqueryAlias: __correlated_sq_1 |\ + \n| | Projection: |\ + \n| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] |\ + \n| | TableScan: t2 projection=[] |\ + \n| physical_plan | NestedLoopJoinExec: join_type=RightSemi |\ + \n| | ProjectionExec: expr=[] |\ + \n| | PlaceholderRowExec |\ + \n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ + \n| | |\ + \n+---------------+---------------------------------------------------------+"; + + assert_eq!( + actual_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + let df_results = ctx .table("t1") .await? @@ -2545,9 +2650,24 @@ async fn test_count_wildcard_on_where_exist() -> Result<()> { .collect() .await?; - //make sure sql plan same with df plan + let actual_df_result = "+---------------+---------------------------------------------------------------------+\ + \n| plan_type | plan |\ + \n+---------------+---------------------------------------------------------------------+\ + \n| logical_plan | LeftSemi Join: |\ + \n| | TableScan: t1 projection=[a, b] |\ + \n| | SubqueryAlias: __correlated_sq_1 |\ + \n| | Projection: |\ + \n| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] |\ + \n| | TableScan: t2 projection=[] |\ + \n| physical_plan | NestedLoopJoinExec: join_type=RightSemi |\ + \n| | ProjectionExec: expr=[] |\ + \n| | PlaceholderRowExec |\ + \n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ + \n| | |\ + \n+---------------+---------------------------------------------------------------------+"; + assert_eq!( - pretty_format_batches(&sql_results)?.to_string(), + actual_df_result, pretty_format_batches(&df_results)?.to_string() ); @@ -2559,34 +2679,62 @@ async fn test_count_wildcard_on_window() -> Result<()> { let ctx = create_join_context()?; let sql_results = ctx - .sql("select count(1) OVER(ORDER BY a DESC RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING) from t1") + .sql("select count(*) OVER(ORDER BY a DESC RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING) from t1") .await? .explain(false, false)? .collect() .await?; + + let actual_sql_result = "+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | Projection: count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING AS count(*) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING |\ +\n| | WindowAggr: windowExpr=[[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING]] |\ +\n| | TableScan: t1 projection=[a] |\ +\n| physical_plan | ProjectionExec: expr=[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING@1 as count(*) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING] |\ +\n| | BoundedWindowAggExec: wdw=[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING: Ok(Field { name: \"count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(UInt32(6)), end_bound: Following(UInt32(2)), is_causal: false }], mode=[Sorted] |\ +\n| | SortExec: expr=[a@0 DESC], preserve_partitioning=[false] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+"; + + assert_eq!( + actual_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + let df_results = ctx .table("t1") .await? - .select(vec![Expr::WindowFunction(WindowFunction::new( - WindowFunctionDefinition::AggregateUDF(count_udaf()), - vec![Expr::Literal(COUNT_STAR_EXPANSION)], - )) - .order_by(vec![Sort::new(col("a"), false, true)]) - .window_frame(WindowFrame::new_bounds( - WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::UInt32(Some(6))), - WindowFrameBound::Following(ScalarValue::UInt32(Some(2))), - )) - .build() - .unwrap()])? + .select(vec![count_all_window() + .order_by(vec![Sort::new(col("a"), false, true)]) + .window_frame(WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::UInt32(Some(6))), + WindowFrameBound::Following(ScalarValue::UInt32(Some(2))), + )) + .build() + .unwrap()])? .explain(false, false)? .collect() .await?; - //make sure sql plan same with df plan + let actual_df_result = "+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | Projection: count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING |\ +\n| | WindowAggr: windowExpr=[[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING]] |\ +\n| | TableScan: t1 projection=[a] |\ +\n| physical_plan | ProjectionExec: expr=[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING@1 as count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING] |\ +\n| | BoundedWindowAggExec: wdw=[count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING: Ok(Field { name: \"count(Int64(1)) ORDER BY [t1.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(UInt32(6)), end_bound: Following(UInt32(2)), is_causal: false }], mode=[Sorted] |\ +\n| | SortExec: expr=[a@0 DESC], preserve_partitioning=[false] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+"; + assert_eq!( - pretty_format_batches(&df_results)?.to_string(), - pretty_format_batches(&sql_results)?.to_string() + actual_df_result, + pretty_format_batches(&df_results)?.to_string() ); Ok(()) @@ -2598,12 +2746,28 @@ async fn test_count_wildcard_on_aggregate() -> Result<()> { register_alltypes_tiny_pages_parquet(&ctx).await?; let sql_results = ctx - .sql("select count(1) from t1") + .sql("select count(*) from t1") .await? .explain(false, false)? .collect() .await?; + let actual_sql_result = + "+---------------+-----------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+-----------------------------------------------------+\ +\n| logical_plan | Projection: count(Int64(1)) AS count(*) |\ +\n| | Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] |\ +\n| | TableScan: t1 projection=[] |\ +\n| physical_plan | ProjectionExec: expr=[4 as count(*)] |\ +\n| | PlaceholderRowExec |\ +\n| | |\ +\n+---------------+-----------------------------------------------------+"; + assert_eq!( + actual_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + // add `.select(vec![count_wildcard()])?` to make sure we can analyze all node instead of just top node. let df_results = ctx .table("t1") @@ -2614,9 +2778,17 @@ async fn test_count_wildcard_on_aggregate() -> Result<()> { .collect() .await?; - //make sure sql plan same with df plan + let actual_df_result = "+---------------+---------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+---------------------------------------------------------------+\ +\n| logical_plan | Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] |\ +\n| | TableScan: t1 projection=[] |\ +\n| physical_plan | ProjectionExec: expr=[4 as count(*)] |\ +\n| | PlaceholderRowExec |\ +\n| | |\ +\n+---------------+---------------------------------------------------------------+"; assert_eq!( - pretty_format_batches(&sql_results)?.to_string(), + actual_df_result, pretty_format_batches(&df_results)?.to_string() ); @@ -2628,16 +2800,51 @@ async fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> { let ctx = create_join_context()?; let sql_results = ctx - .sql("select a,b from t1 where (select count(1) from t2 where t1.a = t2.a)>0;") + .sql("select a,b from t1 where (select count(*) from t2 where t1.a = t2.a)>0;") .await? .explain(false, false)? .collect() .await?; + let actual_sql_result = "+---------------+---------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | Projection: t1.a, t1.b |\ +\n| | Filter: CASE WHEN __scalar_sq_1.__always_true IS NULL THEN Int64(0) ELSE __scalar_sq_1.count(*) END > Int64(0) |\ +\n| | Projection: t1.a, t1.b, __scalar_sq_1.count(*), __scalar_sq_1.__always_true |\ +\n| | Left Join: t1.a = __scalar_sq_1.a |\ +\n| | TableScan: t1 projection=[a, b] |\ +\n| | SubqueryAlias: __scalar_sq_1 |\ +\n| | Projection: count(Int64(1)) AS count(*), t2.a, Boolean(true) AS __always_true |\ +\n| | Aggregate: groupBy=[[t2.a]], aggr=[[count(Int64(1))]] |\ +\n| | TableScan: t2 projection=[a] |\ +\n| physical_plan | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | FilterExec: CASE WHEN __always_true@3 IS NULL THEN 0 ELSE count(*)@2 END > 0, projection=[a@0, b@1] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | HashJoinExec: mode=Partitioned, join_type=Left, on=[(a@0, a@1)], projection=[a@0, b@1, count(*)@2, __always_true@4] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | ProjectionExec: expr=[count(Int64(1))@1 as count(*), a@0 as a, true as __always_true] |\ +\n| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(Int64(1))] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 |\ +\n| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 |\ +\n| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(Int64(1))] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------+"; + assert_eq!( + actual_sql_result, + pretty_format_batches(&sql_results)?.to_string() + ); + // In the same SessionContext, AliasGenerator will increase subquery_alias id by 1 // https://github.com/apache/datafusion/blame/cf45eb9020092943b96653d70fafb143cc362e19/datafusion/optimizer/src/alias.rs#L40-L43 // for compare difference between sql and df logical plan, we need to create a new SessionContext here let ctx = create_join_context()?; + let agg_expr = count_all(); + let agg_expr_col = col(agg_expr.schema_name().to_string()); let df_results = ctx .table("t1") .await? @@ -2646,8 +2853,8 @@ async fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> { ctx.table("t2") .await? .filter(out_ref_col(DataType::UInt32, "t1.a").eq(col("t2.a")))? - .aggregate(vec![], vec![count_all()])? - .select(vec![col(count_all().to_string())])? + .aggregate(vec![], vec![agg_expr])? + .select(vec![agg_expr_col])? .into_unoptimized_plan(), )) .gt(lit(ScalarValue::UInt8(Some(0)))), @@ -2657,9 +2864,36 @@ async fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> { .collect() .await?; - //make sure sql plan same with df plan + let actual_df_result = "+---------------+---------------------------------------------------------------------------------------------------------------------------+\ +\n| plan_type | plan |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------+\ +\n| logical_plan | Projection: t1.a, t1.b |\ +\n| | Filter: CASE WHEN __scalar_sq_1.__always_true IS NULL THEN Int64(0) ELSE __scalar_sq_1.count(*) END > Int64(0) |\ +\n| | Projection: t1.a, t1.b, __scalar_sq_1.count(*), __scalar_sq_1.__always_true |\ +\n| | Left Join: t1.a = __scalar_sq_1.a |\ +\n| | TableScan: t1 projection=[a, b] |\ +\n| | SubqueryAlias: __scalar_sq_1 |\ +\n| | Projection: count(*), t2.a, Boolean(true) AS __always_true |\ +\n| | Aggregate: groupBy=[[t2.a]], aggr=[[count(Int64(1)) AS count(*)]] |\ +\n| | TableScan: t2 projection=[a] |\ +\n| physical_plan | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | FilterExec: CASE WHEN __always_true@3 IS NULL THEN 0 ELSE count(*)@2 END > 0, projection=[a@0, b@1] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | HashJoinExec: mode=Partitioned, join_type=Left, on=[(a@0, a@1)], projection=[a@0, b@1, count(*)@2, __always_true@4] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | ProjectionExec: expr=[count(*)@1 as count(*), a@0 as a, true as __always_true] |\ +\n| | AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(*)] |\ +\n| | CoalesceBatchesExec: target_batch_size=8192 |\ +\n| | RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 |\ +\n| | RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 |\ +\n| | AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(*)] |\ +\n| | DataSourceExec: partitions=1, partition_sizes=[1] |\ +\n| | |\ +\n+---------------+---------------------------------------------------------------------------------------------------------------------------+"; assert_eq!( - pretty_format_batches(&sql_results)?.to_string(), + actual_df_result, pretty_format_batches(&df_results)?.to_string() ); @@ -4228,7 +4462,9 @@ fn create_join_context() -> Result { ], )?; - let ctx = SessionContext::new(); + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config(config); + // let ctx = SessionContext::new(); ctx.register_batch("t1", batch1)?; ctx.register_batch("t2", batch2)?; @@ -5263,6 +5499,7 @@ async fn register_non_csv_file() { ); } +#[cfg(feature = "avro")] #[tokio::test] async fn register_non_avro_file() { let ctx = SessionContext::new(); @@ -5270,7 +5507,7 @@ async fn register_non_avro_file() { .register_avro( "data", "tests/data/test_binary.parquet", - AvroReadOptions::default(), + datafusion::prelude::AvroReadOptions::default(), ) .await; assert_contains!( @@ -5436,3 +5673,75 @@ async fn test_fill_null_all_columns() -> Result<()> { assert_batches_sorted_eq!(expected, &results); Ok(()) } + +#[tokio::test] +async fn test_insert_into_casting_support() -> Result<()> { + // Testing case1: + // Inserting query schema mismatch: Expected table field 'a' with type Float16, but got 'a' with type Utf8. + // And the cast is not supported from Utf8 to Float16. + + // Create a new schema with one field called "a" of type Float16, and setting nullable to false + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float16, false)])); + + let session_ctx = SessionContext::new(); + + // Create and register the initial table with the provided schema and data + let initial_table = Arc::new(MemTable::try_new(schema.clone(), vec![vec![]])?); + session_ctx.register_table("t", initial_table.clone())?; + + let mut write_df = session_ctx.sql("values ('a123'), ('b456')").await.unwrap(); + + write_df = write_df + .clone() + .with_column_renamed("column1", "a") + .unwrap(); + + let e = write_df + .write_table("t", DataFrameWriteOptions::new()) + .await + .unwrap_err(); + + assert_contains!(e.to_string(), "Inserting query schema mismatch: Expected table field 'a' with type Float16, but got 'a' with type Utf8."); + + // Testing case2: + // Inserting query schema mismatch: Expected table field 'a' with type Utf8View, but got 'a' with type Utf8. + // And the cast is supported from Utf8 to Utf8View. + + // Create a new schema with one field called "a" of type Utf8View, and setting nullable to false + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Utf8View, + false, + )])); + + let initial_table = Arc::new(MemTable::try_new(schema.clone(), vec![vec![]])?); + + session_ctx.register_table("t2", initial_table.clone())?; + + let mut write_df = session_ctx.sql("values ('a123'), ('b456')").await.unwrap(); + + write_df = write_df + .clone() + .with_column_renamed("column1", "a") + .unwrap(); + + write_df + .write_table("t2", DataFrameWriteOptions::new()) + .await?; + + let res = session_ctx + .sql("select * from t2") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // The result should be the same as the input which is ['a123', 'b456'] + let expected = [ + "+------+", "| a |", "+------+", "| a123 |", "| b456 |", "+------+", + ]; + + assert_batches_eq!(expected, &res); + Ok(()) +} diff --git a/datafusion/core/tests/expr_api/mod.rs b/datafusion/core/tests/expr_api/mod.rs index 7c0119e8ae83..aef10379da07 100644 --- a/datafusion/core/tests/expr_api/mod.rs +++ b/datafusion/core/tests/expr_api/mod.rs @@ -23,11 +23,14 @@ use arrow::datatypes::{DataType, Field}; use arrow::util::pretty::{pretty_format_batches, pretty_format_columns}; use datafusion::prelude::*; use datafusion_common::{DFSchema, ScalarValue}; +use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::ExprFunctionExt; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::first_last::first_value_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_functions_nested::expr_ext::{IndexAccessor, SliceAccessor}; +use datafusion_optimizer::simplify_expressions::ExprSimplifier; use sqlparser::ast::NullTreatment; /// Tests of using and evaluating `Expr`s outside the context of a LogicalPlan use std::sync::{Arc, LazyLock}; @@ -304,6 +307,37 @@ async fn test_aggregate_ext_null_treatment() { .await; } +#[tokio::test] +async fn test_create_physical_expr() { + // create_physical_expr does not simplify the expression + // 1 + 1 + create_expr_test(lit(1i32) + lit(2i32), "1 + 2"); + // However, you can run the simplifier before creating the physical + // expression. This mimics what delta.rs and other non-sql libraries do to + // create predicates + // + // 1 + 1 + create_simplified_expr_test(lit(1i32) + lit(2i32), "3"); +} + +#[tokio::test] +async fn test_create_physical_expr_coercion() { + // create_physical_expr does apply type coercion and unwrapping in cast + // + // expect the cast on the literals + // compare string function to int `id = 1` + create_expr_test(col("id").eq(lit(1i32)), "id@0 = CAST(1 AS Utf8)"); + create_expr_test(lit(1i32).eq(col("id")), "CAST(1 AS Utf8) = id@0"); + // compare int col to string literal `i = '202410'` + // Note this casts the column (not the field) + create_expr_test(col("i").eq(lit("202410")), "CAST(i@1 AS Utf8) = 202410"); + create_expr_test(lit("202410").eq(col("i")), "202410 = CAST(i@1 AS Utf8)"); + // however, when simplified the casts on i should removed + // https://github.com/apache/datafusion/issues/14944 + create_simplified_expr_test(col("i").eq(lit("202410")), "CAST(i@1 AS Utf8) = 202410"); + create_simplified_expr_test(lit("202410").eq(col("i")), "CAST(i@1 AS Utf8) = 202410"); +} + /// Evaluates the specified expr as an aggregate and compares the result to the /// expected result. async fn evaluate_agg_test(expr: Expr, expected_lines: Vec<&str>) { @@ -350,6 +384,38 @@ fn evaluate_expr_test(expr: Expr, expected_lines: Vec<&str>) { ); } +/// Creates the physical expression from Expr and compares the Debug expression +/// to the expected result. +fn create_expr_test(expr: Expr, expected_expr: &str) { + let batch = &TEST_BATCH; + let df_schema = DFSchema::try_from(batch.schema()).unwrap(); + let physical_expr = SessionContext::new() + .create_physical_expr(expr, &df_schema) + .unwrap(); + + assert_eq!(physical_expr.to_string(), expected_expr); +} + +/// Creates the physical expression from Expr and runs the expr simplifier +fn create_simplified_expr_test(expr: Expr, expected_expr: &str) { + let batch = &TEST_BATCH; + let df_schema = DFSchema::try_from(batch.schema()).unwrap(); + + // Simplify the expression first + let props = ExecutionProps::new(); + let simplify_context = + SimplifyContext::new(&props).with_schema(df_schema.clone().into()); + let simplifier = ExprSimplifier::new(simplify_context).with_max_cycles(10); + let simplified = simplifier.simplify(expr).unwrap(); + create_expr_test(simplified, expected_expr); +} + +/// Returns a Batch with 3 rows and 4 columns: +/// +/// id: Utf8 +/// i: Int64 +/// props: Struct +/// list: List static TEST_BATCH: LazyLock = LazyLock::new(|| { let string_array: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3"])); let int_array: ArrayRef = diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs index 4d4c6aa79357..54c5744c861b 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/data_generator.rs @@ -100,7 +100,28 @@ impl DatasetGeneratorConfig { /// Dataset generator /// -/// It will generate one random [`Dataset`] when `generate` function is called. +/// It will generate random [`Dataset`]s when the `generate` function is called. For each +/// sort key in `sort_keys_set`, an additional sorted dataset will be generated, and the +/// dataset will be chunked into staggered batches. +/// +/// # Example +/// For `DatasetGenerator` with `sort_keys_set = [["a"], ["b"]]`, it will generate 2 +/// datasets. The first one will be sorted by column `a` and get randomly chunked +/// into staggered batches. It might look like the following: +/// ```text +/// a b +/// ---- +/// 1 2 <-- batch 1 +/// 1 1 +/// +/// 2 1 <-- batch 2 +/// +/// 3 3 <-- batch 3 +/// 4 3 +/// 4 1 +/// ``` +/// +/// # Implementation details: /// /// The generation logic in `generate`: /// diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs index 7c5b25e4a0e0..1e42ac1f4b30 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/mod.rs @@ -15,6 +15,26 @@ // specific language governing permissions and limitations // under the License. +//! Fuzzer for aggregation functions +//! +//! The main idea behind aggregate fuzzing is: for aggregation, DataFusion has many +//! specialized implementations for performance. For example, when the group cardinality +//! is high, DataFusion will skip the first stage of two-stage hash aggregation; when +//! the input is ordered by the group key, there is a separate implementation to perform +//! streaming group by. +//! This fuzzer checks the results of different specialized implementations and +//! ensures their results are consistent. The execution path can be controlled by +//! changing the input ordering or by setting related configuration parameters in +//! `SessionContext`. +//! +//! # Architecture +//! - `aggregate_fuzz.rs` includes the entry point for fuzzer runs. +//! - `QueryBuilder` is used to generate candidate queries. +//! - `DatasetGenerator` is used to generate random datasets. +//! - `SessionContextGenerator` is used to generate `SessionContext` with +//! different configuration parameters to control the execution path of aggregate +//! queries. + use arrow::array::RecordBatch; use arrow::util::pretty::pretty_format_batches; use datafusion::prelude::SessionContext; diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 2deb8fde2da6..8f690edc5426 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -468,6 +468,31 @@ async fn test_stringview_external_sort() { let _ = df.collect().await.expect("Query execution failed"); } +/// This test case is for a previously detected bug: +/// When `ExternalSorter` has read all input batches +/// - It has spilled many sorted runs to disk +/// - Its in-memory buffer for batches is almost full +/// The previous implementation will try to merge the spills and in-memory batches +/// together, without spilling the in-memory batches first, causing OOM. +#[tokio::test] +async fn test_in_mem_buffer_almost_full() { + let config = SessionConfig::new() + .with_sort_spill_reservation_bytes(3000000) + .with_target_partitions(1); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(10 * 1024 * 1024))) + .build_arc() + .unwrap(); + + let ctx = SessionContext::new_with_config_rt(config, runtime); + + let query = "select * from generate_series(1,9000000) as t1(v1) order by v1;"; + let df = ctx.sql(query).await.unwrap(); + + // Check not fail + let _ = df.collect().await.unwrap(); +} + /// Run the query with the specified memory limit, /// and verifies the expected errors are returned #[derive(Clone, Debug)] diff --git a/datafusion/core/tests/parquet/external_access_plan.rs b/datafusion/core/tests/parquet/external_access_plan.rs index 1eacbe42c525..31c685378a21 100644 --- a/datafusion/core/tests/parquet/external_access_plan.rs +++ b/datafusion/core/tests/parquet/external_access_plan.rs @@ -27,10 +27,10 @@ use arrow::datatypes::SchemaRef; use arrow::util::pretty::pretty_format_batches; use datafusion::common::Result; use datafusion::datasource::listing::PartitionedFile; -use datafusion::datasource::physical_plan::parquet::{ParquetAccessPlan, RowGroupAccess}; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; use datafusion::prelude::SessionContext; use datafusion_common::{assert_contains, DFSchema}; +use datafusion_datasource_parquet::{ParquetAccessPlan, RowGroupAccess}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{col, lit, Expr}; use datafusion_physical_plan::metrics::MetricsSet; diff --git a/datafusion/core/tests/parquet/schema_coercion.rs b/datafusion/core/tests/parquet/schema_coercion.rs index 4cbbcf12f32b..bb20246bf9d5 100644 --- a/datafusion/core/tests/parquet/schema_coercion.rs +++ b/datafusion/core/tests/parquet/schema_coercion.rs @@ -26,10 +26,10 @@ use datafusion::assert_batches_sorted_eq; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; use datafusion::physical_plan::collect; use datafusion::prelude::SessionContext; +use datafusion::test::object_store::local_unpartitioned_file; use datafusion_common::Result; use datafusion_execution::object_store::ObjectStoreUrl; -use object_store::path::Path; use object_store::ObjectMeta; use parquet::arrow::ArrowWriter; use parquet::file::properties::WriterProperties; @@ -168,16 +168,3 @@ pub async fn store_parquet( let meta: Vec<_> = files.iter().map(local_unpartitioned_file).collect(); Ok((meta, files)) } - -/// Helper method to fetch the file size and date at given path and create a `ObjectMeta` -pub fn local_unpartitioned_file(path: impl AsRef) -> ObjectMeta { - let location = Path::from_filesystem_path(path.as_ref()).unwrap(); - let metadata = std::fs::metadata(path).expect("Local file metadata"); - ObjectMeta { - location, - last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), - size: metadata.len() as usize, - e_tag: None, - version: None, - } -} diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index ac1bef6b13d2..b71724b8f7cd 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -19,11 +19,11 @@ use std::fmt::Debug; use std::ops::Deref; use std::sync::Arc; +use crate::physical_optimizer::test_utils::parquet_exec_with_sort; use crate::physical_optimizer::test_utils::{ check_integrity, coalesce_partitions_exec, repartition_exec, schema, sort_merge_join_exec, sort_preserving_merge_exec, }; -use crate::physical_optimizer::test_utils::{parquet_exec_with_sort, trim_plan_display}; use arrow::compute::SortOptions; use datafusion::config::ConfigOptions; @@ -61,7 +61,9 @@ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeE use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::PlanProperties; -use datafusion_physical_plan::{displayable, DisplayAs, DisplayFormatType, Statistics}; +use datafusion_physical_plan::{ + get_plan_string, DisplayAs, DisplayFormatType, Statistics, +}; /// Models operators like BoundedWindowExec that require an input /// ordering but is easy to construct @@ -99,10 +101,18 @@ impl SortRequiredExec { impl DisplayAs for SortRequiredExec { fn fmt_as( &self, - _t: DisplayFormatType, + t: DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - write!(f, "SortRequiredExec: [{}]", self.expr) + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "SortRequiredExec: [{}]", self.expr) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } } } @@ -220,7 +230,7 @@ fn csv_exec_multiple_sorted(output_ordering: Vec) -> Arc, alias_pairs: Vec<(String, String)>, ) -> Arc { @@ -350,8 +360,7 @@ fn ensure_distribution_helper( macro_rules! plans_matches_expected { ($EXPECTED_LINES: expr, $PLAN: expr) => { let physical_plan = $PLAN; - let formatted = displayable(physical_plan.as_ref()).indent(true).to_string(); - let actual: Vec<&str> = formatted.trim().lines().collect(); + let actual = get_plan_string(&physical_plan); let expected_plan_lines: Vec<&str> = $EXPECTED_LINES .iter().map(|s| *s).collect(); @@ -363,65 +372,103 @@ macro_rules! plans_matches_expected { } } -/// Runs the repartition optimizer and asserts the plan against the expected -/// Arguments -/// * `EXPECTED_LINES` - Expected output plan -/// * `PLAN` - Input plan -/// * `FIRST_ENFORCE_DIST` - -/// true: (EnforceDistribution, EnforceDistribution, EnforceSorting) -/// false: else runs (EnforceSorting, EnforceDistribution, EnforceDistribution) -/// * `PREFER_EXISTING_SORT` (optional) - if true, will not repartition / resort data if it is already sorted -/// * `TARGET_PARTITIONS` (optional) - number of partitions to repartition to -/// * `REPARTITION_FILE_SCANS` (optional) - if true, will repartition file scans -/// * `REPARTITION_FILE_MIN_SIZE` (optional) - minimum file size to repartition -/// * `PREFER_EXISTING_UNION` (optional) - if true, will not attempt to convert Union to Interleave -macro_rules! assert_optimized { - ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr) => { - assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, false, 10, false, 1024, false); - }; +fn test_suite_default_config_options() -> ConfigOptions { + let mut config = ConfigOptions::new(); - ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $PREFER_EXISTING_SORT: expr) => { - assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, $PREFER_EXISTING_SORT, 10, false, 1024, false); - }; + // By default, will not repartition / resort data if it is already sorted. + config.optimizer.prefer_existing_sort = false; - ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $PREFER_EXISTING_SORT: expr, $PREFER_EXISTING_UNION: expr) => { - assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, $PREFER_EXISTING_SORT, 10, false, 1024, $PREFER_EXISTING_UNION); - }; + // By default, will attempt to convert Union to Interleave. + config.optimizer.prefer_existing_union = false; - ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $PREFER_EXISTING_SORT: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS: expr, $REPARTITION_FILE_MIN_SIZE: expr) => { - assert_optimized!($EXPECTED_LINES, $PLAN, $FIRST_ENFORCE_DIST, $PREFER_EXISTING_SORT, $TARGET_PARTITIONS, $REPARTITION_FILE_SCANS, $REPARTITION_FILE_MIN_SIZE, false); - }; + // By default, will not repartition file scans. + config.optimizer.repartition_file_scans = false; + config.optimizer.repartition_file_min_size = 1024; - ($EXPECTED_LINES: expr, $PLAN: expr, $FIRST_ENFORCE_DIST: expr, $PREFER_EXISTING_SORT: expr, $TARGET_PARTITIONS: expr, $REPARTITION_FILE_SCANS: expr, $REPARTITION_FILE_MIN_SIZE: expr, $PREFER_EXISTING_UNION: expr) => { - let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect(); + // By default, set query execution concurrency to 10. + config.execution.target_partitions = 10; + + // Use a small batch size, to trigger RoundRobin in tests + config.execution.batch_size = 1; + + config +} + +#[derive(PartialEq, Clone)] +enum Run { + Distribution, + Sorting, +} - let mut config = ConfigOptions::new(); - config.execution.target_partitions = $TARGET_PARTITIONS; - config.optimizer.repartition_file_scans = $REPARTITION_FILE_SCANS; - config.optimizer.repartition_file_min_size = $REPARTITION_FILE_MIN_SIZE; - config.optimizer.prefer_existing_sort = $PREFER_EXISTING_SORT; - config.optimizer.prefer_existing_union = $PREFER_EXISTING_UNION; - // Use a small batch size, to trigger RoundRobin in tests - config.execution.batch_size = 1; - - // NOTE: These tests verify the joint `EnforceDistribution` + `EnforceSorting` cascade - // because they were written prior to the separation of `BasicEnforcement` into - // `EnforceSorting` and `EnforceDistribution`. - // TODO: Orthogonalize the tests here just to verify `EnforceDistribution` and create - // new tests for the cascade. +/// Standard sets of the series of optimizer runs: +const DISTRIB_DISTRIB_SORT: [Run; 3] = + [Run::Distribution, Run::Distribution, Run::Sorting]; +const SORT_DISTRIB_DISTRIB: [Run; 3] = + [Run::Sorting, Run::Distribution, Run::Distribution]; + +#[derive(Clone)] +struct TestConfig { + config: ConfigOptions, +} + +impl Default for TestConfig { + fn default() -> Self { + Self { + config: test_suite_default_config_options(), + } + } +} + +impl TestConfig { + /// If preferred, will not repartition / resort data if it is already sorted. + fn with_prefer_existing_sort(mut self) -> Self { + self.config.optimizer.prefer_existing_sort = true; + self + } + + /// If preferred, will not attempt to convert Union to Interleave. + fn with_prefer_existing_union(mut self) -> Self { + self.config.optimizer.prefer_existing_union = true; + self + } + + /// If preferred, will repartition file scans. + /// Accepts a minimum file size to repartition. + fn with_prefer_repartition_file_scans(mut self, file_min_size: usize) -> Self { + self.config.optimizer.repartition_file_scans = true; + self.config.optimizer.repartition_file_min_size = file_min_size; + self + } + + /// Set the preferred target partitions for query execution concurrency. + fn with_query_execution_partitions(mut self, target_partitions: usize) -> Self { + self.config.execution.target_partitions = target_partitions; + self + } + + /// Perform a series of runs using the current [`TestConfig`], + /// assert the expected plan result, + /// and return the result plan (for potentional subsequent runs). + fn run( + &self, + expected_lines: &[&str], + plan: Arc, + optimizers_to_run: &[Run], + ) -> Result> { + let expected_lines: Vec<&str> = expected_lines.to_vec(); // Add the ancillary output requirements operator at the start: let optimizer = OutputRequirements::new_add_mode(); - let optimized = optimizer.optimize($PLAN.clone(), &config)?; + let mut optimized = optimizer.optimize(plan.clone(), &self.config)?; // This file has 2 rules that use tree node, apply these rules to original plan consecutively // After these operations tree nodes should be in a consistent state. // This code block makes sure that these rules doesn't violate tree node integrity. { - let adjusted = if config.optimizer.top_down_join_key_reordering { + let adjusted = if self.config.optimizer.top_down_join_key_reordering { // Run adjust_input_keys_ordering rule let plan_requirements = - PlanWithKeyRequirements::new_default($PLAN.clone()); + PlanWithKeyRequirements::new_default(plan.clone()); let adjusted = plan_requirements .transform_down(adjust_input_keys_ordering) .data() @@ -430,70 +477,58 @@ macro_rules! assert_optimized { adjusted.plan } else { // Run reorder_join_keys_to_inputs rule - $PLAN.clone().transform_up(|plan| { - Ok(Transformed::yes(reorder_join_keys_to_inputs(plan)?)) - }) - .data()? + plan.clone() + .transform_up(|plan| { + Ok(Transformed::yes(reorder_join_keys_to_inputs(plan)?)) + }) + .data()? }; // Then run ensure_distribution rule DistributionContext::new_default(adjusted) .transform_up(|distribution_context| { - ensure_distribution(distribution_context, &config) + ensure_distribution(distribution_context, &self.config) }) .data() .and_then(check_integrity)?; // TODO: End state payloads will be checked here. } - let optimized = if $FIRST_ENFORCE_DIST { - // Run enforce distribution rule first: - let optimizer = EnforceDistribution::new(); - let optimized = optimizer.optimize(optimized, &config)?; - // The rule should be idempotent. - // Re-running this rule shouldn't introduce unnecessary operators. - let optimizer = EnforceDistribution::new(); - let optimized = optimizer.optimize(optimized, &config)?; - // Run the enforce sorting rule: - let optimizer = EnforceSorting::new(); - let optimized = optimizer.optimize(optimized, &config)?; - optimized - } else { - // Run the enforce sorting rule first: - let optimizer = EnforceSorting::new(); - let optimized = optimizer.optimize(optimized, &config)?; - // Run enforce distribution rule: - let optimizer = EnforceDistribution::new(); - let optimized = optimizer.optimize(optimized, &config)?; - // The rule should be idempotent. - // Re-running this rule shouldn't introduce unnecessary operators. - let optimizer = EnforceDistribution::new(); - let optimized = optimizer.optimize(optimized, &config)?; - optimized - }; + for run in optimizers_to_run { + optimized = match run { + Run::Distribution => { + let optimizer = EnforceDistribution::new(); + optimizer.optimize(optimized, &self.config)? + } + Run::Sorting => { + let optimizer = EnforceSorting::new(); + optimizer.optimize(optimized, &self.config)? + } + }; + } // Remove the ancillary output requirements operator when done: let optimizer = OutputRequirements::new_remove_mode(); - let optimized = optimizer.optimize(optimized, &config)?; + let optimized = optimizer.optimize(optimized, &self.config)?; // Now format correctly - let plan = displayable(optimized.as_ref()).indent(true).to_string(); - let actual_lines = trim_plan_display(&plan); + let actual_lines = get_plan_string(&optimized); assert_eq!( &expected_lines, &actual_lines, "\n\nexpected:\n\n{:#?}\nactual:\n\n{:#?}\n\n", expected_lines, actual_lines ); - }; + + Ok(optimized) + } } macro_rules! assert_plan_txt { ($EXPECTED_LINES: expr, $PLAN: expr) => { let expected_lines: Vec<&str> = $EXPECTED_LINES.iter().map(|s| *s).collect(); // Now format correctly - let plan = displayable($PLAN.as_ref()).indent(true).to_string(); - let actual_lines = trim_plan_display(&plan); + let actual_lines = get_plan_string(&$PLAN); assert_eq!( &expected_lines, &actual_lines, @@ -534,9 +569,11 @@ fn multi_hash_joins() -> Result<()> { for join_type in join_types { let join = hash_join_exec(left.clone(), right.clone(), &join_on, &join_type); - let join_plan = format!( - "HashJoinExec: mode=Partitioned, join_type={join_type}, on=[(a@0, b1@1)]" - ); + let join_plan = |shift| -> String { + format!("{}HashJoinExec: mode=Partitioned, join_type={join_type}, on=[(a@0, b1@1)]", " ".repeat(shift)) + }; + let join_plan_indent2 = join_plan(2); + let join_plan_indent4 = join_plan(4); match join_type { JoinType::Inner @@ -564,37 +601,39 @@ fn multi_hash_joins() -> Result<()> { // Should include 3 RepartitionExecs JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 4 RepartitionExecs _ => vec![ top_join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + &join_plan_indent4, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], }; - assert_optimized!(expected, top_join.clone(), true); - assert_optimized!(expected, top_join, false); + + let test_config = TestConfig::default(); + test_config.run(&expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected, top_join, &SORT_DISTRIB_DISTRIB)?; } JoinType::RightSemi | JoinType::RightAnti => {} } @@ -627,38 +666,40 @@ fn multi_hash_joins() -> Result<()> { JoinType::Inner | JoinType::Right | JoinType::RightSemi | JoinType::RightAnti => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 4 RepartitionExecs _ => vec![ top_join_plan.as_str(), - "RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10", - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10", + &join_plan_indent4, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], }; - assert_optimized!(expected, top_join.clone(), true); - assert_optimized!(expected, top_join, false); + + let test_config = TestConfig::default(); + test_config.run(&expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected, top_join, &SORT_DISTRIB_DISTRIB)?; } JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => {} } @@ -702,20 +743,21 @@ fn multi_joins_after_alias() -> Result<()> { // Output partition need to respect the Alias and should not introduce additional RepartitionExec let expected = &[ "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a1@0, c@2)]", - "ProjectionExec: expr=[a@0 as a1, a@0 as a2]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, top_join.clone(), true); - assert_optimized!(expected, top_join, false); + " ProjectionExec: expr=[a@0 as a1, a@0 as a2]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + let test_config = TestConfig::default(); + test_config.run(expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, top_join, &SORT_DISTRIB_DISTRIB)?; // Join on (a2 == c) let top_join_on = vec![( @@ -728,20 +770,21 @@ fn multi_joins_after_alias() -> Result<()> { // Output partition need to respect the Alias and should not introduce additional RepartitionExec let expected = &[ "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a2@1, c@2)]", - "ProjectionExec: expr=[a@0 as a1, a@0 as a2]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, top_join.clone(), true); - assert_optimized!(expected, top_join, false); + " ProjectionExec: expr=[a@0 as a1, a@0 as a2]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + let test_config = TestConfig::default(); + test_config.run(expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, top_join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -779,23 +822,24 @@ fn multi_joins_after_multi_alias() -> Result<()> { // The original Output partition can not satisfy the Join requirements and need to add an additional RepartitionExec let expected = &[ "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, c@2)]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "ProjectionExec: expr=[c1@0 as a]", - "ProjectionExec: expr=[c@2 as c1]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - - assert_optimized!(expected, top_join.clone(), true); - assert_optimized!(expected, top_join, false); + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " ProjectionExec: expr=[c1@0 as a]", + " ProjectionExec: expr=[c@2 as c1]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + + let test_config = TestConfig::default(); + test_config.run(expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, top_join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -823,19 +867,20 @@ fn join_after_agg_alias() -> Result<()> { // Only two RepartitionExecs added let expected = &[ "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a1@0, a2@0)]", - "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]", - "RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a2], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]", + " RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a2], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, join.clone(), true); - assert_optimized!(expected, join, false); + let test_config = TestConfig::default(); + test_config.run(expected, join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -875,20 +920,21 @@ fn hash_join_key_ordering() -> Result<()> { // Only two RepartitionExecs added let expected = &[ "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b1@1, b@0), (a1@0, a@1)]", - "ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", - "AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, join.clone(), true); - assert_optimized!(expected, join, false); + " ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", + " AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + let test_config = TestConfig::default(); + test_config.run(expected, join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -994,27 +1040,28 @@ fn multi_hash_join_key_ordering() -> Result<()> { // The bottom joins' join key ordering is adjusted based on the top join. And the top join should not introduce additional RepartitionExec let expected = &[ "FilterExec: c@6 > 1", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(B@2, b1@6), (C@3, c@2), (AA@1, a1@5)]", - "ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]", - "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]", - "RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, filter_top_join.clone(), true); - assert_optimized!(expected, filter_top_join, false); + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(B@2, b1@6), (C@3, c@2), (AA@1, a1@5)]", + " ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]", + " RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(b@1, b1@1), (c@2, c1@2), (a@0, a1@0)]", + " RepartitionExec: partitioning=Hash([b@1, c@2, a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1, c1@2, a1@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + let test_config = TestConfig::default(); + test_config.run(expected, filter_top_join.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, filter_top_join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1133,23 +1180,23 @@ fn reorder_join_keys_to_left_input() -> Result<()> { // The top joins' join key ordering is adjusted based on the children inputs. let expected = &[ top_join_plan.as_str(), - "ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1), (c@2, c1@2)]", - "RepartitionExec: partitioning=Hash([a@0, b@1, c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([a1@0, b1@1, c1@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]", - "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1), (c@2, c1@2)]", + " RepartitionExec: partitioning=Hash([a@0, b@1, c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a1@0, b1@1, c1@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]", + " RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; assert_plan_txt!(expected, reordered); @@ -1267,23 +1314,23 @@ fn reorder_join_keys_to_right_input() -> Result<()> { // The top joins' join key ordering is adjusted based on the children inputs. let expected = &[ top_join_plan.as_str(), - "ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1)]", - "RepartitionExec: partitioning=Hash([a@0, b@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([a1@0, b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]", - "RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " ProjectionExec: expr=[a@0 as A, a@0 as AA, b@1 as B, c@2 as C]", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0), (b@1, b1@1)]", + " RepartitionExec: partitioning=Hash([a@0, b@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a1@0, b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c@2, c1@2), (b@1, b1@1), (a@0, a1@0)]", + " RepartitionExec: partitioning=Hash([c@2, b@1, a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c1@2, b1@1, a1@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; assert_plan_txt!(expected, reordered); @@ -1292,8 +1339,11 @@ fn reorder_join_keys_to_right_input() -> Result<()> { Ok(()) } +/// These test cases use [`TestConfig::with_prefer_existing_sort`]. #[test] fn multi_smj_joins() -> Result<()> { + let test_config = TestConfig::default().with_prefer_existing_sort(); + let left = parquet_exec(); let alias_pairs: Vec<(String, String)> = vec![ ("a".to_string(), "a1".to_string()), @@ -1323,7 +1373,15 @@ fn multi_smj_joins() -> Result<()> { for join_type in join_types { let join = sort_merge_join_exec(left.clone(), right.clone(), &join_on, &join_type); - let join_plan = format!("SortMergeJoin: join_type={join_type}, on=[(a@0, b1@1)]"); + let join_plan = |shift| -> String { + format!( + "{}SortMergeJoin: join_type={join_type}, on=[(a@0, b1@1)]", + " ".repeat(shift) + ) + }; + let join_plan_indent2 = join_plan(2); + let join_plan_indent6 = join_plan(6); + let join_plan_indent10 = join_plan(10); // Top join on (a == c) let top_join_on = vec![( @@ -1340,20 +1398,20 @@ fn multi_smj_joins() -> Result<()> { JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 7 RepartitionExecs (4 hash, 3 round-robin), 4 SortExecs // Since ordering of the left child is not preserved after SortMergeJoin @@ -1367,45 +1425,46 @@ fn multi_smj_joins() -> Result<()> { _ => vec![ top_join_plan.as_str(), // Below 2 operators are differences introduced, when join mode is changed - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - join_plan.as_str(), - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + &join_plan_indent6, + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], }; - assert_optimized!(expected, top_join.clone(), true, true); + // TODO(wiedld): show different test result if enforce sorting first. + test_config.run(&expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; let expected_first_sort_enforcement = match join_type { // Should include 6 RepartitionExecs (3 hash, 3 round-robin), 3 SortExecs JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 8 RepartitionExecs (4 hash, 8 round-robin), 4 SortExecs // Since ordering of the left child is not preserved after SortMergeJoin @@ -1419,27 +1478,32 @@ fn multi_smj_joins() -> Result<()> { _ => vec![ top_join_plan.as_str(), // Below 4 operators are differences introduced, when join mode is changed - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " CoalescePartitionsExec", + &join_plan_indent10, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], }; - assert_optimized!(expected_first_sort_enforcement, top_join, false, true); + // TODO(wiedld): show different test result if enforce distribution first. + test_config.run( + &expected_first_sort_enforcement, + top_join, + &SORT_DISTRIB_DISTRIB, + )?; match join_type { JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { @@ -1458,91 +1522,98 @@ fn multi_smj_joins() -> Result<()> { // Should include 6 RepartitionExecs(3 hash, 3 round-robin) and 3 SortExecs JoinType::Inner | JoinType::Right => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 7 RepartitionExecs (4 hash, 3 round-robin) and 4 SortExecs JoinType::Left | JoinType::Full => vec![ top_join_plan.as_str(), - "SortExec: expr=[b1@6 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10", - join_plan.as_str(), - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b1@6 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10", + &join_plan_indent6, + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // this match arm cannot be reached _ => unreachable!() }; - assert_optimized!(expected, top_join.clone(), true, true); + // TODO(wiedld): show different test result if enforce sorting first. + test_config.run(&expected, top_join.clone(), &DISTRIB_DISTRIB_SORT)?; let expected_first_sort_enforcement = match join_type { // Should include 6 RepartitionExecs (3 of them preserves order) and 3 SortExecs JoinType::Inner | JoinType::Right => vec![ top_join_plan.as_str(), - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + &join_plan_indent2, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // Should include 8 RepartitionExecs (4 of them preserves order) and 4 SortExecs JoinType::Left | JoinType::Full => vec![ top_join_plan.as_str(), - "RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b1@6 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - join_plan.as_str(), - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", - "ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b1@6 ASC], preserve_partitioning=[false]", + " CoalescePartitionsExec", + &join_plan_indent10, + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@1 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]", + " ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=c@2 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ], // this match arm cannot be reached _ => unreachable!() }; - assert_optimized!(expected_first_sort_enforcement, top_join, false, true); + + // TODO(wiedld): show different test result if enforce distribution first. + test_config.run( + &expected_first_sort_enforcement, + top_join, + &SORT_DISTRIB_DISTRIB, + )?; } _ => {} } @@ -1551,6 +1622,7 @@ fn multi_smj_joins() -> Result<()> { Ok(()) } +/// These test cases use [`TestConfig::with_prefer_existing_sort`]. #[test] fn smj_join_key_ordering() -> Result<()> { // group by (a as a1, b as b1) @@ -1597,52 +1669,57 @@ fn smj_join_key_ordering() -> Result<()> { ]; let join = sort_merge_join_exec(left, right.clone(), &join_on, &JoinType::Inner); + // TestConfig: Prefer existing sort. + let test_config = TestConfig::default().with_prefer_existing_sort(); + + // Test: run EnforceDistribution, then EnforceSort. // Only two RepartitionExecs added let expected = &[ "SortMergeJoin: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)]", - "SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[true]", - "ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]", - "ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", - "AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[true]", - "ProjectionExec: expr=[a@1 as a2, b@0 as b2]", - "AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, join.clone(), true, true); - + " SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[true]", + " ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]", + " ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", + " AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[true]", + " ProjectionExec: expr=[a@1 as a2, b@0 as b2]", + " AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + test_config.run(expected, join.clone(), &DISTRIB_DISTRIB_SORT)?; + + // Test: result IS DIFFERENT, if EnforceSorting is run first: let expected_first_sort_enforcement = &[ "SortMergeJoin: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)]", - "RepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=10, preserve_order=true, sort_exprs=b3@1 ASC, a3@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]", - "ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", - "AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "RepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=10, preserve_order=true, sort_exprs=b2@1 ASC, a2@0 ASC", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "ProjectionExec: expr=[a@1 as a2, b@0 as b2]", - "AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected_first_sort_enforcement, join, false, true); + " RepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=10, preserve_order=true, sort_exprs=b3@1 ASC, a3@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[false]", + " CoalescePartitionsExec", + " ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]", + " ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]", + " AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=10, preserve_order=true, sort_exprs=b2@1 ASC, a2@0 ASC", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[false]", + " CoalescePartitionsExec", + " ProjectionExec: expr=[a@1 as a2, b@0 as b2]", + " AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + test_config.run(expected_first_sort_enforcement, join, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1666,26 +1743,31 @@ fn merge_does_not_need_sort() -> Result<()> { let exec: Arc = Arc::new(SortPreservingMergeExec::new(sort_key, exec)); + // Test: run EnforceDistribution, then EnforceSort. + // // The optimizer should not add an additional SortExec as the // data is already sorted let expected = &[ "SortPreservingMergeExec: [a@0 ASC]", - "CoalesceBatchesExec: target_batch_size=4096", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " CoalesceBatchesExec: target_batch_size=4096", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; - assert_optimized!(expected, exec, true); + let test_config = TestConfig::default(); + test_config.run(expected, exec.clone(), &DISTRIB_DISTRIB_SORT)?; + // Test: result IS DIFFERENT, if EnforceSorting is run first: + // // In this case preserving ordering through order preserving operators is not desirable // (according to flag: PREFER_EXISTING_SORT) // hence in this case ordering lost during CoalescePartitionsExec and re-introduced with // SortExec at the top. - let expected = &[ + let expected_first_sort_enforcement = &[ "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "CoalesceBatchesExec: target_batch_size=4096", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " CoalescePartitionsExec", + " CoalesceBatchesExec: target_batch_size=4096", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; - assert_optimized!(expected, exec, false); + test_config.run(expected_first_sort_enforcement, exec, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1713,21 +1795,23 @@ fn union_to_interleave() -> Result<()> { // Only two RepartitionExecs added, no final RepartitionExec required let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]", - "AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]", - "InterleaveExec", - "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan.clone(), false); + " AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]", + " InterleaveExec", + " AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1755,39 +1839,26 @@ fn union_not_to_interleave() -> Result<()> { // Only two RepartitionExecs added, no final RepartitionExec required let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a2@0 as a2], aggr=[]", - "RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=20", - "AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]", - "UnionExec", - "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", - "RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; - // no sort in the plan but since we need it as a parameter, make it default false - let prefer_existing_sort = false; - let first_enforce_distribution = true; - let prefer_existing_union = true; - - assert_optimized!( - expected, - plan.clone(), - first_enforce_distribution, - prefer_existing_sort, - prefer_existing_union - ); - assert_optimized!( - expected, - plan, - !first_enforce_distribution, - prefer_existing_sort, - prefer_existing_union - ); + " RepartitionExec: partitioning=Hash([a2@0], 10), input_partitions=20", + " AggregateExec: mode=Partial, gby=[a1@0 as a2], aggr=[]", + " UnionExec", + " AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " AggregateExec: mode=FinalPartitioned, gby=[a1@0 as a1], aggr=[]", + " RepartitionExec: partitioning=Hash([a1@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a1], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + + // TestConfig: Prefer existing union. + let test_config = TestConfig::default().with_prefer_existing_union(); + + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1799,13 +1870,15 @@ fn added_repartition_to_single_partition() -> Result<()> { let expected = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(&expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1817,35 +1890,37 @@ fn repartition_deepest_node() -> Result<()> { let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } #[test] - fn repartition_unsorted_limit() -> Result<()> { let plan = limit_exec(filter_exec(parquet_exec())); let expected = &[ "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // nothing sorts the data, so the local limit doesn't require sorted data either - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1861,13 +1936,15 @@ fn repartition_sorted_limit() -> Result<()> { let expected = &[ "GlobalLimitExec: skip=0, fetch=100", - "LocalLimitExec: fetch=100", + " LocalLimitExec: fetch=100", // data is sorted so can't repartition here - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1886,16 +1963,17 @@ fn repartition_sorted_limit_with_filter() -> Result<()> { let expected = &[ "SortRequiredExec: [c@2 ASC]", - "FilterExec: c@2 = 0", + " FilterExec: c@2 = 0", // We can use repartition here, ordering requirement by SortRequiredExec // is still satisfied. - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1910,22 +1988,24 @@ fn repartition_ignores_limit() -> Result<()> { let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // repartition should happen prior to the filter to maximize parallelism - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", - "LocalLimitExec: fetch=100", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", + " LocalLimitExec: fetch=100", // Expect no repartition to happen for local limit - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1937,15 +2017,16 @@ fn repartition_ignores_union() -> Result<()> { let expected = &[ "UnionExec", // Expect no repartition of DataSourceExec - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1963,10 +2044,12 @@ fn repartition_through_sort_preserving_merge() -> Result<()> { // need resort as the data was not sorted correctly let expected = &[ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -1984,21 +2067,24 @@ fn repartition_ignores_sort_preserving_merge() -> Result<()> { parquet_exec_multiple_sorted(vec![sort_key]), ); + // Test: run EnforceDistribution, then EnforceSort + // // should not sort (as the data was already sorted) // should not repartition, since increased parallelism is not beneficial for SortPReservingMerge let expected = &[ "SortPreservingMergeExec: [c@2 ASC]", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, plan.clone(), true); - - let expected = &[ + // Test: result IS DIFFERENT, if EnforceSorting is run first: + let expected_first_sort_enforcement = &[ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " CoalescePartitionsExec", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, plan, false); + test_config.run(expected_first_sort_enforcement, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2014,28 +2100,32 @@ fn repartition_ignores_sort_preserving_merge_with_union() -> Result<()> { let input = union_exec(vec![parquet_exec_with_sort(vec![sort_key.clone()]); 2]); let plan = sort_preserving_merge_exec(sort_key, input); + // Test: run EnforceDistribution, then EnforceSort. + // // should not repartition / sort (as the data was already sorted) let expected = &[ "SortPreservingMergeExec: [c@2 ASC]", - "UnionExec", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, plan.clone(), true); - - let expected = &[ + // test: result IS DIFFERENT, if EnforceSorting is run first: + let expected_first_sort_enforcement = &[ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "UnionExec", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " CoalescePartitionsExec", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, plan, false); + test_config.run(expected_first_sort_enforcement, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } +/// These test cases use [`TestConfig::with_prefer_existing_sort`]. #[test] fn repartition_does_not_destroy_sort() -> Result<()> { // SortRequired @@ -2050,16 +2140,19 @@ fn repartition_does_not_destroy_sort() -> Result<()> { sort_key, ); + // TestConfig: Prefer existing sort. + let test_config = TestConfig::default().with_prefer_existing_sort(); + // during repartitioning ordering is preserved let expected = &[ "SortRequiredExec: [d@3 ASC]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[d@3 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[d@3 ASC], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true, true); - assert_optimized!(expected, plan, false, true); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2092,15 +2185,17 @@ fn repartition_does_not_destroy_sort_more_complex() -> Result<()> { let expected = &[ "UnionExec", // union input 1: no repartitioning - "SortRequiredExec: [c@2 ASC]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " SortRequiredExec: [c@2 ASC]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", // union input 2: should repartition - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2124,26 +2219,28 @@ fn repartition_transitively_with_projection() -> Result<()> { }]); let plan = sort_preserving_merge_exec(sort_key, proj); + // Test: run EnforceDistribution, then EnforceSort. let expected = &[ "SortPreservingMergeExec: [sum@0 ASC]", - "SortExec: expr=[sum@0 ASC], preserve_partitioning=[true]", + " SortExec: expr=[sum@0 ASC], preserve_partitioning=[true]", // Since this projection is not trivial, increasing parallelism is beneficial - "ProjectionExec: expr=[a@0 + b@1 as sum]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " ProjectionExec: expr=[a@0 + b@1 as sum]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, plan.clone(), true); - + // Test: result IS DIFFERENT, if EnforceSorting is run first: let expected_first_sort_enforcement = &[ "SortExec: expr=[sum@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", + " CoalescePartitionsExec", // Since this projection is not trivial, increasing parallelism is beneficial - "ProjectionExec: expr=[a@0 + b@1 as sum]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " ProjectionExec: expr=[a@0 + b@1 as sum]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected_first_sort_enforcement, plan, false); + test_config.run(expected_first_sort_enforcement, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2172,11 +2269,13 @@ fn repartition_ignores_transitively_with_projection() -> Result<()> { let expected = &[ "SortRequiredExec: [c@2 ASC]", // Since this projection is trivial, increasing parallelism is not beneficial - "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2203,13 +2302,15 @@ fn repartition_transitively_past_sort_with_projection() -> Result<()> { ); let expected = &[ - "SortExec: expr=[c@2 ASC], preserve_partitioning=[true]", + "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", // Since this projection is trivial, increasing parallelism is not beneficial - "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, plan.clone(), true); - assert_optimized!(expected, plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2223,26 +2324,28 @@ fn repartition_transitively_past_sort_with_filter() -> Result<()> { }]); let plan = sort_exec(sort_key, filter_exec(parquet_exec()), false); + // Test: run EnforceDistribution, then EnforceSort. let expected = &[ "SortPreservingMergeExec: [a@0 ASC]", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", // Expect repartition on the input to the sort (as it can benefit from additional parallelism) - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, plan.clone(), true); - + // Test: result IS DIFFERENT, if EnforceSorting is run first: let expected_first_sort_enforcement = &[ "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "FilterExec: c@2 = 0", + " CoalescePartitionsExec", + " FilterExec: c@2 = 0", // Expect repartition on the input of the filter (as it can benefit from additional parallelism) - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected_first_sort_enforcement, plan, false); + test_config.run(expected_first_sort_enforcement, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2268,28 +2371,30 @@ fn repartition_transitively_past_sort_with_projection_and_filter() -> Result<()> false, ); + // Test: run EnforceDistribution, then EnforceSort. let expected = &[ "SortPreservingMergeExec: [a@0 ASC]", // Expect repartition on the input to the sort (as it can benefit from additional parallelism) - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", - "FilterExec: c@2 = 0", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", + " FilterExec: c@2 = 0", // repartition is lowest down - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + let test_config = TestConfig::default(); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, plan.clone(), true); - + // Test: result IS DIFFERENT, if EnforceSorting is run first: let expected_first_sort_enforcement = &[ "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " CoalescePartitionsExec", + " ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected_first_sort_enforcement, plan, false); + test_config.run(expected_first_sort_enforcement, plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2300,20 +2405,33 @@ fn parallelization_single_partition() -> Result<()> { let plan_parquet = aggregate_exec_with_alias(parquet_exec(), alias.clone()); let plan_csv = aggregate_exec_with_alias(csv_exec(), alias); + let test_config = TestConfig::default() + .with_prefer_repartition_file_scans(10) + .with_query_execution_partitions(2); + + // Test: with parquet let expected_parquet = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + test_config.run( + &expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(&expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true, false, 2, true, 10); - assert_optimized!(expected_csv, plan_csv, true, false, 2, true, 10); + test_config.run(&expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2329,43 +2447,47 @@ fn parallelization_multiple_files() -> Result<()> { let plan = filter_exec(parquet_exec_multiple_sorted(vec![sort_key.clone()])); let plan = sort_required_exec_with_req(plan, sort_key); + let test_config = TestConfig::default() + .with_prefer_existing_sort() + .with_prefer_repartition_file_scans(1); + // The groups must have only contiguous ranges of rows from the same file // if any group has rows from multiple files, the data is no longer sorted destroyed // https://github.com/apache/datafusion/issues/8451 - let expected = [ + let expected_with_3_target_partitions = [ "SortRequiredExec: [a@0 ASC]", - "FilterExec: c@2 = 0", - "DataSourceExec: file_groups={3 groups: [[x:0..50], [y:0..100], [x:50..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; - let target_partitions = 3; - let repartition_size = 1; - assert_optimized!( - expected, - plan, - true, - true, - target_partitions, - true, - repartition_size, - false - ); + " FilterExec: c@2 = 0", + " DataSourceExec: file_groups={3 groups: [[x:0..50], [y:0..100], [x:50..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + ]; + let test_config_concurrency_3 = + test_config.clone().with_query_execution_partitions(3); + test_config_concurrency_3.run( + &expected_with_3_target_partitions, + plan.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config_concurrency_3.run( + &expected_with_3_target_partitions, + plan.clone(), + &SORT_DISTRIB_DISTRIB, + )?; - let expected = [ + let expected_with_8_target_partitions = [ "SortRequiredExec: [a@0 ASC]", - "FilterExec: c@2 = 0", - "DataSourceExec: file_groups={8 groups: [[x:0..25], [y:0..25], [x:25..50], [y:25..50], [x:50..75], [y:50..75], [x:75..100], [y:75..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " DataSourceExec: file_groups={8 groups: [[x:0..25], [y:0..25], [x:25..50], [y:25..50], [x:50..75], [y:50..75], [x:75..100], [y:75..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; - let target_partitions = 8; - let repartition_size = 1; - assert_optimized!( - expected, + let test_config_concurrency_8 = test_config.with_query_execution_partitions(8); + test_config_concurrency_8.run( + &expected_with_8_target_partitions, + plan.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config_concurrency_8.run( + &expected_with_8_target_partitions, plan, - true, - true, - target_partitions, - true, - repartition_size, - false - ); + &SORT_DISTRIB_DISTRIB, + )?; Ok(()) } @@ -2384,17 +2506,17 @@ fn parallelization_compressed_csv() -> Result<()> { let expected_not_partitioned = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; let expected_partitioned = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " DataSourceExec: file_groups={2 groups: [[x:0..50], [x:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; for compression_type in compression_types { @@ -2415,7 +2537,11 @@ fn parallelization_compressed_csv() -> Result<()> { .build(), vec![("a".to_string(), "a".to_string())], ); - assert_optimized!(expected, plan, true, false, 2, true, 10, false); + let test_config = TestConfig::default() + .with_query_execution_partitions(2) + .with_prefer_repartition_file_scans(10); + test_config.run(expected, plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, plan, &SORT_DISTRIB_DISTRIB)?; } Ok(()) } @@ -2426,22 +2552,36 @@ fn parallelization_two_partitions() -> Result<()> { let plan_parquet = aggregate_exec_with_alias(parquet_exec_multiple(), alias.clone()); let plan_csv = aggregate_exec_with_alias(csv_exec_multiple(), alias); + let test_config = TestConfig::default() + .with_query_execution_partitions(2) + .with_prefer_repartition_file_scans(10); + + // Test: with parquet let expected_parquet = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", // Plan already has two partitions - "DataSourceExec: file_groups={2 groups: [[x:0..100], [y:0..100]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={2 groups: [[x:0..100], [y:0..100]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + test_config.run( + &expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(&expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([a@0], 2), input_partitions=2", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", // Plan already has two partitions - "DataSourceExec: file_groups={2 groups: [[x:0..100], [y:0..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={2 groups: [[x:0..100], [y:0..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true, false, 2, true, 10); - assert_optimized!(expected_csv, plan_csv, true, false, 2, true, 10); + test_config.run(&expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; + Ok(()) } @@ -2451,22 +2591,35 @@ fn parallelization_two_partitions_into_four() -> Result<()> { let plan_parquet = aggregate_exec_with_alias(parquet_exec_multiple(), alias.clone()); let plan_csv = aggregate_exec_with_alias(csv_exec_multiple(), alias); + let test_config = TestConfig::default() + .with_query_execution_partitions(4) + .with_prefer_repartition_file_scans(10); + + // Test: with parquet let expected_parquet = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", // Multiple source files splitted across partitions - "DataSourceExec: file_groups={4 groups: [[x:0..50], [x:50..100], [y:0..50], [y:50..100]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={4 groups: [[x:0..50], [x:50..100], [y:0..50], [y:50..100]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + test_config.run( + &expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(&expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = [ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", // Multiple source files splitted across partitions - "DataSourceExec: file_groups={4 groups: [[x:0..50], [x:50..100], [y:0..50], [y:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={4 groups: [[x:0..50], [x:50..100], [y:0..50], [y:50..100]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true, false, 4, true, 10); - assert_optimized!(expected_csv, plan_csv, true, false, 4, true, 10); + test_config.run(&expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(&expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2481,24 +2634,35 @@ fn parallelization_sorted_limit() -> Result<()> { let plan_parquet = limit_exec(sort_exec(sort_key.clone(), parquet_exec(), false)); let plan_csv = limit_exec(sort_exec(sort_key, csv_exec(), false)); + let test_config = TestConfig::default(); + + // Test: with parquet let expected_parquet = &[ "GlobalLimitExec: skip=0, fetch=100", - "LocalLimitExec: fetch=100", + " LocalLimitExec: fetch=100", // data is sorted so can't repartition here - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", // Doesn't parallelize for SortExec without preserve_partitioning - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "GlobalLimitExec: skip=0, fetch=100", - "LocalLimitExec: fetch=100", + " LocalLimitExec: fetch=100", // data is sorted so can't repartition here - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", // Doesn't parallelize for SortExec without preserve_partitioning - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2517,32 +2681,43 @@ fn parallelization_limit_with_filter() -> Result<()> { ))); let plan_csv = limit_exec(filter_exec(sort_exec(sort_key, csv_exec(), false))); + let test_config = TestConfig::default(); + + // Test: with parquet let expected_parquet = &[ "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // even though data is sorted, we can use repartition here. Since // ordering is not used in subsequent stages anyway. - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", // SortExec doesn't benefit from input partitioning - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // even though data is sorted, we can use repartition here. Since // ordering is not used in subsequent stages anyway. - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", // SortExec doesn't benefit from input partitioning - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2557,40 +2732,51 @@ fn parallelization_ignores_limit() -> Result<()> { let plan_csv = aggregate_exec_with_alias(limit_exec(filter_exec(limit_exec(csv_exec()))), alias); + let test_config = TestConfig::default(); + + // Test: with parquet let expected_parquet = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // repartition should happen prior to the filter to maximize parallelism - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", // Limit doesn't benefit from input partitioning - no parallelism - "LocalLimitExec: fetch=100", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " LocalLimitExec: fetch=100", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", - "CoalescePartitionsExec", - "LocalLimitExec: fetch=100", - "FilterExec: c@2 = 0", + " RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", + " CoalescePartitionsExec", + " LocalLimitExec: fetch=100", + " FilterExec: c@2 = 0", // repartition should happen prior to the filter to maximize parallelism - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "GlobalLimitExec: skip=0, fetch=100", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " GlobalLimitExec: skip=0, fetch=100", // Limit doesn't benefit from input partitioning - no parallelism - "LocalLimitExec: fetch=100", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " LocalLimitExec: fetch=100", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2600,26 +2786,37 @@ fn parallelization_union_inputs() -> Result<()> { let plan_parquet = union_exec(vec![parquet_exec(); 5]); let plan_csv = union_exec(vec![csv_exec(); 5]); + let test_config = TestConfig::default(); + + // Test: with parquet let expected_parquet = &[ "UnionExec", // Union doesn't benefit from input partitioning - no parallelism - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ]; + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "UnionExec", // Union doesn't benefit from input partitioning - no parallelism - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2639,15 +2836,28 @@ fn parallelization_prior_to_sort_preserving_merge() -> Result<()> { let plan_csv = sort_preserving_merge_exec(sort_key.clone(), csv_exec_with_sort(vec![sort_key])); + let test_config = TestConfig::default(); + + // Expected Outcome: // parallelization is not beneficial for SortPreservingMerge + + // Test: with parquet let expected_parquet = &[ "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2666,22 +2876,61 @@ fn parallelization_sort_preserving_merge_with_union() -> Result<()> { let plan_parquet = sort_preserving_merge_exec(sort_key.clone(), input_parquet); let plan_csv = sort_preserving_merge_exec(sort_key, input_csv); + let test_config = TestConfig::default(); + + // Expected Outcome: // should not repartition (union doesn't benefit from increased parallelism) // should not sort (as the data was already sorted) + + // Test: with parquet let expected_parquet = &[ "SortPreservingMergeExec: [c@2 ASC]", - "UnionExec", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + let expected_parquet_first_sort_enforcement = &[ + // no SPM + "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + // has coalesce + " CoalescePartitionsExec", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + ]; + test_config.run( + expected_parquet_first_sort_enforcement, + plan_parquet, + &SORT_DISTRIB_DISTRIB, + )?; + + // Test: with csv let expected_csv = &[ "SortPreservingMergeExec: [c@2 ASC]", - "UnionExec", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + let expected_csv_first_sort_enforcement = &[ + // no SPM + "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + // has coalesce + " CoalescePartitionsExec", + " UnionExec", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + ]; + test_config.run( + expected_csv_first_sort_enforcement, + plan_csv.clone(), + &SORT_DISTRIB_DISTRIB, + )?; Ok(()) } @@ -2702,17 +2951,30 @@ fn parallelization_does_not_benefit() -> Result<()> { let plan_csv = sort_required_exec_with_req(csv_exec_with_sort(vec![sort_key.clone()]), sort_key); + let test_config = TestConfig::default(); + + // Expected Outcome: // no parallelization, because SortRequiredExec doesn't benefit from increased parallelism + + // Test: with parquet let expected_parquet = &[ "SortRequiredExec: [c@2 ASC]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; + + // Test: with csv let expected_csv = &[ "SortRequiredExec: [c@2 ASC]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", ]; - assert_optimized!(expected_parquet, plan_parquet, true); - assert_optimized!(expected_csv, plan_csv, true); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2746,12 +3008,19 @@ fn parallelization_ignores_transitively_with_projection_parquet() -> Result<()> ]; plans_matches_expected!(expected, &plan_parquet); + // Expected Outcome: // data should not be repartitioned / resorted let expected_parquet = &[ "ProjectionExec: expr=[a@0 as a2, c@2 as c2]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected_parquet, plan_parquet, true); + let test_config = TestConfig::default(); + test_config.run( + expected_parquet, + plan_parquet.clone(), + &DISTRIB_DISTRIB_SORT, + )?; + test_config.run(expected_parquet, plan_parquet, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2785,12 +3054,15 @@ fn parallelization_ignores_transitively_with_projection_csv() -> Result<()> { ]; plans_matches_expected!(expected, &plan_csv); + // Expected Outcome: // data should not be repartitioned / resorted let expected_csv = &[ "ProjectionExec: expr=[a@0 as a2, c@2 as c2]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false", ]; - assert_optimized!(expected_csv, plan_csv, true); + let test_config = TestConfig::default(); + test_config.run(expected_csv, plan_csv.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected_csv, plan_csv, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2811,15 +3083,18 @@ fn remove_redundant_roundrobins() -> Result<()> { let expected = &[ "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, physical_plan.clone(), true); - assert_optimized!(expected, physical_plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } +/// This test case uses [`TestConfig::with_prefer_existing_sort`]. #[test] fn remove_unnecessary_spm_after_filter() -> Result<()> { let schema = schema(); @@ -2830,21 +3105,26 @@ fn remove_unnecessary_spm_after_filter() -> Result<()> { let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); let physical_plan = sort_preserving_merge_exec(sort_key, filter_exec(input)); + // TestConfig: Prefer existing sort. + let test_config = TestConfig::default().with_prefer_existing_sort(); + + // Expected Outcome: // Original plan expects its output to be ordered by c@2 ASC. // This is still satisfied since, after filter that column is constant. let expected = &[ "CoalescePartitionsExec", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c@2 ASC", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c@2 ASC", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - // last flag sets config.optimizer.PREFER_EXISTING_SORT - assert_optimized!(expected, physical_plan.clone(), true, true); - assert_optimized!(expected, physical_plan, false, true); + + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } +/// This test case uses [`TestConfig::with_prefer_existing_sort`]. #[test] fn preserve_ordering_through_repartition() -> Result<()> { let schema = schema(); @@ -2855,15 +3135,17 @@ fn preserve_ordering_through_repartition() -> Result<()> { let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); let physical_plan = sort_preserving_merge_exec(sort_key, filter_exec(input)); + // TestConfig: Prefer existing sort. + let test_config = TestConfig::default().with_prefer_existing_sort(); + let expected = &[ "SortPreservingMergeExec: [d@3 ASC]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=d@3 ASC", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[d@3 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=d@3 ASC", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[d@3 ASC], file_type=parquet", ]; - // last flag sets config.optimizer.PREFER_EXISTING_SORT - assert_optimized!(expected, physical_plan.clone(), true, true); - assert_optimized!(expected, physical_plan, false, true); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2878,24 +3160,31 @@ fn do_not_preserve_ordering_through_repartition() -> Result<()> { let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); let physical_plan = sort_preserving_merge_exec(sort_key, filter_exec(input)); + let test_config = TestConfig::default(); + + // Test: run EnforceDistribution, then EnforceSort. let expected = &[ "SortPreservingMergeExec: [a@0 ASC]", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, physical_plan.clone(), true); - - let expected = &[ + // Test: result IS DIFFERENT, if EnforceSorting is run first: + let expected_first_sort_enforcement = &[ "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", - ]; - assert_optimized!(expected, physical_plan, false); + " CoalescePartitionsExec", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + ]; + test_config.run( + expected_first_sort_enforcement, + physical_plan, + &SORT_DISTRIB_DISTRIB, + )?; Ok(()) } @@ -2914,12 +3203,13 @@ fn no_need_for_sort_after_filter() -> Result<()> { // After CoalescePartitionsExec c is still constant. Hence c@2 ASC ordering is already satisfied. "CoalescePartitionsExec", // Since after this stage c is constant. c@2 ASC ordering is already satisfied. - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, physical_plan.clone(), true); - assert_optimized!(expected, physical_plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2939,25 +3229,32 @@ fn do_not_preserve_ordering_through_repartition2() -> Result<()> { }]); let physical_plan = sort_preserving_merge_exec(sort_req, filter_exec(input)); + let test_config = TestConfig::default(); + + // Test: run EnforceDistribution, then EnforceSort. let expected = &[ "SortPreservingMergeExec: [a@0 ASC]", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; - assert_optimized!(expected, physical_plan.clone(), true); - - let expected = &[ + // Test: result IS DIFFERENT, if EnforceSorting is run first: + let expected_first_sort_enforcement = &[ "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - "CoalescePartitionsExec", - "SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " CoalescePartitionsExec", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, physical_plan, false); + test_config.run( + expected_first_sort_enforcement, + physical_plan, + &SORT_DISTRIB_DISTRIB, + )?; Ok(()) } @@ -2974,11 +3271,12 @@ fn do_not_preserve_ordering_through_repartition3() -> Result<()> { let expected = &[ "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - assert_optimized!(expected, physical_plan.clone(), true); - assert_optimized!(expected, physical_plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -2996,8 +3294,8 @@ fn do_not_put_sort_when_input_is_invalid() -> Result<()> { // Ordering requirement of sort required exec is NOT satisfied // by existing ordering at the source. "SortRequiredExec: [a@0 ASC]", - "FilterExec: c@2 = 0", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " FilterExec: c@2 = 0", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; assert_plan_txt!(expected, physical_plan); @@ -3005,9 +3303,9 @@ fn do_not_put_sort_when_input_is_invalid() -> Result<()> { "SortRequiredExec: [a@0 ASC]", // Since at the start of the rule ordering requirement is not satisfied // EnforceDistribution rule doesn't satisfy this requirement either. - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; let mut config = ConfigOptions::new(); @@ -3034,8 +3332,8 @@ fn put_sort_when_input_is_valid() -> Result<()> { // Ordering requirement of sort required exec is satisfied // by existing ordering at the source. "SortRequiredExec: [a@0 ASC]", - "FilterExec: c@2 = 0", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; assert_plan_txt!(expected, physical_plan); @@ -3043,8 +3341,8 @@ fn put_sort_when_input_is_valid() -> Result<()> { // Since at the start of the rule ordering requirement is satisfied // EnforceDistribution rule satisfy this requirement also. "SortRequiredExec: [a@0 ASC]", - "FilterExec: c@2 = 0", - "DataSourceExec: file_groups={10 groups: [[x:0..20], [y:0..20], [x:20..40], [y:20..40], [x:40..60], [y:40..60], [x:60..80], [y:60..80], [x:80..100], [y:80..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", + " FilterExec: c@2 = 0", + " DataSourceExec: file_groups={10 groups: [[x:0..20], [y:0..20], [x:20..40], [y:20..40], [x:40..60], [y:40..60], [x:60..80], [y:60..80], [x:80..100], [y:80..100]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet", ]; let mut config = ConfigOptions::new(); @@ -3068,14 +3366,17 @@ fn do_not_add_unnecessary_hash() -> Result<()> { let input = parquet_exec_with_sort(vec![sort_key]); let physical_plan = aggregate_exec_with_alias(input, alias); + // TestConfig: + // Make sure target partition number is 1. In this case hash repartition is unnecessary. + let test_config = TestConfig::default().with_query_execution_partitions(1); + let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - // Make sure target partition number is 1. In this case hash repartition is unnecessary - assert_optimized!(expected, physical_plan.clone(), true, false, 1, false, 1024); - assert_optimized!(expected, physical_plan, false, false, 1, false, 1024); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -3092,20 +3393,23 @@ fn do_not_add_unnecessary_hash2() -> Result<()> { let aggregate = aggregate_exec_with_alias(input, alias.clone()); let physical_plan = aggregate_exec_with_alias(aggregate, alias); + // TestConfig: + // Make sure target partition number is larger than 2 (e.g partition number at the source). + let test_config = TestConfig::default().with_query_execution_partitions(4); + let expected = &[ "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", // Since hash requirements of this operator is satisfied. There shouldn't be // a hash repartition here - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", - "AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", - "RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2", - "DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4", + " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[]", + " RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2", + " DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet", ]; - // Make sure target partition number is larger than 2 (e.g partition number at the source). - assert_optimized!(expected, physical_plan.clone(), true, false, 4, false, 1024); - assert_optimized!(expected, physical_plan, false, false, 4, false, 1024); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -3122,8 +3426,10 @@ fn optimize_away_unnecessary_repartition() -> Result<()> { let expected = &["DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet"]; - assert_optimized!(expected, physical_plan.clone(), true); - assert_optimized!(expected, physical_plan, false); + + let test_config = TestConfig::default(); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } @@ -3145,12 +3451,13 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> { let expected = &[ "FilterExec: c@2 = 0", - "FilterExec: c@2 = 0", - "RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", - "DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", + " FilterExec: c@2 = 0", + " RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", ]; - assert_optimized!(expected, physical_plan.clone(), true); - assert_optimized!(expected, physical_plan, false); + let test_config = TestConfig::default(); + test_config.run(expected, physical_plan.clone(), &DISTRIB_DISTRIB_SORT)?; + test_config.run(expected, physical_plan, &SORT_DISTRIB_DISTRIB)?; Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 3223768acb74..bb77192e05b8 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -17,19 +17,14 @@ use std::sync::Arc; -use crate::physical_optimizer::enforce_distribution::projection_exec_with_alias; -use crate::physical_optimizer::sanity_checker::{ - assert_sanity_check, assert_sanity_check_err, -}; use crate::physical_optimizer::test_utils::{ aggregate_exec, bounded_window_exec, check_integrity, coalesce_batches_exec, coalesce_partitions_exec, create_test_schema, create_test_schema2, create_test_schema3, filter_exec, global_limit_exec, hash_join_exec, limit_exec, - local_limit_exec, memory_exec, parquet_exec, parquet_exec_with_stats, - repartition_exec, schema, single_partitioned_aggregate, sort_exec, sort_expr, - sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, - sort_preserving_merge_exec_with_fetch, spr_repartition_exec, stream_exec_ordered, - union_exec, RequirementsTestExec, + local_limit_exec, memory_exec, parquet_exec, repartition_exec, sort_exec, + sort_exec_with_fetch, sort_expr, sort_expr_options, sort_merge_join_exec, + sort_preserving_merge_exec, sort_preserving_merge_exec_with_fetch, + spr_repartition_exec, stream_exec_ordered, union_exec, RequirementsTestExec, }; use arrow::compute::SortOptions; @@ -2247,7 +2242,7 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> { " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], expected_plan: vec![ - "SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false]", + "SortExec: expr=[non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false]", " WindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]", " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], @@ -2264,7 +2259,7 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> { " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], expected_plan: vec![ - "SortExec: expr=[non_nullable_col@1 DESC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[false]", + "SortExec: expr=[non_nullable_col@1 DESC NULLS LAST], preserve_partitioning=[false]", " WindowAggExec: wdw=[max: Ok(Field { name: \"max\", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]", " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], @@ -2281,7 +2276,7 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> { " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], expected_plan: vec![ - "SortExec: expr=[min@2 ASC NULLS LAST, non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false]", + "SortExec: expr=[non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false]", " WindowAggExec: wdw=[min: Ok(Field { name: \"min\", data_type: Int32, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]", " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], @@ -2298,7 +2293,7 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> { " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], expected_plan: vec![ - "SortExec: expr=[avg@2 DESC NULLS LAST, nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]", + "SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]", " WindowAggExec: wdw=[avg: Ok(Field { name: \"avg\", data_type: Float64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]", " DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet", ], @@ -3352,61 +3347,88 @@ async fn test_window_partial_constant_and_set_monotonicity() -> Result<()> { Ok(()) } -#[tokio::test] -async fn test_preserve_needed_coalesce() -> Result<()> { - // Input to EnforceSorting, from our test case. - let plan = projection_exec_with_alias( - union_exec(vec![parquet_exec_with_stats(); 2]), - vec![ - ("a".to_string(), "a".to_string()), - ("b".to_string(), "value".to_string()), - ], - ); - let plan = Arc::new(CoalescePartitionsExec::new(plan)); - let schema = schema(); - let sort_key = LexOrdering::new(vec![PhysicalSortExpr { - expr: col("a", &schema).unwrap(), - options: SortOptions::default(), - }]); - let plan: Arc = - single_partitioned_aggregate(plan, vec![("a".to_string(), "a1".to_string())]); - let plan = sort_exec(sort_key, plan); - - // Starting plan: as in our test case. - assert_eq!( - get_plan_string(&plan), - vec![ - "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", - " AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]", - " CoalescePartitionsExec", - " ProjectionExec: expr=[a@0 as a, b@1 as value]", - " UnionExec", - " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ], - ); - // Test: plan is valid. - assert_sanity_check(&plan, true); +#[test] +fn test_removes_unused_orthogonal_sort() -> Result<()> { + let schema = create_test_schema3()?; + let input_sort_exprs = vec![sort_expr("b", &schema), sort_expr("c", &schema)]; + let unbounded_input = stream_exec_ordered(&schema, input_sort_exprs.clone()); - // EnforceSorting will remove the coalesce, and add an SPM further up (above the aggregate). - let optimizer = EnforceSorting::new(); - let optimized = optimizer.optimize(plan, &Default::default())?; - assert_eq!( - get_plan_string(&optimized), - vec![ - "SortPreservingMergeExec: [a@0 ASC]", - " SortExec: expr=[a@0 ASC], preserve_partitioning=[true]", - " AggregateExec: mode=SinglePartitioned, gby=[a@0 as a1], aggr=[]", - " ProjectionExec: expr=[a@0 as a, b@1 as value]", - " UnionExec", - " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet", - ], - ); + let orthogonal_sort = sort_exec(vec![sort_expr("a", &schema)], unbounded_input); + let output_sort = sort_exec(input_sort_exprs, orthogonal_sort); // same sort as data source + + // Test scenario/input has an orthogonal sort: + let expected_input = [ + "SortExec: expr=[b@1 ASC, c@2 ASC], preserve_partitioning=[false]", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[b@1 ASC, c@2 ASC]" + ]; + assert_eq!(get_plan_string(&output_sort), expected_input,); + + // Test: should remove orthogonal sort, and the uppermost (unneeded) sort: + let expected_optimized = [ + "StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[b@1 ASC, c@2 ASC]" + ]; + assert_optimized!(expected_input, expected_optimized, output_sort, true); + + Ok(()) +} + +#[test] +fn test_keeps_used_orthogonal_sort() -> Result<()> { + let schema = create_test_schema3()?; + let input_sort_exprs = vec![sort_expr("b", &schema), sort_expr("c", &schema)]; + let unbounded_input = stream_exec_ordered(&schema, input_sort_exprs.clone()); + + let orthogonal_sort = + sort_exec_with_fetch(vec![sort_expr("a", &schema)], Some(3), unbounded_input); // has fetch, so this orthogonal sort changes the output + let output_sort = sort_exec(input_sort_exprs, orthogonal_sort); + + // Test scenario/input has an orthogonal sort: + let expected_input = [ + "SortExec: expr=[b@1 ASC, c@2 ASC], preserve_partitioning=[false]", + " SortExec: TopK(fetch=3), expr=[a@0 ASC], preserve_partitioning=[false]", + " StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[b@1 ASC, c@2 ASC]" + ]; + assert_eq!(get_plan_string(&output_sort), expected_input,); + + // Test: should keep the orthogonal sort, since it modifies the output: + let expected_optimized = expected_input; + assert_optimized!(expected_input, expected_optimized, output_sort, true); - // Bug: Plan is now invalid. - let err = "does not satisfy distribution requirements: HashPartitioned[[a@0]]). Child-0 output partitioning: UnknownPartitioning(2)"; - assert_sanity_check_err(&optimized, err); + Ok(()) +} + +#[test] +fn test_handles_multiple_orthogonal_sorts() -> Result<()> { + let schema = create_test_schema3()?; + let input_sort_exprs = vec![sort_expr("b", &schema), sort_expr("c", &schema)]; + let unbounded_input = stream_exec_ordered(&schema, input_sort_exprs.clone()); + + let orthogonal_sort_0 = sort_exec(vec![sort_expr("c", &schema)], unbounded_input); // has no fetch, so can be removed + let orthogonal_sort_1 = + sort_exec_with_fetch(vec![sort_expr("a", &schema)], Some(3), orthogonal_sort_0); // has fetch, so this orthogonal sort changes the output + let orthogonal_sort_2 = sort_exec(vec![sort_expr("c", &schema)], orthogonal_sort_1); // has no fetch, so can be removed + let orthogonal_sort_3 = sort_exec(vec![sort_expr("a", &schema)], orthogonal_sort_2); // has no fetch, so can be removed + let output_sort = sort_exec(input_sort_exprs, orthogonal_sort_3); // final sort + + // Test scenario/input has an orthogonal sort: + let expected_input = [ + "SortExec: expr=[b@1 ASC, c@2 ASC], preserve_partitioning=[false]", + " SortExec: expr=[a@0 ASC], preserve_partitioning=[false]", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " SortExec: TopK(fetch=3), expr=[a@0 ASC], preserve_partitioning=[false]", + " SortExec: expr=[c@2 ASC], preserve_partitioning=[false]", + " StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[b@1 ASC, c@2 ASC]", + ]; + assert_eq!(get_plan_string(&output_sort), expected_input,); + + // Test: should keep only the needed orthogonal sort, and remove the unneeded ones: + let expected_optimized = [ + "SortExec: expr=[b@1 ASC, c@2 ASC], preserve_partitioning=[false]", + " SortExec: TopK(fetch=3), expr=[a@0 ASC], preserve_partitioning=[false]", + " StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[b@1 ASC, c@2 ASC]", + ]; + assert_optimized!(expected_input, expected_optimized, output_sort, true); Ok(()) } diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 375af94acaf4..d3b6ec700bee 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -924,6 +924,10 @@ impl DisplayAs for UnboundedExec { self.batch_produce.is_none(), ) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -1019,6 +1023,11 @@ impl DisplayAs for StatisticsExec { self.stats.num_rows, ) } + + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 836758b21318..b0b5f731063f 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -31,7 +31,7 @@ use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::{Operator, ScalarUDF, ScalarUDFImpl, Signature, Volatility}; use datafusion_physical_expr::expressions::{ - binary, col, BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr, + binary, cast, col, BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr, }; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::{ @@ -1382,3 +1382,97 @@ fn test_union_after_projection() -> Result<()> { Ok(()) } + +/// Returns a DataSourceExec that scans a file with (int_col, string_col) +/// and has a partitioning column partition_col (Utf8) +fn partitioned_data_source() -> Arc { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("int_col", DataType::Int32, true), + Field::new("string_col", DataType::Utf8, true), + ])); + + FileScanConfig::new( + ObjectStoreUrl::parse("test:///").unwrap(), + file_schema.clone(), + Arc::new(CsvSource::default()), + ) + .with_file(PartitionedFile::new("x".to_string(), 100)) + .with_table_partition_cols(vec![Field::new("partition_col", DataType::Utf8, true)]) + .with_projection(Some(vec![0, 1, 2])) + .build() +} + +#[test] +fn test_partition_col_projection_pushdown() -> Result<()> { + let source = partitioned_data_source(); + let partitioned_schema = source.schema(); + + let projection = Arc::new(ProjectionExec::try_new( + vec![ + ( + col("string_col", partitioned_schema.as_ref())?, + "string_col".to_string(), + ), + ( + col("partition_col", partitioned_schema.as_ref())?, + "partition_col".to_string(), + ), + ( + col("int_col", partitioned_schema.as_ref())?, + "int_col".to_string(), + ), + ], + source, + )?); + + let after_optimize = + ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?; + + let expected = [ + "ProjectionExec: expr=[string_col@1 as string_col, partition_col@2 as partition_col, int_col@0 as int_col]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[int_col, string_col, partition_col], file_type=csv, has_header=false" + ]; + assert_eq!(get_plan_string(&after_optimize), expected); + + Ok(()) +} + +#[test] +fn test_partition_col_projection_pushdown_expr() -> Result<()> { + let source = partitioned_data_source(); + let partitioned_schema = source.schema(); + + let projection = Arc::new(ProjectionExec::try_new( + vec![ + ( + col("string_col", partitioned_schema.as_ref())?, + "string_col".to_string(), + ), + ( + // CAST(partition_col, Utf8View) + cast( + col("partition_col", partitioned_schema.as_ref())?, + partitioned_schema.as_ref(), + DataType::Utf8View, + )?, + "partition_col".to_string(), + ), + ( + col("int_col", partitioned_schema.as_ref())?, + "int_col".to_string(), + ), + ], + source, + )?); + + let after_optimize = + ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?; + + let expected = [ + "ProjectionExec: expr=[string_col@1 as string_col, CAST(partition_col@2 AS Utf8View) as partition_col, int_col@0 as int_col]", + " DataSourceExec: file_groups={1 group: [[x]]}, projection=[int_col, string_col, partition_col], file_type=csv, has_header=false" + ]; + assert_eq!(get_plan_string(&after_optimize), expected); + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index ee9cb032c341..a73d084a081f 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -388,7 +388,7 @@ fn create_test_schema2() -> SchemaRef { } /// Check if sanity checker should accept or reject plans. -pub(crate) fn assert_sanity_check(plan: &Arc, is_sane: bool) { +fn assert_sanity_check(plan: &Arc, is_sane: bool) { let sanity_checker = SanityCheckPlan::new(); let opts = ConfigOptions::default(); assert_eq!( @@ -397,14 +397,6 @@ pub(crate) fn assert_sanity_check(plan: &Arc, is_sane: bool) ); } -/// Assert reason for sanity check failure. -pub(crate) fn assert_sanity_check_err(plan: &Arc, err: &str) { - let sanity_checker = SanityCheckPlan::new(); - let opts = ConfigOptions::default(); - let error = sanity_checker.optimize(plan.clone(), &opts).unwrap_err(); - assert!(error.message().contains(err)); -} - /// Check if the plan we created is as expected by comparing the plan /// formatted as a string. fn assert_plan(plan: &dyn ExecutionPlan, expected_lines: Vec<&str>) { diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 0b9c3b80bb93..99a75e6e5067 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -30,10 +30,9 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::source::DataSourceExec; use datafusion_common::config::ConfigOptions; -use datafusion_common::stats::Precision; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; -use datafusion_common::{ColumnStatistics, JoinType, Result, Statistics}; +use datafusion_common::{JoinType, Result}; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -103,44 +102,6 @@ pub fn schema() -> SchemaRef { ])) } -fn int64_stats() -> ColumnStatistics { - ColumnStatistics { - null_count: Precision::Absent, - sum_value: Precision::Absent, - max_value: Precision::Exact(1_000_000.into()), - min_value: Precision::Exact(0.into()), - distinct_count: Precision::Absent, - } -} - -fn column_stats() -> Vec { - vec![ - int64_stats(), // a - int64_stats(), // b - int64_stats(), // c - ColumnStatistics::default(), - ColumnStatistics::default(), - ] -} - -/// Create parquet datasource exec using schema from [`schema`]. -pub(crate) fn parquet_exec_with_stats() -> Arc { - let mut statistics = Statistics::new_unknown(&schema()); - statistics.num_rows = Precision::Inexact(10); - statistics.column_statistics = column_stats(); - - let config = FileScanConfig::new( - ObjectStoreUrl::parse("test:///").unwrap(), - schema(), - Arc::new(ParquetSource::new(Default::default())), - ) - .with_file(PartitionedFile::new("x".to_string(), 10000)) - .with_statistics(statistics); - assert_eq!(config.statistics.num_rows, Precision::Inexact(10)); - - config.build() -} - pub fn create_test_schema() -> Result { let nullable_column = Field::new("nullable_col", DataType::Int32, true); let non_nullable_column = Field::new("non_nullable_col", DataType::Int32, false); @@ -334,9 +295,17 @@ pub fn coalesce_batches_exec(input: Arc) -> Arc, input: Arc, +) -> Arc { + sort_exec_with_fetch(sort_exprs, None, input) +} + +pub fn sort_exec_with_fetch( + sort_exprs: impl IntoIterator, + fetch: Option, + input: Arc, ) -> Arc { let sort_exprs = sort_exprs.into_iter().collect(); - Arc::new(SortExec::new(sort_exprs, input)) + Arc::new(SortExec::new(sort_exprs, input).with_fetch(fetch)) } /// A test [`ExecutionPlan`] whose requirements can be configured. @@ -378,8 +347,16 @@ impl RequirementsTestExec { } impl DisplayAs for RequirementsTestExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - write!(f, "RequiredInputOrderingExec") + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "RequiredInputOrderingExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } } } @@ -561,30 +538,6 @@ pub fn build_group_by(input_schema: &SchemaRef, columns: Vec) -> Physica PhysicalGroupBy::new_single(group_by_expr.clone()) } -pub(crate) fn single_partitioned_aggregate( - input: Arc, - alias_pairs: Vec<(String, String)>, -) -> Arc { - let schema = schema(); - let group_by = alias_pairs - .iter() - .map(|(column, alias)| (col(column, &input.schema()).unwrap(), alias.to_string())) - .collect::>(); - let group_by = PhysicalGroupBy::new_single(group_by); - - Arc::new( - AggregateExec::try_new( - AggregateMode::SinglePartitioned, - group_by, - vec![], - vec![], - input, - schema, - ) - .unwrap(), - ) -} - pub fn assert_plan_matches_expected( plan: &Arc, expected: &[&str], diff --git a/datafusion/core/tests/serde/mod.rs b/datafusion/core/tests/serde/mod.rs new file mode 100644 index 000000000000..05dde7a54186 --- /dev/null +++ b/datafusion/core/tests/serde/mod.rs @@ -0,0 +1,34 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +/// Ensure `serde` feature from `arrow-schema` crate is re-exported. +#[test] +#[cfg(feature = "serde")] +fn ensure_serde_support() { + use datafusion::arrow::datatypes::DataType; + + #[derive(Debug, PartialEq, serde::Serialize, serde::Deserialize)] + struct WrappingStruct(DataType); + + let boolean = WrappingStruct(DataType::Boolean); + + let serialized = serde_json::to_string(&boolean).unwrap(); + assert_eq!("\"Boolean\"", serialized); + + let deserialized = serde_json::from_str(&serialized).unwrap(); + assert_eq!(boolean, deserialized); +} diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index 3bdc71a8eb99..e8ef34c2afe7 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/datafusion/core/tests/sql/explain_analyze.rs @@ -355,7 +355,8 @@ async fn csv_explain_verbose() { async fn csv_explain_inlist_verbose() { let ctx = SessionContext::new(); register_aggregate_csv_by_sql(&ctx).await; - let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4)"; + // Inlist len <=3 case will be transformed to OR List so we test with len=4 + let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4,5)"; let actual = execute(&ctx, sql).await; // Optimized by PreCastLitInComparisonExpressions rule @@ -368,12 +369,12 @@ async fn csv_explain_inlist_verbose() { // before optimization (Int64 literals) assert_contains!( &actual, - "aggregate_test_100.c2 IN ([Int64(1), Int64(2), Int64(4)])" + "aggregate_test_100.c2 IN ([Int64(1), Int64(2), Int64(4), Int64(5)])" ); // after optimization (casted to Int8) assert_contains!( &actual, - "aggregate_test_100.c2 IN ([Int8(1), Int8(2), Int8(4)])" + "aggregate_test_100.c2 IN ([Int8(1), Int8(2), Int8(4), Int8(5)])" ); } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index fae4b2cd82ab..915d61712074 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -700,6 +700,10 @@ impl DisplayAs for TopKExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "TopKExec: k={}", self.k) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs index 43e7ec9e45e4..87694d83bb99 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -795,7 +795,7 @@ impl ScalarUDFImpl for TakeUDF { &self.signature } fn return_type(&self, _args: &[DataType]) -> Result { - not_impl_err!("Not called because the return_type_from_exprs is implemented") + not_impl_err!("Not called because the return_type_from_args is implemented") } /// This function returns the type of the first or second argument based on @@ -1228,12 +1228,8 @@ impl ScalarUDFImpl for MyRegexUdf { } } - fn invoke_batch( - &self, - args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { - match args { + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + match args.args.as_slice() { [ColumnarValue::Scalar(ScalarValue::Utf8(value))] => { Ok(ColumnarValue::Scalar(ScalarValue::Boolean( self.matches(value.as_deref()), diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml new file mode 100644 index 000000000000..e6bb2ef4d5a9 --- /dev/null +++ b/datafusion/datasource-avro/Cargo.toml @@ -0,0 +1,60 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-datasource-avro" +description = "datafusion-datasource-avro" +authors.workspace = true +edition.workspace = true +homepage.workspace = true +license.workspace = true +readme.workspace = true +repository.workspace = true +rust-version.workspace = true +version.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +apache-avro = { workspace = true } +arrow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +chrono = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-common = { workspace = true, features = ["object_store", "avro"] } +datafusion-datasource = { workspace = true } +datafusion-execution = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } +num-traits = { version = "0.2" } +object_store = { workspace = true } +tokio = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +serde_json = { workspace = true } + +[lints] +workspace = true + +[lib] +name = "datafusion_datasource_avro" +path = "src/mod.rs" diff --git a/datafusion/datasource-avro/LICENSE.txt b/datafusion/datasource-avro/LICENSE.txt new file mode 120000 index 000000000000..1ef648f64b34 --- /dev/null +++ b/datafusion/datasource-avro/LICENSE.txt @@ -0,0 +1 @@ +../../LICENSE.txt \ No newline at end of file diff --git a/datafusion/datasource-avro/NOTICE.txt b/datafusion/datasource-avro/NOTICE.txt new file mode 120000 index 000000000000..fb051c92b10b --- /dev/null +++ b/datafusion/datasource-avro/NOTICE.txt @@ -0,0 +1 @@ +../../NOTICE.txt \ No newline at end of file diff --git a/datafusion/datasource-avro/README.md b/datafusion/datasource-avro/README.md new file mode 100644 index 000000000000..f8d7aebdcad1 --- /dev/null +++ b/datafusion/datasource-avro/README.md @@ -0,0 +1,26 @@ + + +# DataFusion datasource + +[DataFusion][df] is an extensible query execution framework, written in Rust, that uses Apache Arrow as its in-memory format. + +This crate is a submodule of DataFusion that defines a Avro based file source. + +[df]: https://crates.io/crates/datafusion diff --git a/datafusion/core/src/datasource/avro_to_arrow/arrow_array_reader.rs b/datafusion/datasource-avro/src/avro_to_arrow/arrow_array_reader.rs similarity index 99% rename from datafusion/core/src/datasource/avro_to_arrow/arrow_array_reader.rs rename to datafusion/datasource-avro/src/avro_to_arrow/arrow_array_reader.rs index 8f0e3792ffec..9a1b54b872ad 100644 --- a/datafusion/core/src/datasource/avro_to_arrow/arrow_array_reader.rs +++ b/datafusion/datasource-avro/src/avro_to_arrow/arrow_array_reader.rs @@ -17,13 +17,20 @@ //! Avro to Arrow array readers -use crate::arrow::array::{ +use apache_avro::schema::RecordSchema; +use apache_avro::{ + schema::{Schema as AvroSchema, SchemaKind}, + types::Value, + AvroResult, Error as AvroError, Reader as AvroReader, +}; +use arrow::array::{ make_array, Array, ArrayBuilder, ArrayData, ArrayDataBuilder, ArrayRef, BooleanBuilder, LargeStringArray, ListBuilder, NullArray, OffsetSizeTrait, PrimitiveArray, StringArray, StringBuilder, StringDictionaryBuilder, }; -use crate::arrow::buffer::{Buffer, MutableBuffer}; -use crate::arrow::datatypes::{ +use arrow::array::{BinaryArray, FixedSizeBinaryArray, GenericListArray}; +use arrow::buffer::{Buffer, MutableBuffer}; +use arrow::datatypes::{ ArrowDictionaryKeyType, ArrowNumericType, ArrowPrimitiveType, DataType, Date32Type, Date64Type, Field, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, Schema, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, @@ -31,21 +38,14 @@ use crate::arrow::datatypes::{ TimestampNanosecondType, TimestampSecondType, UInt16Type, UInt32Type, UInt64Type, UInt8Type, }; -use crate::arrow::error::ArrowError; -use crate::arrow::record_batch::RecordBatch; -use crate::arrow::util::bit_util; -use crate::error::{DataFusionError, Result}; -use apache_avro::schema::RecordSchema; -use apache_avro::{ - schema::{Schema as AvroSchema, SchemaKind}, - types::Value, - AvroResult, Error as AvroError, Reader as AvroReader, -}; -use arrow::array::{BinaryArray, FixedSizeBinaryArray, GenericListArray}; use arrow::datatypes::{Fields, SchemaRef}; +use arrow::error::ArrowError; use arrow::error::ArrowError::SchemaError; use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; +use arrow::util::bit_util; use datafusion_common::arrow_err; +use datafusion_common::error::{DataFusionError, Result}; use num_traits::NumCast; use std::collections::BTreeMap; use std::io::Read; @@ -1071,10 +1071,10 @@ where #[cfg(test)] mod test { - use crate::arrow::array::Array; - use crate::arrow::datatypes::{Field, TimeUnit}; - use crate::datasource::avro_to_arrow::{Reader, ReaderBuilder}; + use crate::avro_to_arrow::{Reader, ReaderBuilder}; + use arrow::array::Array; use arrow::datatypes::DataType; + use arrow::datatypes::{Field, TimeUnit}; use datafusion_common::assert_batches_eq; use datafusion_common::cast::{ as_int32_array, as_int64_array, as_list_array, as_timestamp_microsecond_array, @@ -1083,7 +1083,7 @@ mod test { use std::sync::Arc; fn build_reader(name: &str, batch_size: usize) -> Reader { - let testdata = crate::test_util::arrow_test_data(); + let testdata = datafusion_common::test_util::arrow_test_data(); let filename = format!("{testdata}/avro/{name}"); let builder = ReaderBuilder::new() .read_schema() diff --git a/datafusion/core/src/datasource/avro_to_arrow/mod.rs b/datafusion/datasource-avro/src/avro_to_arrow/mod.rs similarity index 67% rename from datafusion/core/src/datasource/avro_to_arrow/mod.rs rename to datafusion/datasource-avro/src/avro_to_arrow/mod.rs index 71184a78c96f..c1530a488020 100644 --- a/datafusion/core/src/datasource/avro_to_arrow/mod.rs +++ b/datafusion/datasource-avro/src/avro_to_arrow/mod.rs @@ -19,33 +19,21 @@ //! //! [Avro]: https://avro.apache.org/docs/1.2.0/ -#[cfg(feature = "avro")] mod arrow_array_reader; -#[cfg(feature = "avro")] mod reader; -#[cfg(feature = "avro")] mod schema; -use crate::arrow::datatypes::Schema; -use crate::error::Result; -#[cfg(feature = "avro")] +use arrow::datatypes::Schema; pub use reader::{Reader, ReaderBuilder}; -#[cfg(feature = "avro")] + pub use schema::to_arrow_schema; use std::io::Read; -#[cfg(feature = "avro")] /// Read Avro schema given a reader -pub fn read_avro_schema_from_reader(reader: &mut R) -> Result { +pub fn read_avro_schema_from_reader( + reader: &mut R, +) -> datafusion_common::Result { let avro_reader = apache_avro::Reader::new(reader)?; let schema = avro_reader.writer_schema(); to_arrow_schema(schema) } - -#[cfg(not(feature = "avro"))] -/// Read Avro schema given a reader (requires the avro feature) -pub fn read_avro_schema_from_reader(_: &mut R) -> Result { - Err(crate::error::DataFusionError::NotImplemented( - "cannot read avro schema without the 'avro' feature enabled".to_string(), - )) -} diff --git a/datafusion/core/src/datasource/avro_to_arrow/reader.rs b/datafusion/datasource-avro/src/avro_to_arrow/reader.rs similarity index 95% rename from datafusion/core/src/datasource/avro_to_arrow/reader.rs rename to datafusion/datasource-avro/src/avro_to_arrow/reader.rs index dbc24da46366..bc7b50a9cdc3 100644 --- a/datafusion/core/src/datasource/avro_to_arrow/reader.rs +++ b/datafusion/datasource-avro/src/avro_to_arrow/reader.rs @@ -16,10 +16,10 @@ // under the License. use super::arrow_array_reader::AvroArrowArrayReader; -use crate::arrow::datatypes::SchemaRef; -use crate::arrow::record_batch::RecordBatch; -use crate::error::Result; +use arrow::datatypes::SchemaRef; use arrow::error::Result as ArrowResult; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; use std::io::{Read, Seek}; use std::sync::Arc; @@ -58,7 +58,7 @@ impl ReaderBuilder { /// ``` /// use std::fs::File; /// - /// use datafusion::datasource::avro_to_arrow::{Reader, ReaderBuilder}; + /// use datafusion_datasource_avro::avro_to_arrow::{Reader, ReaderBuilder}; /// /// fn example() -> Reader<'static, File> { /// let file = File::open("test/data/basic.avro").unwrap(); @@ -170,13 +170,17 @@ impl Iterator for Reader<'_, R> { #[cfg(test)] mod tests { use super::*; - use crate::arrow::array::*; - use crate::arrow::datatypes::{DataType, Field}; + use arrow::array::*; + use arrow::array::{ + BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, + TimestampMicrosecondArray, + }; use arrow::datatypes::TimeUnit; + use arrow::datatypes::{DataType, Field}; use std::fs::File; fn build_reader(name: &str) -> Reader { - let testdata = crate::test_util::arrow_test_data(); + let testdata = datafusion_common::test_util::arrow_test_data(); let filename = format!("{testdata}/avro/{name}"); let builder = ReaderBuilder::new().read_schema().with_batch_size(64); builder.build(File::open(filename).unwrap()).unwrap() diff --git a/datafusion/core/src/datasource/avro_to_arrow/schema.rs b/datafusion/datasource-avro/src/avro_to_arrow/schema.rs similarity index 98% rename from datafusion/core/src/datasource/avro_to_arrow/schema.rs rename to datafusion/datasource-avro/src/avro_to_arrow/schema.rs index 991f648e58bd..276056c24c01 100644 --- a/datafusion/core/src/datasource/avro_to_arrow/schema.rs +++ b/datafusion/datasource-avro/src/avro_to_arrow/schema.rs @@ -15,14 +15,14 @@ // specific language governing permissions and limitations // under the License. -use crate::arrow::datatypes::{DataType, IntervalUnit, Schema, TimeUnit, UnionMode}; -use crate::error::{DataFusionError, Result}; use apache_avro::schema::{ Alias, DecimalSchema, EnumSchema, FixedSchema, Name, RecordSchema, }; use apache_avro::types::Value; use apache_avro::Schema as AvroSchema; +use arrow::datatypes::{DataType, IntervalUnit, Schema, TimeUnit, UnionMode}; use arrow::datatypes::{Field, UnionFields}; +use datafusion_common::error::{DataFusionError, Result}; use std::collections::HashMap; use std::sync::Arc; @@ -309,12 +309,12 @@ pub fn aliased( #[cfg(test)] mod test { use super::{aliased, external_props, to_arrow_schema}; - use crate::arrow::datatypes::DataType::{Binary, Float32, Float64, Timestamp, Utf8}; - use crate::arrow::datatypes::TimeUnit::Microsecond; - use crate::arrow::datatypes::{Field, Schema}; use apache_avro::schema::{Alias, EnumSchema, FixedSchema, Name, RecordSchema}; use apache_avro::Schema as AvroSchema; + use arrow::datatypes::DataType::{Binary, Float32, Float64, Timestamp, Utf8}; use arrow::datatypes::DataType::{Boolean, Int32, Int64}; + use arrow::datatypes::TimeUnit::Microsecond; + use arrow::datatypes::{Field, Schema}; fn alias(name: &str) -> Alias { Alias::new(name).unwrap() diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs new file mode 100644 index 000000000000..00a96121aa3b --- /dev/null +++ b/datafusion/datasource-avro/src/file_format.rs @@ -0,0 +1,160 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Apache Avro [`FileFormat`] abstractions + +use std::any::Any; +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use datafusion_common::{Result, Statistics}; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; + +use arrow::datatypes::Schema; +use arrow::datatypes::SchemaRef; +use async_trait::async_trait; +use datafusion_catalog::Session; +use datafusion_common::internal_err; +use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_common::GetExt; +use datafusion_common::DEFAULT_AVRO_EXTENSION; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_plan::ExecutionPlan; +use object_store::{GetResultPayload, ObjectMeta, ObjectStore}; + +use crate::avro_to_arrow::read_avro_schema_from_reader; +use crate::source::AvroSource; + +#[derive(Default)] +/// Factory struct used to create [`AvroFormat`] +pub struct AvroFormatFactory; + +impl AvroFormatFactory { + /// Creates an instance of [`AvroFormatFactory`] + pub fn new() -> Self { + Self {} + } +} + +impl FileFormatFactory for AvroFormatFactory { + fn create( + &self, + _state: &dyn Session, + _format_options: &HashMap, + ) -> Result> { + Ok(Arc::new(AvroFormat)) + } + + fn default(&self) -> Arc { + Arc::new(AvroFormat) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl fmt::Debug for AvroFormatFactory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AvroFormatFactory").finish() + } +} + +impl GetExt for AvroFormatFactory { + fn get_ext(&self) -> String { + // Removes the dot, i.e. ".avro" -> "avro" + DEFAULT_AVRO_EXTENSION[1..].to_string() + } +} + +/// Avro [`FileFormat`] implementation. +#[derive(Default, Debug)] +pub struct AvroFormat; + +#[async_trait] +impl FileFormat for AvroFormat { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_ext(&self) -> String { + AvroFormatFactory::new().get_ext() + } + + fn get_ext_with_compression( + &self, + file_compression_type: &FileCompressionType, + ) -> Result { + let ext = self.get_ext(); + match file_compression_type.get_variant() { + CompressionTypeVariant::UNCOMPRESSED => Ok(ext), + _ => internal_err!("Avro FileFormat does not support compression."), + } + } + + async fn infer_schema( + &self, + _state: &dyn Session, + store: &Arc, + objects: &[ObjectMeta], + ) -> Result { + let mut schemas = vec![]; + for object in objects { + let r = store.as_ref().get(&object.location).await?; + let schema = match r.payload { + GetResultPayload::File(mut file, _) => { + read_avro_schema_from_reader(&mut file)? + } + GetResultPayload::Stream(_) => { + // TODO: Fetching entire file to get schema is potentially wasteful + let data = r.bytes().await?; + read_avro_schema_from_reader(&mut data.as_ref())? + } + }; + schemas.push(schema); + } + let merged_schema = Schema::try_merge(schemas)?; + Ok(Arc::new(merged_schema)) + } + + async fn infer_stats( + &self, + _state: &dyn Session, + _store: &Arc, + table_schema: SchemaRef, + _object: &ObjectMeta, + ) -> Result { + Ok(Statistics::new_unknown(&table_schema)) + } + + async fn create_physical_plan( + &self, + _state: &dyn Session, + conf: FileScanConfig, + _filters: Option<&Arc>, + ) -> Result> { + Ok(conf.with_source(self.file_source()).build()) + } + + fn file_source(&self) -> Arc { + Arc::new(AvroSource::new()) + } +} diff --git a/datafusion/datasource-avro/src/mod.rs b/datafusion/datasource-avro/src/mod.rs new file mode 100644 index 000000000000..7d00b14e5119 --- /dev/null +++ b/datafusion/datasource-avro/src/mod.rs @@ -0,0 +1,30 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + +//! An [Avro](https://avro.apache.org/) based [`FileSource`](datafusion_datasource::file::FileSource) implementation and related functionality. + +pub mod avro_to_arrow; +pub mod file_format; +pub mod source; + +pub use file_format::*; diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs new file mode 100644 index 000000000000..ce3722e7b11e --- /dev/null +++ b/datafusion/datasource-avro/src/source.rs @@ -0,0 +1,282 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for reading line-delimited Avro files + +use std::any::Any; +use std::fmt::Formatter; +use std::sync::Arc; + +use crate::avro_to_arrow::Reader as AvroReader; + +use datafusion_common::error::Result; + +use arrow::datatypes::SchemaRef; +use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file_stream::FileOpener; +use datafusion_datasource::source::DataSourceExec; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, +}; + +use object_store::ObjectStore; + +/// Execution plan for scanning Avro data source +#[derive(Debug, Clone)] +#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] +pub struct AvroExec { + inner: DataSourceExec, + base_config: FileScanConfig, +} + +#[allow(unused, deprecated)] +impl AvroExec { + /// Create a new Avro reader execution plan provided base configurations + pub fn new(base_config: FileScanConfig) -> Self { + let ( + projected_schema, + projected_constraints, + projected_statistics, + projected_output_ordering, + ) = base_config.project(); + let cache = Self::compute_properties( + Arc::clone(&projected_schema), + &projected_output_ordering, + projected_constraints, + &base_config, + ); + let base_config = base_config.with_source(Arc::new(AvroSource::default())); + Self { + inner: DataSourceExec::new(Arc::new(base_config.clone())), + base_config, + } + } + + /// Ref to the base configs + pub fn base_config(&self) -> &FileScanConfig { + &self.base_config + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: &[LexOrdering], + constraints: Constraints, + file_scan_config: &FileScanConfig, + ) -> PlanProperties { + // Equivalence Properties + let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) + .with_constraints(constraints); + let n_partitions = file_scan_config.file_groups.len(); + + PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(n_partitions), // Output Partitioning + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +#[allow(unused, deprecated)] +impl DisplayAs for AvroExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(unused, deprecated)] +impl ExecutionPlan for AvroExec { + fn name(&self) -> &'static str { + "AvroExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &PlanProperties { + self.inner.properties() + } + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn statistics(&self) -> Result { + self.inner.statistics() + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + self.inner.with_fetch(limit) + } +} + +/// AvroSource holds the extra configuration that is necessary for opening avro files +#[derive(Clone, Default)] +pub struct AvroSource { + schema: Option, + batch_size: Option, + projection: Option>, + metrics: ExecutionPlanMetricsSet, + projected_statistics: Option, +} + +impl AvroSource { + /// Initialize an AvroSource with default values + pub fn new() -> Self { + Self::default() + } + + fn open(&self, reader: R) -> Result> { + AvroReader::try_new( + reader, + Arc::clone(self.schema.as_ref().expect("Schema must set before open")), + self.batch_size.expect("Batch size must set before open"), + self.projection.clone(), + ) + } +} + +impl FileSource for AvroSource { + fn create_file_opener( + &self, + object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Arc { + Arc::new(private::AvroOpener { + config: Arc::new(self.clone()), + object_store, + }) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn with_batch_size(&self, batch_size: usize) -> Arc { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn with_schema(&self, schema: SchemaRef) -> Arc { + let mut conf = self.clone(); + conf.schema = Some(schema); + Arc::new(conf) + } + fn with_statistics(&self, statistics: Statistics) -> Arc { + let mut conf = self.clone(); + conf.projected_statistics = Some(statistics); + Arc::new(conf) + } + + fn with_projection(&self, config: &FileScanConfig) -> Arc { + let mut conf = self.clone(); + conf.projection = config.projected_file_column_names(); + Arc::new(conf) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn statistics(&self) -> Result { + let statistics = &self.projected_statistics; + Ok(statistics + .clone() + .expect("projected_statistics must be set")) + } + + fn file_type(&self) -> &str { + "avro" + } + + fn repartitioned( + &self, + _target_partitions: usize, + _repartition_file_min_size: usize, + _output_ordering: Option, + _config: &FileScanConfig, + ) -> Result> { + Ok(None) + } +} + +mod private { + use super::*; + + use bytes::Buf; + use datafusion_datasource::{file_meta::FileMeta, file_stream::FileOpenFuture}; + use futures::StreamExt; + use object_store::{GetResultPayload, ObjectStore}; + + pub struct AvroOpener { + pub config: Arc, + pub object_store: Arc, + } + + impl FileOpener for AvroOpener { + fn open(&self, file_meta: FileMeta) -> Result { + let config = Arc::clone(&self.config); + let object_store = Arc::clone(&self.object_store); + Ok(Box::pin(async move { + let r = object_store.get(file_meta.location()).await?; + match r.payload { + GetResultPayload::File(file, _) => { + let reader = config.open(file)?; + Ok(futures::stream::iter(reader).boxed()) + } + GetResultPayload::Stream(_) => { + let bytes = r.bytes().await?; + let reader = config.open(bytes.reader())?; + Ok(futures::stream::iter(reader).boxed()) + } + } + })) + } + } +} diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml new file mode 100644 index 000000000000..689531758cad --- /dev/null +++ b/datafusion/datasource-csv/Cargo.toml @@ -0,0 +1,60 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-datasource-csv" +description = "datafusion-datasource-csv" +authors.workspace = true +edition.workspace = true +homepage.workspace = true +license.workspace = true +readme.workspace = true +repository.workspace = true +rust-version.workspace = true +version.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common-runtime = { workspace = true } +datafusion-datasource = { workspace = true } +datafusion-execution = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } +itertools = { workspace = true } +log = { workspace = true } +object_store = { workspace = true } +rand = { workspace = true } +regex = { workspace = true } +tokio = { workspace = true } +url = { workspace = true } + +[lints] +workspace = true + +[lib] +name = "datafusion_datasource_csv" +path = "src/mod.rs" diff --git a/datafusion/datasource-csv/LICENSE.txt b/datafusion/datasource-csv/LICENSE.txt new file mode 120000 index 000000000000..1ef648f64b34 --- /dev/null +++ b/datafusion/datasource-csv/LICENSE.txt @@ -0,0 +1 @@ +../../LICENSE.txt \ No newline at end of file diff --git a/datafusion/datasource-csv/NOTICE.txt b/datafusion/datasource-csv/NOTICE.txt new file mode 120000 index 000000000000..fb051c92b10b --- /dev/null +++ b/datafusion/datasource-csv/NOTICE.txt @@ -0,0 +1 @@ +../../NOTICE.txt \ No newline at end of file diff --git a/datafusion/datasource-csv/README.md b/datafusion/datasource-csv/README.md new file mode 100644 index 000000000000..c5944f9e438f --- /dev/null +++ b/datafusion/datasource-csv/README.md @@ -0,0 +1,26 @@ + + +# DataFusion datasource + +[DataFusion][df] is an extensible query execution framework, written in Rust, that uses Apache Arrow as its in-memory format. + +This crate is a submodule of DataFusion that defines a CSV based file source. + +[df]: https://crates.io/crates/datafusion diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs new file mode 100644 index 000000000000..cab561d163b3 --- /dev/null +++ b/datafusion/datasource-csv/src/file_format.rs @@ -0,0 +1,740 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`CsvFormat`], Comma Separated Value (CSV) [`FileFormat`] abstractions + +use std::any::Any; +use std::collections::{HashMap, HashSet}; +use std::fmt::{self, Debug}; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::csv::WriterBuilder; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::error::ArrowError; +use datafusion_catalog::Session; +use datafusion_common::config::{ConfigField, ConfigFileType, CsvOptions}; +use datafusion_common::file_options::csv_writer::CsvWriterOptions; +use datafusion_common::{ + exec_err, not_impl_err, DataFusionError, GetExt, Result, Statistics, + DEFAULT_CSV_EXTENSION, +}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_datasource::decoder::Decoder; +use datafusion_datasource::display::FileGroupDisplay; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_format::{ + FileFormat, FileFormatFactory, DEFAULT_SCHEMA_INFER_MAX_RECORD, +}; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::write::demux::DemuxedStreamReceiver; +use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_datasource::write::BatchSerializer; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::dml::InsertOp; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::LexRequirement; + +use async_trait::async_trait; +use bytes::{Buf, Bytes}; +use datafusion_physical_plan::insert::{DataSink, DataSinkExec}; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use futures::stream::BoxStream; +use futures::{pin_mut, Stream, StreamExt, TryStreamExt}; +use object_store::{delimited::newline_delimited_stream, ObjectMeta, ObjectStore}; +use regex::Regex; + +use crate::source::CsvSource; + +#[derive(Default)] +/// Factory used to create [`CsvFormat`] +pub struct CsvFormatFactory { + /// the options for csv file read + pub options: Option, +} + +impl CsvFormatFactory { + /// Creates an instance of [`CsvFormatFactory`] + pub fn new() -> Self { + Self { options: None } + } + + /// Creates an instance of [`CsvFormatFactory`] with customized default options + pub fn new_with_options(options: CsvOptions) -> Self { + Self { + options: Some(options), + } + } +} + +impl Debug for CsvFormatFactory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CsvFormatFactory") + .field("options", &self.options) + .finish() + } +} + +impl FileFormatFactory for CsvFormatFactory { + fn create( + &self, + state: &dyn Session, + format_options: &HashMap, + ) -> Result> { + let csv_options = match &self.options { + None => { + let mut table_options = state.default_table_options(); + table_options.set_config_format(ConfigFileType::CSV); + table_options.alter_with_string_hash_map(format_options)?; + table_options.csv + } + Some(csv_options) => { + let mut csv_options = csv_options.clone(); + for (k, v) in format_options { + csv_options.set(k, v)?; + } + csv_options + } + }; + + Ok(Arc::new(CsvFormat::default().with_options(csv_options))) + } + + fn default(&self) -> Arc { + Arc::new(CsvFormat::default()) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl GetExt for CsvFormatFactory { + fn get_ext(&self) -> String { + // Removes the dot, i.e. ".csv" -> "csv" + DEFAULT_CSV_EXTENSION[1..].to_string() + } +} + +/// Character Separated Value [`FileFormat`] implementation. +#[derive(Debug, Default)] +pub struct CsvFormat { + options: CsvOptions, +} + +impl CsvFormat { + /// Return a newline delimited stream from the specified file on + /// Stream, decompressing if necessary + /// Each returned `Bytes` has a whole number of newline delimited rows + async fn read_to_delimited_chunks<'a>( + &self, + store: &Arc, + object: &ObjectMeta, + ) -> BoxStream<'a, Result> { + // stream to only read as many rows as needed into memory + let stream = store + .get(&object.location) + .await + .map_err(DataFusionError::ObjectStore); + let stream = match stream { + Ok(stream) => self + .read_to_delimited_chunks_from_stream( + stream + .into_stream() + .map_err(DataFusionError::ObjectStore) + .boxed(), + ) + .await + .map_err(DataFusionError::from) + .left_stream(), + Err(e) => { + futures::stream::once(futures::future::ready(Err(e))).right_stream() + } + }; + stream.boxed() + } + + /// Convert a stream of bytes into a stream of of [`Bytes`] containing newline + /// delimited CSV records, while accounting for `\` and `"`. + pub async fn read_to_delimited_chunks_from_stream<'a>( + &self, + stream: BoxStream<'a, Result>, + ) -> BoxStream<'a, Result> { + let file_compression_type: FileCompressionType = self.options.compression.into(); + let decoder = file_compression_type.convert_stream(stream); + let stream = match decoder { + Ok(decoded_stream) => { + newline_delimited_stream(decoded_stream.map_err(|e| match e { + DataFusionError::ObjectStore(e) => e, + err => object_store::Error::Generic { + store: "read to delimited chunks failed", + source: Box::new(err), + }, + })) + .map_err(DataFusionError::from) + .left_stream() + } + Err(e) => { + futures::stream::once(futures::future::ready(Err(e))).right_stream() + } + }; + stream.boxed() + } + + /// Set the csv options + pub fn with_options(mut self, options: CsvOptions) -> Self { + self.options = options; + self + } + + /// Retrieve the csv options + pub fn options(&self) -> &CsvOptions { + &self.options + } + + /// Set a limit in terms of records to scan to infer the schema + /// - default to `DEFAULT_SCHEMA_INFER_MAX_RECORD` + pub fn with_schema_infer_max_rec(mut self, max_rec: usize) -> Self { + self.options.schema_infer_max_rec = Some(max_rec); + self + } + + /// Set true to indicate that the first line is a header. + /// - default to true + pub fn with_has_header(mut self, has_header: bool) -> Self { + self.options.has_header = Some(has_header); + self + } + + /// Set the regex to use for null values in the CSV reader. + /// - default to treat empty values as null. + pub fn with_null_regex(mut self, null_regex: Option) -> Self { + self.options.null_regex = null_regex; + self + } + + /// Returns `Some(true)` if the first line is a header, `Some(false)` if + /// it is not, and `None` if it is not specified. + pub fn has_header(&self) -> Option { + self.options.has_header + } + + /// Lines beginning with this byte are ignored. + pub fn with_comment(mut self, comment: Option) -> Self { + self.options.comment = comment; + self + } + + /// The character separating values within a row. + /// - default to ',' + pub fn with_delimiter(mut self, delimiter: u8) -> Self { + self.options.delimiter = delimiter; + self + } + + /// The quote character in a row. + /// - default to '"' + pub fn with_quote(mut self, quote: u8) -> Self { + self.options.quote = quote; + self + } + + /// The escape character in a row. + /// - default is None + pub fn with_escape(mut self, escape: Option) -> Self { + self.options.escape = escape; + self + } + + /// The character used to indicate the end of a row. + /// - default to None (CRLF) + pub fn with_terminator(mut self, terminator: Option) -> Self { + self.options.terminator = terminator; + self + } + + /// Specifies whether newlines in (quoted) values are supported. + /// + /// Parsing newlines in quoted values may be affected by execution behaviour such as + /// parallel file scanning. Setting this to `true` ensures that newlines in values are + /// parsed successfully, which may reduce performance. + /// + /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting. + pub fn with_newlines_in_values(mut self, newlines_in_values: bool) -> Self { + self.options.newlines_in_values = Some(newlines_in_values); + self + } + + /// Set a `FileCompressionType` of CSV + /// - defaults to `FileCompressionType::UNCOMPRESSED` + pub fn with_file_compression_type( + mut self, + file_compression_type: FileCompressionType, + ) -> Self { + self.options.compression = file_compression_type.into(); + self + } + + /// The delimiter character. + pub fn delimiter(&self) -> u8 { + self.options.delimiter + } + + /// The quote character. + pub fn quote(&self) -> u8 { + self.options.quote + } + + /// The escape character. + pub fn escape(&self) -> Option { + self.options.escape + } +} + +#[derive(Debug)] +pub struct CsvDecoder { + inner: arrow::csv::reader::Decoder, +} + +impl CsvDecoder { + pub fn new(decoder: arrow::csv::reader::Decoder) -> Self { + Self { inner: decoder } + } +} + +impl Decoder for CsvDecoder { + fn decode(&mut self, buf: &[u8]) -> Result { + self.inner.decode(buf) + } + + fn flush(&mut self) -> Result, ArrowError> { + self.inner.flush() + } + + fn can_flush_early(&self) -> bool { + self.inner.capacity() == 0 + } +} + +impl Debug for CsvSerializer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CsvSerializer") + .field("header", &self.header) + .finish() + } +} + +#[async_trait] +impl FileFormat for CsvFormat { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_ext(&self) -> String { + CsvFormatFactory::new().get_ext() + } + + fn get_ext_with_compression( + &self, + file_compression_type: &FileCompressionType, + ) -> Result { + let ext = self.get_ext(); + Ok(format!("{}{}", ext, file_compression_type.get_ext())) + } + + async fn infer_schema( + &self, + state: &dyn Session, + store: &Arc, + objects: &[ObjectMeta], + ) -> Result { + let mut schemas = vec![]; + + let mut records_to_read = self + .options + .schema_infer_max_rec + .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); + + for object in objects { + let stream = self.read_to_delimited_chunks(store, object).await; + let (schema, records_read) = self + .infer_schema_from_stream(state, records_to_read, stream) + .await + .map_err(|err| { + DataFusionError::Context( + format!("Error when processing CSV file {}", &object.location), + Box::new(err), + ) + })?; + records_to_read -= records_read; + schemas.push(schema); + if records_to_read == 0 { + break; + } + } + + let merged_schema = Schema::try_merge(schemas)?; + Ok(Arc::new(merged_schema)) + } + + async fn infer_stats( + &self, + _state: &dyn Session, + _store: &Arc, + table_schema: SchemaRef, + _object: &ObjectMeta, + ) -> Result { + Ok(Statistics::new_unknown(&table_schema)) + } + + async fn create_physical_plan( + &self, + state: &dyn Session, + mut conf: FileScanConfig, + _filters: Option<&Arc>, + ) -> Result> { + conf.file_compression_type = self.options.compression.into(); + // Consult configuration options for default values + let has_header = self + .options + .has_header + .unwrap_or(state.config_options().catalog.has_header); + let newlines_in_values = self + .options + .newlines_in_values + .unwrap_or(state.config_options().catalog.newlines_in_values); + conf.new_lines_in_values = newlines_in_values; + + let source = Arc::new( + CsvSource::new(has_header, self.options.delimiter, self.options.quote) + .with_escape(self.options.escape) + .with_terminator(self.options.terminator) + .with_comment(self.options.comment), + ); + Ok(conf.with_source(source).build()) + } + + async fn create_writer_physical_plan( + &self, + input: Arc, + state: &dyn Session, + conf: FileSinkConfig, + order_requirements: Option, + ) -> Result> { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for CSV"); + } + + // `has_header` and `newlines_in_values` fields of CsvOptions may inherit + // their values from session from configuration settings. To support + // this logic, writer options are built from the copy of `self.options` + // with updated values of these special fields. + let has_header = self + .options() + .has_header + .unwrap_or(state.config_options().catalog.has_header); + let newlines_in_values = self + .options() + .newlines_in_values + .unwrap_or(state.config_options().catalog.newlines_in_values); + + let options = self + .options() + .clone() + .with_has_header(has_header) + .with_newlines_in_values(newlines_in_values); + + let writer_options = CsvWriterOptions::try_from(&options)?; + + let sink = Arc::new(CsvSink::new(conf, writer_options)); + + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + } + + fn file_source(&self) -> Arc { + Arc::new(CsvSource::default()) + } +} + +impl CsvFormat { + /// Return the inferred schema reading up to records_to_read from a + /// stream of delimited chunks returning the inferred schema and the + /// number of lines that were read + pub async fn infer_schema_from_stream( + &self, + state: &dyn Session, + mut records_to_read: usize, + stream: impl Stream>, + ) -> Result<(Schema, usize)> { + let mut total_records_read = 0; + let mut column_names = vec![]; + let mut column_type_possibilities = vec![]; + let mut record_number = -1; + + pin_mut!(stream); + + while let Some(chunk) = stream.next().await.transpose()? { + record_number += 1; + let first_chunk = record_number == 0; + let mut format = arrow::csv::reader::Format::default() + .with_header( + first_chunk + && self + .options + .has_header + .unwrap_or(state.config_options().catalog.has_header), + ) + .with_delimiter(self.options.delimiter) + .with_quote(self.options.quote); + + if let Some(null_regex) = &self.options.null_regex { + let regex = Regex::new(null_regex.as_str()) + .expect("Unable to parse CSV null regex."); + format = format.with_null_regex(regex); + } + + if let Some(escape) = self.options.escape { + format = format.with_escape(escape); + } + + if let Some(comment) = self.options.comment { + format = format.with_comment(comment); + } + + let (Schema { fields, .. }, records_read) = + format.infer_schema(chunk.reader(), Some(records_to_read))?; + + records_to_read -= records_read; + total_records_read += records_read; + + if first_chunk { + // set up initial structures for recording inferred schema across chunks + (column_names, column_type_possibilities) = fields + .into_iter() + .map(|field| { + let mut possibilities = HashSet::new(); + if records_read > 0 { + // at least 1 data row read, record the inferred datatype + possibilities.insert(field.data_type().clone()); + } + (field.name().clone(), possibilities) + }) + .unzip(); + } else { + if fields.len() != column_type_possibilities.len() { + return exec_err!( + "Encountered unequal lengths between records on CSV file whilst inferring schema. \ + Expected {} fields, found {} fields at record {}", + column_type_possibilities.len(), + fields.len(), + record_number + 1 + ); + } + + column_type_possibilities.iter_mut().zip(&fields).for_each( + |(possibilities, field)| { + possibilities.insert(field.data_type().clone()); + }, + ); + } + + if records_to_read == 0 { + break; + } + } + + let schema = build_schema_helper(column_names, &column_type_possibilities); + Ok((schema, total_records_read)) + } +} + +fn build_schema_helper(names: Vec, types: &[HashSet]) -> Schema { + let fields = names + .into_iter() + .zip(types) + .map(|(field_name, data_type_possibilities)| { + // ripped from arrow::csv::reader::infer_reader_schema_with_csv_options + // determine data type based on possible types + // if there are incompatible types, use DataType::Utf8 + match data_type_possibilities.len() { + 1 => Field::new( + field_name, + data_type_possibilities.iter().next().unwrap().clone(), + true, + ), + 2 => { + if data_type_possibilities.contains(&DataType::Int64) + && data_type_possibilities.contains(&DataType::Float64) + { + // we have an integer and double, fall down to double + Field::new(field_name, DataType::Float64, true) + } else { + // default to Utf8 for conflicting datatypes (e.g bool and int) + Field::new(field_name, DataType::Utf8, true) + } + } + _ => Field::new(field_name, DataType::Utf8, true), + } + }) + .collect::(); + Schema::new(fields) +} + +impl Default for CsvSerializer { + fn default() -> Self { + Self::new() + } +} + +/// Define a struct for serializing CSV records to a stream +pub struct CsvSerializer { + // CSV writer builder + builder: WriterBuilder, + // Flag to indicate whether there will be a header + header: bool, +} + +impl CsvSerializer { + /// Constructor for the CsvSerializer object + pub fn new() -> Self { + Self { + builder: WriterBuilder::new(), + header: true, + } + } + + /// Method for setting the CSV writer builder + pub fn with_builder(mut self, builder: WriterBuilder) -> Self { + self.builder = builder; + self + } + + /// Method for setting the CSV writer header status + pub fn with_header(mut self, header: bool) -> Self { + self.header = header; + self + } +} + +impl BatchSerializer for CsvSerializer { + fn serialize(&self, batch: RecordBatch, initial: bool) -> Result { + let mut buffer = Vec::with_capacity(4096); + let builder = self.builder.clone(); + let header = self.header && initial; + let mut writer = builder.with_header(header).build(&mut buffer); + writer.write(&batch)?; + drop(writer); + Ok(Bytes::from(buffer)) + } +} + +/// Implements [`DataSink`] for writing to a CSV file. +pub struct CsvSink { + /// Config options for writing data + config: FileSinkConfig, + writer_options: CsvWriterOptions, +} + +impl Debug for CsvSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CsvSink").finish() + } +} + +impl DisplayAs for CsvSink { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CsvSink(file_groups=",)?; + FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; + write!(f, ")") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl CsvSink { + /// Create from config. + pub fn new(config: FileSinkConfig, writer_options: CsvWriterOptions) -> Self { + Self { + config, + writer_options, + } + } + + /// Retrieve the writer options + pub fn writer_options(&self) -> &CsvWriterOptions { + &self.writer_options + } +} + +#[async_trait] +impl FileSink for CsvSink { + fn config(&self) -> &FileSinkConfig { + &self.config + } + + async fn spawn_writer_tasks_and_join( + &self, + context: &Arc, + demux_task: SpawnedTask>, + file_stream_rx: DemuxedStreamReceiver, + object_store: Arc, + ) -> Result { + let builder = self.writer_options.writer_options.clone(); + let header = builder.header(); + let serializer = Arc::new( + CsvSerializer::new() + .with_builder(builder) + .with_header(header), + ) as _; + spawn_writer_tasks_and_join( + context, + serializer, + self.writer_options.compression.into(), + object_store, + demux_task, + file_stream_rx, + ) + .await + } +} + +#[async_trait] +impl DataSink for CsvSink { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> &SchemaRef { + self.config.output_schema() + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + context: &Arc, + ) -> Result { + FileSink::write_all(self, data, context).await + } +} diff --git a/datafusion/datasource-csv/src/mod.rs b/datafusion/datasource-csv/src/mod.rs new file mode 100644 index 000000000000..4117d1fee5fc --- /dev/null +++ b/datafusion/datasource-csv/src/mod.rs @@ -0,0 +1,38 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod file_format; +pub mod source; + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion_datasource::{ + file::FileSource, file_scan_config::FileScanConfig, PartitionedFile, +}; +use datafusion_execution::object_store::ObjectStoreUrl; +pub use file_format::*; + +/// Returns a [`FileScanConfig`] for given `file_groups` +pub fn partitioned_csv_config( + schema: SchemaRef, + file_groups: Vec>, + file_source: Arc, +) -> FileScanConfig { + FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, file_source) + .with_file_groups(file_groups) +} diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs new file mode 100644 index 000000000000..bb584433d1a4 --- /dev/null +++ b/datafusion/datasource-csv/src/source.rs @@ -0,0 +1,786 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for reading CSV files + +use std::any::Any; +use std::fmt; +use std::io::{Read, Seek, SeekFrom}; +use std::sync::Arc; +use std::task::Poll; + +use datafusion_datasource::decoder::{deserialize_stream, DecoderDeserializer}; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_meta::FileMeta; +use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; +use datafusion_datasource::{ + calculate_range, FileRange, ListingTableUrl, PartitionedFile, RangeCalculation, +}; + +use arrow::csv; +use arrow::datatypes::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{Constraints, DataFusionError, Result, Statistics}; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; + +use futures::{StreamExt, TryStreamExt}; +use object_store::buffered::BufWriter; +use object_store::{GetOptions, GetResultPayload, ObjectStore}; +use tokio::io::AsyncWriteExt; +use tokio::task::JoinSet; + +use crate::file_format::CsvDecoder; + +/// Old Csv source, deprecated with DataSourceExec implementation and CsvSource +/// +/// See examples on `CsvSource` +#[derive(Debug, Clone)] +#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] +pub struct CsvExec { + base_config: FileScanConfig, + inner: DataSourceExec, +} + +/// Builder for [`CsvExec`]. +/// +/// See example on [`CsvExec`]. +#[derive(Debug, Clone)] +#[deprecated(since = "46.0.0", note = "use FileScanConfig instead")] +pub struct CsvExecBuilder { + file_scan_config: FileScanConfig, + file_compression_type: FileCompressionType, + // TODO: it seems like these format options could be reused across all the various CSV config + has_header: bool, + delimiter: u8, + quote: u8, + terminator: Option, + escape: Option, + comment: Option, + newlines_in_values: bool, +} + +#[allow(unused, deprecated)] +impl CsvExecBuilder { + /// Create a new builder to read the provided file scan configuration. + pub fn new(file_scan_config: FileScanConfig) -> Self { + Self { + file_scan_config, + // TODO: these defaults are duplicated from `CsvOptions` - should they be computed? + has_header: false, + delimiter: b',', + quote: b'"', + terminator: None, + escape: None, + comment: None, + newlines_in_values: false, + file_compression_type: FileCompressionType::UNCOMPRESSED, + } + } + + /// Set whether the first row defines the column names. + /// + /// The default value is `false`. + pub fn with_has_header(mut self, has_header: bool) -> Self { + self.has_header = has_header; + self + } + + /// Set the column delimeter. + /// + /// The default is `,`. + pub fn with_delimeter(mut self, delimiter: u8) -> Self { + self.delimiter = delimiter; + self + } + + /// Set the quote character. + /// + /// The default is `"`. + pub fn with_quote(mut self, quote: u8) -> Self { + self.quote = quote; + self + } + + /// Set the line terminator. If not set, the default is CRLF. + /// + /// The default is None. + pub fn with_terminator(mut self, terminator: Option) -> Self { + self.terminator = terminator; + self + } + + /// Set the escape character. + /// + /// The default is `None` (i.e. quotes cannot be escaped). + pub fn with_escape(mut self, escape: Option) -> Self { + self.escape = escape; + self + } + + /// Set the comment character. + /// + /// The default is `None` (i.e. comments are not supported). + pub fn with_comment(mut self, comment: Option) -> Self { + self.comment = comment; + self + } + + /// Set whether newlines in (quoted) values are supported. + /// + /// Parsing newlines in quoted values may be affected by execution behaviour such as + /// parallel file scanning. Setting this to `true` ensures that newlines in values are + /// parsed successfully, which may reduce performance. + /// + /// The default value is `false`. + pub fn with_newlines_in_values(mut self, newlines_in_values: bool) -> Self { + self.newlines_in_values = newlines_in_values; + self + } + + /// Set the file compression type. + /// + /// The default is [`FileCompressionType::UNCOMPRESSED`]. + pub fn with_file_compression_type( + mut self, + file_compression_type: FileCompressionType, + ) -> Self { + self.file_compression_type = file_compression_type; + self + } + + /// Build a [`CsvExec`]. + #[must_use] + pub fn build(self) -> CsvExec { + let Self { + file_scan_config: base_config, + file_compression_type, + has_header, + delimiter, + quote, + terminator, + escape, + comment, + newlines_in_values, + } = self; + + let ( + projected_schema, + projected_constraints, + projected_statistics, + projected_output_ordering, + ) = base_config.project(); + let cache = CsvExec::compute_properties( + projected_schema, + &projected_output_ordering, + projected_constraints, + &base_config, + ); + let csv = CsvSource::new(has_header, delimiter, quote) + .with_comment(comment) + .with_escape(escape) + .with_terminator(terminator); + let base_config = base_config + .with_newlines_in_values(newlines_in_values) + .with_file_compression_type(file_compression_type) + .with_source(Arc::new(csv)); + + CsvExec { + inner: DataSourceExec::new(Arc::new(base_config.clone())), + base_config, + } + } +} + +#[allow(unused, deprecated)] +impl CsvExec { + /// Create a new CSV reader execution plan provided base and specific configurations + #[allow(clippy::too_many_arguments)] + pub fn new( + base_config: FileScanConfig, + has_header: bool, + delimiter: u8, + quote: u8, + terminator: Option, + escape: Option, + comment: Option, + newlines_in_values: bool, + file_compression_type: FileCompressionType, + ) -> Self { + CsvExecBuilder::new(base_config) + .with_has_header(has_header) + .with_delimeter(delimiter) + .with_quote(quote) + .with_terminator(terminator) + .with_escape(escape) + .with_comment(comment) + .with_newlines_in_values(newlines_in_values) + .with_file_compression_type(file_compression_type) + .build() + } + + /// Return a [`CsvExecBuilder`]. + /// + /// See example on [`CsvExec`] and [`CsvExecBuilder`] for specifying CSV table options. + pub fn builder(file_scan_config: FileScanConfig) -> CsvExecBuilder { + CsvExecBuilder::new(file_scan_config) + } + + /// Ref to the base configs + pub fn base_config(&self) -> &FileScanConfig { + &self.base_config + } + + fn file_scan_config(&self) -> FileScanConfig { + self.inner + .data_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + fn csv_source(&self) -> CsvSource { + let source = self.file_scan_config(); + source + .file_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + /// true if the first line of each file is a header + pub fn has_header(&self) -> bool { + self.csv_source().has_header() + } + + /// Specifies whether newlines in (quoted) values are supported. + /// + /// Parsing newlines in quoted values may be affected by execution behaviour such as + /// parallel file scanning. Setting this to `true` ensures that newlines in values are + /// parsed successfully, which may reduce performance. + /// + /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting. + pub fn newlines_in_values(&self) -> bool { + let source = self.file_scan_config(); + source.newlines_in_values() + } + + fn output_partitioning_helper(file_scan_config: &FileScanConfig) -> Partitioning { + Partitioning::UnknownPartitioning(file_scan_config.file_groups.len()) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: &[LexOrdering], + constraints: Constraints, + file_scan_config: &FileScanConfig, + ) -> PlanProperties { + // Equivalence Properties + let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) + .with_constraints(constraints); + + PlanProperties::new( + eq_properties, + Self::output_partitioning_helper(file_scan_config), // Output Partitioning + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + + fn with_file_groups(mut self, file_groups: Vec>) -> Self { + self.base_config.file_groups = file_groups.clone(); + let mut file_source = self.file_scan_config(); + file_source = file_source.with_file_groups(file_groups); + self.inner = self.inner.with_data_source(Arc::new(file_source)); + self + } +} + +#[allow(unused, deprecated)] +impl DisplayAs for CsvExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(unused, deprecated)] +impl ExecutionPlan for CsvExec { + fn name(&self) -> &'static str { + "CsvExec" + } + + /// Return a reference to Any that can be used for downcasting + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &PlanProperties { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + // this is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + /// Redistribute files across partitions according to their size + /// See comments on `FileGroupPartitioner` for more detail. + /// + /// Return `None` if can't get repartitioned (empty, compressed file, or `newlines_in_values` set). + fn repartitioned( + &self, + target_partitions: usize, + config: &ConfigOptions, + ) -> Result>> { + self.inner.repartitioned(target_partitions, config) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn statistics(&self) -> Result { + self.inner.statistics() + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + self.inner.with_fetch(limit) + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + self.inner.try_swapping_with_projection(projection) + } +} + +/// A Config for [`CsvOpener`] +/// +/// # Example: create a `DataSourceExec` for CSV +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::Schema; +/// # use datafusion_datasource::file_scan_config::FileScanConfig; +/// # use datafusion_datasource::PartitionedFile; +/// # use datafusion_datasource_csv::source::CsvSource; +/// # use datafusion_execution::object_store::ObjectStoreUrl; +/// # use datafusion_datasource::source::DataSourceExec; +/// +/// # let object_store_url = ObjectStoreUrl::local_filesystem(); +/// # let file_schema = Arc::new(Schema::empty()); +/// +/// let source = Arc::new(CsvSource::new( +/// true, +/// b',', +/// b'"', +/// ) +/// .with_terminator(Some(b'#') +/// )); +/// // Create a DataSourceExec for reading the first 100MB of `file1.csv` +/// let file_scan_config = FileScanConfig::new(object_store_url, file_schema, source) +/// .with_file(PartitionedFile::new("file1.csv", 100*1024*1024)) +/// .with_newlines_in_values(true); // The file contains newlines in values; +/// let exec = file_scan_config.build(); +/// ``` +#[derive(Debug, Clone, Default)] +pub struct CsvSource { + batch_size: Option, + file_schema: Option, + file_projection: Option>, + pub(crate) has_header: bool, + delimiter: u8, + quote: u8, + terminator: Option, + escape: Option, + comment: Option, + metrics: ExecutionPlanMetricsSet, + projected_statistics: Option, +} + +impl CsvSource { + /// Returns a [`CsvSource`] + pub fn new(has_header: bool, delimiter: u8, quote: u8) -> Self { + Self { + has_header, + delimiter, + quote, + ..Self::default() + } + } + + /// true if the first line of each file is a header + pub fn has_header(&self) -> bool { + self.has_header + } + /// A column delimiter + pub fn delimiter(&self) -> u8 { + self.delimiter + } + + /// The quote character + pub fn quote(&self) -> u8 { + self.quote + } + + /// The line terminator + pub fn terminator(&self) -> Option { + self.terminator + } + + /// Lines beginning with this byte are ignored. + pub fn comment(&self) -> Option { + self.comment + } + + /// The escape character + pub fn escape(&self) -> Option { + self.escape + } + + /// Initialize a CsvSource with escape + pub fn with_escape(&self, escape: Option) -> Self { + let mut conf = self.clone(); + conf.escape = escape; + conf + } + + /// Initialize a CsvSource with terminator + pub fn with_terminator(&self, terminator: Option) -> Self { + let mut conf = self.clone(); + conf.terminator = terminator; + conf + } + + /// Initialize a CsvSource with comment + pub fn with_comment(&self, comment: Option) -> Self { + let mut conf = self.clone(); + conf.comment = comment; + conf + } +} + +impl CsvSource { + fn open(&self, reader: R) -> Result> { + Ok(self.builder().build(reader)?) + } + + fn builder(&self) -> csv::ReaderBuilder { + let mut builder = csv::ReaderBuilder::new(Arc::clone( + self.file_schema + .as_ref() + .expect("Schema must be set before initializing builder"), + )) + .with_delimiter(self.delimiter) + .with_batch_size( + self.batch_size + .expect("Batch size must be set before initializing builder"), + ) + .with_header(self.has_header) + .with_quote(self.quote); + if let Some(terminator) = self.terminator { + builder = builder.with_terminator(terminator); + } + if let Some(proj) = &self.file_projection { + builder = builder.with_projection(proj.clone()); + } + if let Some(escape) = self.escape { + builder = builder.with_escape(escape) + } + if let Some(comment) = self.comment { + builder = builder.with_comment(comment); + } + + builder + } +} + +/// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`] +pub struct CsvOpener { + config: Arc, + file_compression_type: FileCompressionType, + object_store: Arc, +} + +impl CsvOpener { + /// Returns a [`CsvOpener`] + pub fn new( + config: Arc, + file_compression_type: FileCompressionType, + object_store: Arc, + ) -> Self { + Self { + config, + file_compression_type, + object_store, + } + } +} + +impl FileSource for CsvSource { + fn create_file_opener( + &self, + object_store: Arc, + base_config: &FileScanConfig, + _partition: usize, + ) -> Arc { + Arc::new(CsvOpener { + config: Arc::new(self.clone()), + file_compression_type: base_config.file_compression_type, + object_store, + }) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn with_batch_size(&self, batch_size: usize) -> Arc { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn with_schema(&self, schema: SchemaRef) -> Arc { + let mut conf = self.clone(); + conf.file_schema = Some(schema); + Arc::new(conf) + } + + fn with_statistics(&self, statistics: Statistics) -> Arc { + let mut conf = self.clone(); + conf.projected_statistics = Some(statistics); + Arc::new(conf) + } + + fn with_projection(&self, config: &FileScanConfig) -> Arc { + let mut conf = self.clone(); + conf.file_projection = config.file_column_projection_indices(); + Arc::new(conf) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + fn statistics(&self) -> Result { + let statistics = &self.projected_statistics; + Ok(statistics + .clone() + .expect("projected_statistics must be set")) + } + fn file_type(&self) -> &str { + "csv" + } + fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, ", has_header={}", self.has_header) + } + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl FileOpener for CsvOpener { + /// Open a partitioned CSV file. + /// + /// If `file_meta.range` is `None`, the entire file is opened. + /// If `file_meta.range` is `Some(FileRange {start, end})`, this signifies that the partition + /// corresponds to the byte range [start, end) within the file. + /// + /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules + /// are applied to determine which lines to read: + /// 1. The first line of the partition is the line in which the index of the first character >= `start`. + /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides. + /// + /// Examples: + /// Consider the following partitions enclosed by braces `{}`: + /// + /// {A,1,2,3,4,5,6,7,8,9\n + /// A,1,2,3,4,5,6,7,8,9\n} + /// A,1,2,3,4,5,6,7,8,9\n + /// The lines read would be: [0, 1] + /// + /// A,{1,2,3,4,5,6,7,8,9\n + /// A,1,2,3,4,5,6,7,8,9\n + /// A},1,2,3,4,5,6,7,8,9\n + /// The lines read would be: [1, 2] + fn open(&self, file_meta: FileMeta) -> Result { + // `self.config.has_header` controls whether to skip reading the 1st line header + // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle + // partition, then don't skip first line + let mut csv_has_header = self.config.has_header; + if let Some(FileRange { start, .. }) = file_meta.range { + if start != 0 { + csv_has_header = false; + } + } + + let config = CsvSource { + has_header: csv_has_header, + ..(*self.config).clone() + }; + + let file_compression_type = self.file_compression_type.to_owned(); + + if file_meta.range.is_some() { + assert!( + !file_compression_type.is_compressed(), + "Reading compressed .csv in parallel is not supported" + ); + } + + let store = Arc::clone(&self.object_store); + let terminator = self.config.terminator; + + Ok(Box::pin(async move { + // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries) + + let calculated_range = + calculate_range(&file_meta, &store, terminator).await?; + + let range = match calculated_range { + RangeCalculation::Range(None) => None, + RangeCalculation::Range(Some(range)) => Some(range.into()), + RangeCalculation::TerminateEarly => { + return Ok( + futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed() + ) + } + }; + + let options = GetOptions { + range, + ..Default::default() + }; + + let result = store.get_opts(file_meta.location(), options).await?; + + match result.payload { + GetResultPayload::File(mut file, _) => { + let is_whole_file_scanned = file_meta.range.is_none(); + let decoder = if is_whole_file_scanned { + // Don't seek if no range as breaks FIFO files + file_compression_type.convert_read(file)? + } else { + file.seek(SeekFrom::Start(result.range.start as _))?; + file_compression_type.convert_read( + file.take((result.range.end - result.range.start) as u64), + )? + }; + + Ok(futures::stream::iter(config.open(decoder)?).boxed()) + } + GetResultPayload::Stream(s) => { + let decoder = config.builder().build_decoder(); + let s = s.map_err(DataFusionError::from); + let input = file_compression_type.convert_stream(s.boxed())?.fuse(); + + Ok(deserialize_stream( + input, + DecoderDeserializer::new(CsvDecoder::new(decoder)), + )) + } + } + })) + } +} + +pub async fn plan_to_csv( + task_ctx: Arc, + plan: Arc, + path: impl AsRef, +) -> Result<()> { + let path = path.as_ref(); + let parsed = ListingTableUrl::parse(path)?; + let object_store_url = parsed.object_store(); + let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let mut join_set = JoinSet::new(); + for i in 0..plan.output_partitioning().partition_count() { + let storeref = Arc::clone(&store); + let plan: Arc = Arc::clone(&plan); + let filename = format!("{}/part-{i}.csv", parsed.prefix()); + let file = object_store::path::Path::parse(filename)?; + + let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; + join_set.spawn(async move { + let mut buf_writer = BufWriter::new(storeref, file.clone()); + let mut buffer = Vec::with_capacity(1024); + //only write headers on first iteration + let mut write_headers = true; + while let Some(batch) = stream.next().await.transpose()? { + let mut writer = csv::WriterBuilder::new() + .with_header(write_headers) + .build(buffer); + writer.write(&batch)?; + buffer = writer.into_inner(); + buf_writer.write_all(&buffer).await?; + buffer.clear(); + //prevent writing headers more than once + write_headers = false; + } + buf_writer.shutdown().await.map_err(DataFusionError::from) + }); + } + + while let Some(result) = join_set.join_next().await { + match result { + Ok(res) => res?, // propagate DataFusion error + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } + } + } + } + + Ok(()) +} diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml new file mode 100644 index 000000000000..78547c592ba3 --- /dev/null +++ b/datafusion/datasource-json/Cargo.toml @@ -0,0 +1,56 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-datasource-json" +description = "datafusion-datasource-json" +authors.workspace = true +edition.workspace = true +homepage.workspace = true +license.workspace = true +readme.workspace = true +repository.workspace = true +rust-version.workspace = true +version.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-common = { workspace = true, features = ["object_store"] } +datafusion-common-runtime = { workspace = true } +datafusion-datasource = { workspace = true } +datafusion-execution = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } +object_store = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } + +[lints] +workspace = true + +[lib] +name = "datafusion_datasource_json" +path = "src/mod.rs" diff --git a/datafusion/datasource-json/LICENSE.txt b/datafusion/datasource-json/LICENSE.txt new file mode 120000 index 000000000000..1ef648f64b34 --- /dev/null +++ b/datafusion/datasource-json/LICENSE.txt @@ -0,0 +1 @@ +../../LICENSE.txt \ No newline at end of file diff --git a/datafusion/datasource-json/NOTICE.txt b/datafusion/datasource-json/NOTICE.txt new file mode 120000 index 000000000000..fb051c92b10b --- /dev/null +++ b/datafusion/datasource-json/NOTICE.txt @@ -0,0 +1 @@ +../../NOTICE.txt \ No newline at end of file diff --git a/datafusion/datasource-json/README.md b/datafusion/datasource-json/README.md new file mode 100644 index 000000000000..64181814736d --- /dev/null +++ b/datafusion/datasource-json/README.md @@ -0,0 +1,26 @@ + + +# DataFusion datasource + +[DataFusion][df] is an extensible query execution framework, written in Rust, that uses Apache Arrow as its in-memory format. + +This crate is a submodule of DataFusion that defines a JSON based file source. + +[df]: https://crates.io/crates/datafusion diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs new file mode 100644 index 000000000000..bec3a524f657 --- /dev/null +++ b/datafusion/datasource-json/src/file_format.rs @@ -0,0 +1,420 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`JsonFormat`]: Line delimited JSON [`FileFormat`] abstractions + +use std::any::Any; +use std::collections::HashMap; +use std::fmt; +use std::fmt::Debug; +use std::io::BufReader; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::json; +use arrow::json::reader::{infer_json_schema_from_iterator, ValueIter}; +use datafusion_catalog::Session; +use datafusion_common::config::{ConfigField, ConfigFileType, JsonOptions}; +use datafusion_common::file_options::json_writer::JsonWriterOptions; +use datafusion_common::{ + not_impl_err, GetExt, Result, Statistics, DEFAULT_JSON_EXTENSION, +}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_datasource::decoder::Decoder; +use datafusion_datasource::display::FileGroupDisplay; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_format::{ + FileFormat, FileFormatFactory, DEFAULT_SCHEMA_INFER_MAX_RECORD, +}; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::write::demux::DemuxedStreamReceiver; +use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_datasource::write::BatchSerializer; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::dml::InsertOp; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_plan::insert::{DataSink, DataSinkExec}; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; + +use async_trait::async_trait; +use bytes::{Buf, Bytes}; +use datafusion_physical_expr_common::sort_expr::LexRequirement; +use object_store::{GetResultPayload, ObjectMeta, ObjectStore}; + +use crate::source::JsonSource; + +#[derive(Default)] +/// Factory struct used to create [JsonFormat] +pub struct JsonFormatFactory { + /// the options carried by format factory + pub options: Option, +} + +impl JsonFormatFactory { + /// Creates an instance of [JsonFormatFactory] + pub fn new() -> Self { + Self { options: None } + } + + /// Creates an instance of [JsonFormatFactory] with customized default options + pub fn new_with_options(options: JsonOptions) -> Self { + Self { + options: Some(options), + } + } +} + +impl FileFormatFactory for JsonFormatFactory { + fn create( + &self, + state: &dyn Session, + format_options: &HashMap, + ) -> Result> { + let json_options = match &self.options { + None => { + let mut table_options = state.default_table_options(); + table_options.set_config_format(ConfigFileType::JSON); + table_options.alter_with_string_hash_map(format_options)?; + table_options.json + } + Some(json_options) => { + let mut json_options = json_options.clone(); + for (k, v) in format_options { + json_options.set(k, v)?; + } + json_options + } + }; + + Ok(Arc::new(JsonFormat::default().with_options(json_options))) + } + + fn default(&self) -> Arc { + Arc::new(JsonFormat::default()) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl GetExt for JsonFormatFactory { + fn get_ext(&self) -> String { + // Removes the dot, i.e. ".parquet" -> "parquet" + DEFAULT_JSON_EXTENSION[1..].to_string() + } +} + +impl Debug for JsonFormatFactory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JsonFormatFactory") + .field("options", &self.options) + .finish() + } +} + +/// New line delimited JSON `FileFormat` implementation. +#[derive(Debug, Default)] +pub struct JsonFormat { + options: JsonOptions, +} + +impl JsonFormat { + /// Set JSON options + pub fn with_options(mut self, options: JsonOptions) -> Self { + self.options = options; + self + } + + /// Retrieve JSON options + pub fn options(&self) -> &JsonOptions { + &self.options + } + + /// Set a limit in terms of records to scan to infer the schema + /// - defaults to `DEFAULT_SCHEMA_INFER_MAX_RECORD` + pub fn with_schema_infer_max_rec(mut self, max_rec: usize) -> Self { + self.options.schema_infer_max_rec = Some(max_rec); + self + } + + /// Set a [`FileCompressionType`] of JSON + /// - defaults to `FileCompressionType::UNCOMPRESSED` + pub fn with_file_compression_type( + mut self, + file_compression_type: FileCompressionType, + ) -> Self { + self.options.compression = file_compression_type.into(); + self + } +} + +#[async_trait] +impl FileFormat for JsonFormat { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_ext(&self) -> String { + JsonFormatFactory::new().get_ext() + } + + fn get_ext_with_compression( + &self, + file_compression_type: &FileCompressionType, + ) -> Result { + let ext = self.get_ext(); + Ok(format!("{}{}", ext, file_compression_type.get_ext())) + } + + async fn infer_schema( + &self, + _state: &dyn Session, + store: &Arc, + objects: &[ObjectMeta], + ) -> Result { + let mut schemas = Vec::new(); + let mut records_to_read = self + .options + .schema_infer_max_rec + .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); + let file_compression_type = FileCompressionType::from(self.options.compression); + for object in objects { + let mut take_while = || { + let should_take = records_to_read > 0; + if should_take { + records_to_read -= 1; + } + should_take + }; + + let r = store.as_ref().get(&object.location).await?; + let schema = match r.payload { + GetResultPayload::File(file, _) => { + let decoder = file_compression_type.convert_read(file)?; + let mut reader = BufReader::new(decoder); + let iter = ValueIter::new(&mut reader, None); + infer_json_schema_from_iterator(iter.take_while(|_| take_while()))? + } + GetResultPayload::Stream(_) => { + let data = r.bytes().await?; + let decoder = file_compression_type.convert_read(data.reader())?; + let mut reader = BufReader::new(decoder); + let iter = ValueIter::new(&mut reader, None); + infer_json_schema_from_iterator(iter.take_while(|_| take_while()))? + } + }; + + schemas.push(schema); + if records_to_read == 0 { + break; + } + } + + let schema = Schema::try_merge(schemas)?; + Ok(Arc::new(schema)) + } + + async fn infer_stats( + &self, + _state: &dyn Session, + _store: &Arc, + table_schema: SchemaRef, + _object: &ObjectMeta, + ) -> Result { + Ok(Statistics::new_unknown(&table_schema)) + } + + async fn create_physical_plan( + &self, + _state: &dyn Session, + mut conf: FileScanConfig, + _filters: Option<&Arc>, + ) -> Result> { + let source = Arc::new(JsonSource::new()); + conf.file_compression_type = FileCompressionType::from(self.options.compression); + Ok(conf.with_source(source).build()) + } + + async fn create_writer_physical_plan( + &self, + input: Arc, + _state: &dyn Session, + conf: FileSinkConfig, + order_requirements: Option, + ) -> Result> { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for Json"); + } + + let writer_options = JsonWriterOptions::try_from(&self.options)?; + + let sink = Arc::new(JsonSink::new(conf, writer_options)); + + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + } + + fn file_source(&self) -> Arc { + Arc::new(JsonSource::default()) + } +} + +impl Default for JsonSerializer { + fn default() -> Self { + Self::new() + } +} + +/// Define a struct for serializing Json records to a stream +pub struct JsonSerializer {} + +impl JsonSerializer { + /// Constructor for the JsonSerializer object + pub fn new() -> Self { + Self {} + } +} + +impl BatchSerializer for JsonSerializer { + fn serialize(&self, batch: RecordBatch, _initial: bool) -> Result { + let mut buffer = Vec::with_capacity(4096); + let mut writer = json::LineDelimitedWriter::new(&mut buffer); + writer.write(&batch)?; + Ok(Bytes::from(buffer)) + } +} + +/// Implements [`DataSink`] for writing to a Json file. +pub struct JsonSink { + /// Config options for writing data + config: FileSinkConfig, + /// Writer options for underlying Json writer + writer_options: JsonWriterOptions, +} + +impl Debug for JsonSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("JsonSink").finish() + } +} + +impl DisplayAs for JsonSink { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "JsonSink(file_groups=",)?; + FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; + write!(f, ")") + } + DisplayFormatType::TreeRender => { + if !self.config.file_groups.is_empty() { + FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; + } + Ok(()) + } + } + } +} + +impl JsonSink { + /// Create from config. + pub fn new(config: FileSinkConfig, writer_options: JsonWriterOptions) -> Self { + Self { + config, + writer_options, + } + } + + /// Retrieve the writer options + pub fn writer_options(&self) -> &JsonWriterOptions { + &self.writer_options + } +} + +#[async_trait] +impl FileSink for JsonSink { + fn config(&self) -> &FileSinkConfig { + &self.config + } + + async fn spawn_writer_tasks_and_join( + &self, + context: &Arc, + demux_task: SpawnedTask>, + file_stream_rx: DemuxedStreamReceiver, + object_store: Arc, + ) -> Result { + let serializer = Arc::new(JsonSerializer::new()) as _; + spawn_writer_tasks_and_join( + context, + serializer, + self.writer_options.compression.into(), + object_store, + demux_task, + file_stream_rx, + ) + .await + } +} + +#[async_trait] +impl DataSink for JsonSink { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> &SchemaRef { + self.config.output_schema() + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + context: &Arc, + ) -> Result { + FileSink::write_all(self, data, context).await + } +} + +#[derive(Debug)] +pub struct JsonDecoder { + inner: json::reader::Decoder, +} + +impl JsonDecoder { + pub fn new(decoder: json::reader::Decoder) -> Self { + Self { inner: decoder } + } +} + +impl Decoder for JsonDecoder { + fn decode(&mut self, buf: &[u8]) -> Result { + self.inner.decode(buf) + } + + fn flush(&mut self) -> Result, ArrowError> { + self.inner.flush() + } + + fn can_flush_early(&self) -> bool { + false + } +} diff --git a/datafusion/datasource-json/src/mod.rs b/datafusion/datasource-json/src/mod.rs new file mode 100644 index 000000000000..35dabfa109fc --- /dev/null +++ b/datafusion/datasource-json/src/mod.rs @@ -0,0 +1,21 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod file_format; +pub mod source; + +pub use file_format::*; diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs new file mode 100644 index 000000000000..249593587b82 --- /dev/null +++ b/datafusion/datasource-json/src/source.rs @@ -0,0 +1,440 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for reading line-delimited JSON files + +use std::any::Any; +use std::io::{BufReader, Read, Seek, SeekFrom}; +use std::sync::Arc; +use std::task::Poll; + +use crate::file_format::JsonDecoder; + +use datafusion_common::error::{DataFusionError, Result}; +use datafusion_datasource::decoder::{deserialize_stream, DecoderDeserializer}; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_meta::FileMeta; +use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; +use datafusion_datasource::{ + calculate_range, ListingTableUrl, PartitionedFile, RangeCalculation, +}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +use arrow::json::ReaderBuilder; +use arrow::{datatypes::SchemaRef, json}; +use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; + +use futures::{StreamExt, TryStreamExt}; +use object_store::buffered::BufWriter; +use object_store::{GetOptions, GetResultPayload, ObjectStore}; +use tokio::io::AsyncWriteExt; +use tokio::task::JoinSet; + +/// Execution plan for scanning NdJson data source +#[derive(Debug, Clone)] +#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] +pub struct NdJsonExec { + inner: DataSourceExec, + base_config: FileScanConfig, + file_compression_type: FileCompressionType, +} + +#[allow(unused, deprecated)] +impl NdJsonExec { + /// Create a new JSON reader execution plan provided base configurations + pub fn new( + base_config: FileScanConfig, + file_compression_type: FileCompressionType, + ) -> Self { + let ( + projected_schema, + projected_constraints, + projected_statistics, + projected_output_ordering, + ) = base_config.project(); + let cache = Self::compute_properties( + projected_schema, + &projected_output_ordering, + projected_constraints, + &base_config, + ); + + let json = JsonSource::default(); + let base_config = base_config + .with_file_compression_type(file_compression_type) + .with_source(Arc::new(json)); + + Self { + inner: DataSourceExec::new(Arc::new(base_config.clone())), + file_compression_type: base_config.file_compression_type, + base_config, + } + } + + /// Ref to the base configs + pub fn base_config(&self) -> &FileScanConfig { + &self.base_config + } + + /// Ref to file compression type + pub fn file_compression_type(&self) -> &FileCompressionType { + &self.file_compression_type + } + + fn file_scan_config(&self) -> FileScanConfig { + self.inner + .data_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + fn json_source(&self) -> JsonSource { + let source = self.file_scan_config(); + source + .file_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + fn output_partitioning_helper(file_scan_config: &FileScanConfig) -> Partitioning { + Partitioning::UnknownPartitioning(file_scan_config.file_groups.len()) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: &[LexOrdering], + constraints: Constraints, + file_scan_config: &FileScanConfig, + ) -> PlanProperties { + // Equivalence Properties + let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings) + .with_constraints(constraints); + + PlanProperties::new( + eq_properties, + Self::output_partitioning_helper(file_scan_config), // Output Partitioning + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + + fn with_file_groups(mut self, file_groups: Vec>) -> Self { + self.base_config.file_groups = file_groups.clone(); + let mut file_source = self.file_scan_config(); + file_source = file_source.with_file_groups(file_groups); + self.inner = self.inner.with_data_source(Arc::new(file_source)); + self + } +} + +#[allow(unused, deprecated)] +impl DisplayAs for NdJsonExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(unused, deprecated)] +impl ExecutionPlan for NdJsonExec { + fn name(&self) -> &'static str { + "NdJsonExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + fn properties(&self) -> &PlanProperties { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn repartitioned( + &self, + target_partitions: usize, + config: &datafusion_common::config::ConfigOptions, + ) -> Result>> { + self.inner.repartitioned(target_partitions, config) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + fn statistics(&self) -> Result { + self.inner.statistics() + } + + fn metrics(&self) -> Option { + self.inner.metrics() + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + self.inner.with_fetch(limit) + } +} + +/// A [`FileOpener`] that opens a JSON file and yields a [`FileOpenFuture`] +pub struct JsonOpener { + batch_size: usize, + projected_schema: SchemaRef, + file_compression_type: FileCompressionType, + object_store: Arc, +} + +impl JsonOpener { + /// Returns a [`JsonOpener`] + pub fn new( + batch_size: usize, + projected_schema: SchemaRef, + file_compression_type: FileCompressionType, + object_store: Arc, + ) -> Self { + Self { + batch_size, + projected_schema, + file_compression_type, + object_store, + } + } +} + +/// JsonSource holds the extra configuration that is necessary for [`JsonOpener`] +#[derive(Clone, Default)] +pub struct JsonSource { + batch_size: Option, + metrics: ExecutionPlanMetricsSet, + projected_statistics: Option, +} + +impl JsonSource { + /// Initialize a JsonSource with default values + pub fn new() -> Self { + Self::default() + } +} + +impl FileSource for JsonSource { + fn create_file_opener( + &self, + object_store: Arc, + base_config: &FileScanConfig, + _partition: usize, + ) -> Arc { + Arc::new(JsonOpener { + batch_size: self + .batch_size + .expect("Batch size must set before creating opener"), + projected_schema: base_config.projected_file_schema(), + file_compression_type: base_config.file_compression_type, + object_store, + }) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn with_batch_size(&self, batch_size: usize) -> Arc { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn with_schema(&self, _schema: SchemaRef) -> Arc { + Arc::new(Self { ..self.clone() }) + } + fn with_statistics(&self, statistics: Statistics) -> Arc { + let mut conf = self.clone(); + conf.projected_statistics = Some(statistics); + Arc::new(conf) + } + + fn with_projection(&self, _config: &FileScanConfig) -> Arc { + Arc::new(Self { ..self.clone() }) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn statistics(&self) -> Result { + let statistics = &self.projected_statistics; + Ok(statistics + .clone() + .expect("projected_statistics must be set to call")) + } + + fn file_type(&self) -> &str { + "json" + } +} + +impl FileOpener for JsonOpener { + /// Open a partitioned NDJSON file. + /// + /// If `file_meta.range` is `None`, the entire file is opened. + /// Else `file_meta.range` is `Some(FileRange{start, end})`, which corresponds to the byte range [start, end) within the file. + /// + /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules + /// are applied to determine which lines to read: + /// 1. The first line of the partition is the line in which the index of the first character >= `start`. + /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides. + fn open(&self, file_meta: FileMeta) -> Result { + let store = Arc::clone(&self.object_store); + let schema = Arc::clone(&self.projected_schema); + let batch_size = self.batch_size; + let file_compression_type = self.file_compression_type.to_owned(); + + Ok(Box::pin(async move { + let calculated_range = calculate_range(&file_meta, &store, None).await?; + + let range = match calculated_range { + RangeCalculation::Range(None) => None, + RangeCalculation::Range(Some(range)) => Some(range.into()), + RangeCalculation::TerminateEarly => { + return Ok( + futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed() + ) + } + }; + + let options = GetOptions { + range, + ..Default::default() + }; + + let result = store.get_opts(file_meta.location(), options).await?; + + match result.payload { + GetResultPayload::File(mut file, _) => { + let bytes = match file_meta.range { + None => file_compression_type.convert_read(file)?, + Some(_) => { + file.seek(SeekFrom::Start(result.range.start as _))?; + let limit = result.range.end - result.range.start; + file_compression_type.convert_read(file.take(limit as u64))? + } + }; + + let reader = ReaderBuilder::new(schema) + .with_batch_size(batch_size) + .build(BufReader::new(bytes))?; + + Ok(futures::stream::iter(reader).boxed()) + } + GetResultPayload::Stream(s) => { + let s = s.map_err(DataFusionError::from); + + let decoder = ReaderBuilder::new(schema) + .with_batch_size(batch_size) + .build_decoder()?; + let input = file_compression_type.convert_stream(s.boxed())?.fuse(); + + Ok(deserialize_stream( + input, + DecoderDeserializer::new(JsonDecoder::new(decoder)), + )) + } + } + })) + } +} + +pub async fn plan_to_json( + task_ctx: Arc, + plan: Arc, + path: impl AsRef, +) -> Result<()> { + let path = path.as_ref(); + let parsed = ListingTableUrl::parse(path)?; + let object_store_url = parsed.object_store(); + let store = task_ctx.runtime_env().object_store(&object_store_url)?; + let mut join_set = JoinSet::new(); + for i in 0..plan.output_partitioning().partition_count() { + let storeref = Arc::clone(&store); + let plan: Arc = Arc::clone(&plan); + let filename = format!("{}/part-{i}.json", parsed.prefix()); + let file = object_store::path::Path::parse(filename)?; + + let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; + join_set.spawn(async move { + let mut buf_writer = BufWriter::new(storeref, file.clone()); + + let mut buffer = Vec::with_capacity(1024); + while let Some(batch) = stream.next().await.transpose()? { + let mut writer = json::LineDelimitedWriter::new(buffer); + writer.write(&batch)?; + buffer = writer.into_inner(); + buf_writer.write_all(&buffer).await?; + buffer.clear(); + } + + buf_writer.shutdown().await.map_err(DataFusionError::from) + }); + } + + while let Some(result) = join_set.join_next().await { + match result { + Ok(res) => res?, // propagate DataFusion error + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } + } + } + } + + Ok(()) +} diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml new file mode 100644 index 000000000000..8aa041b7a4a7 --- /dev/null +++ b/datafusion/datasource-parquet/Cargo.toml @@ -0,0 +1,65 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-datasource-parquet" +description = "datafusion-datasource-parquet" +authors.workspace = true +edition.workspace = true +homepage.workspace = true +license.workspace = true +readme.workspace = true +repository.workspace = true +rust-version.workspace = true +version.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +datafusion-catalog = { workspace = true } +datafusion-common = { workspace = true, features = ["object_store", "parquet"] } +datafusion-common-runtime = { workspace = true } +datafusion-datasource = { workspace = true, features = ["parquet"] } +datafusion-execution = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-functions-aggregate = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-physical-optimizer = { workspace = true } +datafusion-physical-plan = { workspace = true } +futures = { workspace = true } +itertools = { workspace = true } +log = { workspace = true } +object_store = { workspace = true } +parking_lot = { workspace = true } +parquet = { workspace = true } +rand = { workspace = true } +tokio = { workspace = true } + +[dev-dependencies] +chrono = { workspace = true } + +[lints] +workspace = true + +[lib] +name = "datafusion_datasource_parquet" +path = "src/mod.rs" diff --git a/datafusion/datasource-parquet/LICENSE.txt b/datafusion/datasource-parquet/LICENSE.txt new file mode 120000 index 000000000000..1ef648f64b34 --- /dev/null +++ b/datafusion/datasource-parquet/LICENSE.txt @@ -0,0 +1 @@ +../../LICENSE.txt \ No newline at end of file diff --git a/datafusion/datasource-parquet/NOTICE.txt b/datafusion/datasource-parquet/NOTICE.txt new file mode 120000 index 000000000000..fb051c92b10b --- /dev/null +++ b/datafusion/datasource-parquet/NOTICE.txt @@ -0,0 +1 @@ +../../NOTICE.txt \ No newline at end of file diff --git a/datafusion/datasource-parquet/README.md b/datafusion/datasource-parquet/README.md new file mode 100644 index 000000000000..abcdd5ab1340 --- /dev/null +++ b/datafusion/datasource-parquet/README.md @@ -0,0 +1,26 @@ + + +# DataFusion datasource + +[DataFusion][df] is an extensible query execution framework, written in Rust, that uses Apache Arrow as its in-memory format. + +This crate is a submodule of DataFusion that defines a Parquet based file source. + +[df]: https://crates.io/crates/datafusion diff --git a/datafusion/core/src/datasource/physical_plan/parquet/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs similarity index 99% rename from datafusion/core/src/datasource/physical_plan/parquet/access_plan.rs rename to datafusion/datasource-parquet/src/access_plan.rs index d30549708bbd..0c30f3ff85b6 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -35,7 +35,7 @@ use parquet::file::metadata::RowGroupMetaData; /// /// ```rust /// # use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; -/// # use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; +/// # use datafusion_datasource_parquet::ParquetAccessPlan; /// // Default to scan all row groups /// let mut access_plan = ParquetAccessPlan::new_all(4); /// access_plan.skip(0); // skip row group diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs new file mode 100644 index 000000000000..232dd2fbe31c --- /dev/null +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -0,0 +1,1403 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ParquetFormat`]: Parquet [`FileFormat`] abstractions + +use std::any::Any; +use std::fmt; +use std::fmt::Debug; +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::datatypes::{Fields, Schema, SchemaRef}; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::write::{create_writer, get_writer_schema, SharedBuffer}; + +use datafusion_datasource::file_format::{ + FileFormat, FileFormatFactory, FilePushdownSupport, +}; +use datafusion_datasource::write::demux::DemuxedStreamReceiver; + +use arrow::compute::sum; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_catalog::Session; +use datafusion_common::config::{ConfigField, ConfigFileType, TableParquetOptions}; +use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_common::stats::Precision; +use datafusion_common::{ + internal_datafusion_err, internal_err, not_impl_err, ColumnStatistics, + DataFusionError, GetExt, Result, DEFAULT_PARQUET_EXTENSION, +}; +use datafusion_common::{HashMap, Statistics}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_datasource::display::FileGroupDisplay; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::dml::InsertOp; +use datafusion_expr::Expr; +use datafusion_functions_aggregate::min_max::{MaxAccumulator, MinAccumulator}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::Accumulator; + +use async_trait::async_trait; +use bytes::Bytes; +use datafusion_physical_plan::insert::{DataSink, DataSinkExec}; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use futures::future::BoxFuture; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use log::debug; +use object_store::buffered::BufWriter; +use object_store::path::Path; +use object_store::{ObjectMeta, ObjectStore}; +use parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use parquet::arrow::arrow_writer::{ + compute_leaves, get_column_writers, ArrowColumnChunk, ArrowColumnWriter, + ArrowLeafColumn, ArrowWriterOptions, +}; +use parquet::arrow::async_reader::MetadataFetch; +use parquet::arrow::{parquet_to_arrow_schema, ArrowSchemaConverter, AsyncArrowWriter}; +use parquet::errors::ParquetError; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData}; +use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder}; +use parquet::file::writer::SerializedFileWriter; +use parquet::format::FileMetaData; +use tokio::io::{AsyncWrite, AsyncWriteExt}; +use tokio::sync::mpsc::{self, Receiver, Sender}; +use tokio::task::JoinSet; + +use crate::can_expr_be_pushed_down_with_schemas; +use crate::source::ParquetSource; + +/// Initial writing buffer size. Note this is just a size hint for efficiency. It +/// will grow beyond the set value if needed. +const INITIAL_BUFFER_BYTES: usize = 1048576; + +/// When writing parquet files in parallel, if the buffered Parquet data exceeds +/// this size, it is flushed to object store +const BUFFER_FLUSH_BYTES: usize = 1024000; + +#[derive(Default)] +/// Factory struct used to create [ParquetFormat] +pub struct ParquetFormatFactory { + /// inner options for parquet + pub options: Option, +} + +impl ParquetFormatFactory { + /// Creates an instance of [ParquetFormatFactory] + pub fn new() -> Self { + Self { options: None } + } + + /// Creates an instance of [ParquetFormatFactory] with customized default options + pub fn new_with_options(options: TableParquetOptions) -> Self { + Self { + options: Some(options), + } + } +} + +impl FileFormatFactory for ParquetFormatFactory { + fn create( + &self, + state: &dyn Session, + format_options: &std::collections::HashMap, + ) -> Result> { + let parquet_options = match &self.options { + None => { + let mut table_options = state.default_table_options(); + table_options.set_config_format(ConfigFileType::PARQUET); + table_options.alter_with_string_hash_map(format_options)?; + table_options.parquet + } + Some(parquet_options) => { + let mut parquet_options = parquet_options.clone(); + for (k, v) in format_options { + parquet_options.set(k, v)?; + } + parquet_options + } + }; + + Ok(Arc::new( + ParquetFormat::default().with_options(parquet_options), + )) + } + + fn default(&self) -> Arc { + Arc::new(ParquetFormat::default()) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +impl GetExt for ParquetFormatFactory { + fn get_ext(&self) -> String { + // Removes the dot, i.e. ".parquet" -> "parquet" + DEFAULT_PARQUET_EXTENSION[1..].to_string() + } +} + +impl Debug for ParquetFormatFactory { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ParquetFormatFactory") + .field("ParquetFormatFactory", &self.options) + .finish() + } +} +/// The Apache Parquet `FileFormat` implementation +#[derive(Debug, Default)] +pub struct ParquetFormat { + options: TableParquetOptions, +} + +impl ParquetFormat { + /// Construct a new Format with no local overrides + pub fn new() -> Self { + Self::default() + } + + /// Activate statistics based row group level pruning + /// - If `None`, defaults to value on `config_options` + pub fn with_enable_pruning(mut self, enable: bool) -> Self { + self.options.global.pruning = enable; + self + } + + /// Return `true` if pruning is enabled + pub fn enable_pruning(&self) -> bool { + self.options.global.pruning + } + + /// Provide a hint to the size of the file metadata. If a hint is provided + /// the reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. + /// Without a hint, two read are required. One read to fetch the 8-byte parquet footer and then + /// another read to fetch the metadata length encoded in the footer. + /// + /// - If `None`, defaults to value on `config_options` + pub fn with_metadata_size_hint(mut self, size_hint: Option) -> Self { + self.options.global.metadata_size_hint = size_hint; + self + } + + /// Return the metadata size hint if set + pub fn metadata_size_hint(&self) -> Option { + self.options.global.metadata_size_hint + } + + /// Tell the parquet reader to skip any metadata that may be in + /// the file Schema. This can help avoid schema conflicts due to + /// metadata. + /// + /// - If `None`, defaults to value on `config_options` + pub fn with_skip_metadata(mut self, skip_metadata: bool) -> Self { + self.options.global.skip_metadata = skip_metadata; + self + } + + /// Returns `true` if schema metadata will be cleared prior to + /// schema merging. + pub fn skip_metadata(&self) -> bool { + self.options.global.skip_metadata + } + + /// Set Parquet options for the ParquetFormat + pub fn with_options(mut self, options: TableParquetOptions) -> Self { + self.options = options; + self + } + + /// Parquet options + pub fn options(&self) -> &TableParquetOptions { + &self.options + } + + /// Return `true` if should use view types. + /// + /// If this returns true, DataFusion will instruct the parquet reader + /// to read string / binary columns using view `StringView` or `BinaryView` + /// if the table schema specifies those types, regardless of any embedded metadata + /// that may specify an alternate Arrow type. The parquet reader is optimized + /// for reading `StringView` and `BinaryView` and such queries are significantly faster. + /// + /// If this returns false, the parquet reader will read the columns according to the + /// defaults or any embedded Arrow type information. This may result in reading + /// `StringArrays` and then casting to `StringViewArray` which is less efficient. + pub fn force_view_types(&self) -> bool { + self.options.global.schema_force_view_types + } + + /// If true, will use view types. See [`Self::force_view_types`] for details + pub fn with_force_view_types(mut self, use_views: bool) -> Self { + self.options.global.schema_force_view_types = use_views; + self + } + + /// Return `true` if binary types will be read as strings. + /// + /// If this returns true, DataFusion will instruct the parquet reader + /// to read binary columns such as `Binary` or `BinaryView` as the + /// corresponding string type such as `Utf8` or `LargeUtf8`. + /// The parquet reader has special optimizations for `Utf8` and `LargeUtf8` + /// validation, and such queries are significantly faster than reading + /// binary columns and then casting to string columns. + pub fn binary_as_string(&self) -> bool { + self.options.global.binary_as_string + } + + /// If true, will read binary types as strings. See [`Self::binary_as_string`] for details + pub fn with_binary_as_string(mut self, binary_as_string: bool) -> Self { + self.options.global.binary_as_string = binary_as_string; + self + } +} + +/// Clears all metadata (Schema level and field level) on an iterator +/// of Schemas +fn clear_metadata( + schemas: impl IntoIterator, +) -> impl Iterator { + schemas.into_iter().map(|schema| { + let fields = schema + .fields() + .iter() + .map(|field| { + field.as_ref().clone().with_metadata(Default::default()) // clear meta + }) + .collect::(); + Schema::new(fields) + }) +} + +async fn fetch_schema_with_location( + store: &dyn ObjectStore, + file: &ObjectMeta, + metadata_size_hint: Option, +) -> Result<(Path, Schema)> { + let loc_path = file.location.clone(); + let schema = fetch_schema(store, file, metadata_size_hint).await?; + Ok((loc_path, schema)) +} + +#[async_trait] +impl FileFormat for ParquetFormat { + fn as_any(&self) -> &dyn Any { + self + } + + fn get_ext(&self) -> String { + ParquetFormatFactory::new().get_ext() + } + + fn get_ext_with_compression( + &self, + file_compression_type: &FileCompressionType, + ) -> Result { + let ext = self.get_ext(); + match file_compression_type.get_variant() { + CompressionTypeVariant::UNCOMPRESSED => Ok(ext), + _ => internal_err!("Parquet FileFormat does not support compression."), + } + } + + async fn infer_schema( + &self, + state: &dyn Session, + store: &Arc, + objects: &[ObjectMeta], + ) -> Result { + let mut schemas: Vec<_> = futures::stream::iter(objects) + .map(|object| { + fetch_schema_with_location( + store.as_ref(), + object, + self.metadata_size_hint(), + ) + }) + .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 + .buffered(state.config_options().execution.meta_fetch_concurrency) + .try_collect() + .await?; + + // Schema inference adds fields based the order they are seen + // which depends on the order the files are processed. For some + // object stores (like local file systems) the order returned from list + // is not deterministic. Thus, to ensure deterministic schema inference + // sort the files first. + // https://github.com/apache/datafusion/pull/6629 + schemas.sort_by(|(location1, _), (location2, _)| location1.cmp(location2)); + + let schemas = schemas + .into_iter() + .map(|(_, schema)| schema) + .collect::>(); + + let schema = if self.skip_metadata() { + Schema::try_merge(clear_metadata(schemas)) + } else { + Schema::try_merge(schemas) + }?; + + let schema = if self.binary_as_string() { + transform_binary_to_string(&schema) + } else { + schema + }; + + let schema = if self.force_view_types() { + transform_schema_to_view(&schema) + } else { + schema + }; + + Ok(Arc::new(schema)) + } + + async fn infer_stats( + &self, + _state: &dyn Session, + store: &Arc, + table_schema: SchemaRef, + object: &ObjectMeta, + ) -> Result { + let stats = fetch_statistics( + store.as_ref(), + table_schema, + object, + self.metadata_size_hint(), + ) + .await?; + Ok(stats) + } + + async fn create_physical_plan( + &self, + _state: &dyn Session, + conf: FileScanConfig, + filters: Option<&Arc>, + ) -> Result> { + let mut predicate = None; + let mut metadata_size_hint = None; + + // If enable pruning then combine the filters to build the predicate. + // If disable pruning then set the predicate to None, thus readers + // will not prune data based on the statistics. + if self.enable_pruning() { + if let Some(pred) = filters.cloned() { + predicate = Some(pred); + } + } + if let Some(metadata) = self.metadata_size_hint() { + metadata_size_hint = Some(metadata); + } + + let mut source = ParquetSource::new(self.options.clone()); + + if let Some(predicate) = predicate { + source = source.with_predicate(Arc::clone(&conf.file_schema), predicate); + } + if let Some(metadata_size_hint) = metadata_size_hint { + source = source.with_metadata_size_hint(metadata_size_hint) + } + Ok(conf.with_source(Arc::new(source)).build()) + } + + async fn create_writer_physical_plan( + &self, + input: Arc, + _state: &dyn Session, + conf: FileSinkConfig, + order_requirements: Option, + ) -> Result> { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for Parquet"); + } + + let sink = Arc::new(ParquetSink::new(conf, self.options.clone())); + + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + } + + fn supports_filters_pushdown( + &self, + file_schema: &Schema, + table_schema: &Schema, + filters: &[&Expr], + ) -> Result { + if !self.options().global.pushdown_filters { + return Ok(FilePushdownSupport::NoSupport); + } + + let all_supported = filters.iter().all(|filter| { + can_expr_be_pushed_down_with_schemas(filter, file_schema, table_schema) + }); + + Ok(if all_supported { + FilePushdownSupport::Supported + } else { + FilePushdownSupport::NotSupportedForFilter + }) + } + + fn file_source(&self) -> Arc { + Arc::new(ParquetSource::default()) + } +} + +/// Coerces the file schema if the table schema uses a view type. +pub fn coerce_file_schema_to_view_type( + table_schema: &Schema, + file_schema: &Schema, +) -> Option { + let mut transform = false; + let table_fields: HashMap<_, _> = table_schema + .fields + .iter() + .map(|f| { + let dt = f.data_type(); + if dt.equals_datatype(&DataType::Utf8View) + || dt.equals_datatype(&DataType::BinaryView) + { + transform = true; + } + (f.name(), dt) + }) + .collect(); + + if !transform { + return None; + } + + let transformed_fields: Vec> = file_schema + .fields + .iter() + .map( + |field| match (table_fields.get(field.name()), field.data_type()) { + (Some(DataType::Utf8View), DataType::Utf8 | DataType::LargeUtf8) => { + field_with_new_type(field, DataType::Utf8View) + } + ( + Some(DataType::BinaryView), + DataType::Binary | DataType::LargeBinary, + ) => field_with_new_type(field, DataType::BinaryView), + _ => Arc::clone(field), + }, + ) + .collect(); + + Some(Schema::new_with_metadata( + transformed_fields, + file_schema.metadata.clone(), + )) +} + +/// If the table schema uses a string type, coerce the file schema to use a string type. +/// +/// See [ParquetFormat::binary_as_string] for details +pub fn coerce_file_schema_to_string_type( + table_schema: &Schema, + file_schema: &Schema, +) -> Option { + let mut transform = false; + let table_fields: HashMap<_, _> = table_schema + .fields + .iter() + .map(|f| (f.name(), f.data_type())) + .collect(); + let transformed_fields: Vec> = file_schema + .fields + .iter() + .map( + |field| match (table_fields.get(field.name()), field.data_type()) { + // table schema uses string type, coerce the file schema to use string type + ( + Some(DataType::Utf8), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::Utf8) + } + // table schema uses large string type, coerce the file schema to use large string type + ( + Some(DataType::LargeUtf8), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::LargeUtf8) + } + // table schema uses string view type, coerce the file schema to use view type + ( + Some(DataType::Utf8View), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::Utf8View) + } + _ => Arc::clone(field), + }, + ) + .collect(); + + if !transform { + None + } else { + Some(Schema::new_with_metadata( + transformed_fields, + file_schema.metadata.clone(), + )) + } +} + +/// Create a new field with the specified data type, copying the other +/// properties from the input field +fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { + Arc::new(field.as_ref().clone().with_data_type(new_type)) +} + +/// Transform a schema to use view types for Utf8 and Binary +/// +/// See [ParquetFormat::force_view_types] for details +pub fn transform_schema_to_view(schema: &Schema) -> Schema { + let transformed_fields: Vec> = schema + .fields + .iter() + .map(|field| match field.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + field_with_new_type(field, DataType::Utf8View) + } + DataType::Binary | DataType::LargeBinary => { + field_with_new_type(field, DataType::BinaryView) + } + _ => Arc::clone(field), + }) + .collect(); + Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) +} + +/// Transform a schema so that any binary types are strings +pub fn transform_binary_to_string(schema: &Schema) -> Schema { + let transformed_fields: Vec> = schema + .fields + .iter() + .map(|field| match field.data_type() { + DataType::Binary => field_with_new_type(field, DataType::Utf8), + DataType::LargeBinary => field_with_new_type(field, DataType::LargeUtf8), + DataType::BinaryView => field_with_new_type(field, DataType::Utf8View), + _ => Arc::clone(field), + }) + .collect(); + Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) +} + +/// [`MetadataFetch`] adapter for reading bytes from an [`ObjectStore`] +struct ObjectStoreFetch<'a> { + store: &'a dyn ObjectStore, + meta: &'a ObjectMeta, +} + +impl<'a> ObjectStoreFetch<'a> { + fn new(store: &'a dyn ObjectStore, meta: &'a ObjectMeta) -> Self { + Self { store, meta } + } +} + +impl MetadataFetch for ObjectStoreFetch<'_> { + fn fetch( + &mut self, + range: Range, + ) -> BoxFuture<'_, Result> { + async { + self.store + .get_range(&self.meta.location, range) + .await + .map_err(ParquetError::from) + } + .boxed() + } +} + +/// Fetches parquet metadata from ObjectStore for given object +/// +/// This component is a subject to **change** in near future and is exposed for low level integrations +/// through [`ParquetFileReaderFactory`]. +/// +/// [`ParquetFileReaderFactory`]: crate::ParquetFileReaderFactory +pub async fn fetch_parquet_metadata( + store: &dyn ObjectStore, + meta: &ObjectMeta, + size_hint: Option, +) -> Result { + let file_size = meta.size; + let fetch = ObjectStoreFetch::new(store, meta); + + ParquetMetaDataReader::new() + .with_prefetch_hint(size_hint) + .load_and_finish(fetch, file_size) + .await + .map_err(DataFusionError::from) +} + +/// Read and parse the schema of the Parquet file at location `path` +async fn fetch_schema( + store: &dyn ObjectStore, + file: &ObjectMeta, + metadata_size_hint: Option, +) -> Result { + let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?; + let file_metadata = metadata.file_metadata(); + let schema = parquet_to_arrow_schema( + file_metadata.schema_descr(), + file_metadata.key_value_metadata(), + )?; + Ok(schema) +} + +/// Read and parse the statistics of the Parquet file at location `path` +/// +/// See [`statistics_from_parquet_meta_calc`] for more details +pub async fn fetch_statistics( + store: &dyn ObjectStore, + table_schema: SchemaRef, + file: &ObjectMeta, + metadata_size_hint: Option, +) -> Result { + let metadata = fetch_parquet_metadata(store, file, metadata_size_hint).await?; + statistics_from_parquet_meta_calc(&metadata, table_schema) +} + +/// Convert statistics in [`ParquetMetaData`] into [`Statistics`] using ['StatisticsConverter`] +/// +/// The statistics are calculated for each column in the table schema +/// using the row group statistics in the parquet metadata. +pub fn statistics_from_parquet_meta_calc( + metadata: &ParquetMetaData, + table_schema: SchemaRef, +) -> Result { + let row_groups_metadata = metadata.row_groups(); + + let mut statistics = Statistics::new_unknown(&table_schema); + let mut has_statistics = false; + let mut num_rows = 0_usize; + let mut total_byte_size = 0_usize; + for row_group_meta in row_groups_metadata { + num_rows += row_group_meta.num_rows() as usize; + total_byte_size += row_group_meta.total_byte_size() as usize; + + if !has_statistics { + row_group_meta.columns().iter().for_each(|column| { + has_statistics = column.statistics().is_some(); + }); + } + } + statistics.num_rows = Precision::Exact(num_rows); + statistics.total_byte_size = Precision::Exact(total_byte_size); + + let file_metadata = metadata.file_metadata(); + let mut file_schema = parquet_to_arrow_schema( + file_metadata.schema_descr(), + file_metadata.key_value_metadata(), + )?; + if let Some(merged) = coerce_file_schema_to_string_type(&table_schema, &file_schema) { + file_schema = merged; + } + + if let Some(merged) = coerce_file_schema_to_view_type(&table_schema, &file_schema) { + file_schema = merged; + } + + statistics.column_statistics = if has_statistics { + let (mut max_accs, mut min_accs) = create_max_min_accs(&table_schema); + let mut null_counts_array = + vec![Precision::Exact(0); table_schema.fields().len()]; + + table_schema + .fields() + .iter() + .enumerate() + .for_each(|(idx, field)| { + match StatisticsConverter::try_new( + field.name(), + &file_schema, + file_metadata.schema_descr(), + ) { + Ok(stats_converter) => { + summarize_min_max_null_counts( + &mut min_accs, + &mut max_accs, + &mut null_counts_array, + idx, + num_rows, + &stats_converter, + row_groups_metadata, + ) + .ok(); + } + Err(e) => { + debug!("Failed to create statistics converter: {}", e); + null_counts_array[idx] = Precision::Exact(num_rows); + } + } + }); + + get_col_stats( + &table_schema, + null_counts_array, + &mut max_accs, + &mut min_accs, + ) + } else { + Statistics::unknown_column(&table_schema) + }; + + Ok(statistics) +} + +fn get_col_stats( + schema: &Schema, + null_counts: Vec>, + max_values: &mut [Option], + min_values: &mut [Option], +) -> Vec { + (0..schema.fields().len()) + .map(|i| { + let max_value = match max_values.get_mut(i).unwrap() { + Some(max_value) => max_value.evaluate().ok(), + None => None, + }; + let min_value = match min_values.get_mut(i).unwrap() { + Some(min_value) => min_value.evaluate().ok(), + None => None, + }; + ColumnStatistics { + null_count: null_counts[i], + max_value: max_value.map(Precision::Exact).unwrap_or(Precision::Absent), + min_value: min_value.map(Precision::Exact).unwrap_or(Precision::Absent), + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + } + }) + .collect() +} + +fn summarize_min_max_null_counts( + min_accs: &mut [Option], + max_accs: &mut [Option], + null_counts_array: &mut [Precision], + arrow_schema_index: usize, + num_rows: usize, + stats_converter: &StatisticsConverter, + row_groups_metadata: &[RowGroupMetaData], +) -> Result<()> { + let max_values = stats_converter.row_group_maxes(row_groups_metadata)?; + let min_values = stats_converter.row_group_mins(row_groups_metadata)?; + let null_counts = stats_converter.row_group_null_counts(row_groups_metadata)?; + + if let Some(max_acc) = &mut max_accs[arrow_schema_index] { + max_acc.update_batch(&[max_values])?; + } + + if let Some(min_acc) = &mut min_accs[arrow_schema_index] { + min_acc.update_batch(&[min_values])?; + } + + null_counts_array[arrow_schema_index] = Precision::Exact(match sum(&null_counts) { + Some(null_count) => null_count as usize, + None => num_rows, + }); + + Ok(()) +} + +/// Implements [`DataSink`] for writing to a parquet file. +pub struct ParquetSink { + /// Config options for writing data + config: FileSinkConfig, + /// Underlying parquet options + parquet_options: TableParquetOptions, + /// File metadata from successfully produced parquet files. The Mutex is only used + /// to allow inserting to HashMap from behind borrowed reference in DataSink::write_all. + written: Arc>>, +} + +impl Debug for ParquetSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ParquetSink").finish() + } +} + +impl DisplayAs for ParquetSink { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ParquetSink(file_groups=",)?; + FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?; + write!(f, ")") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ParquetSink { + /// Create from config. + pub fn new(config: FileSinkConfig, parquet_options: TableParquetOptions) -> Self { + Self { + config, + parquet_options, + written: Default::default(), + } + } + + /// Retrieve the file metadata for the written files, keyed to the path + /// which may be partitioned (in the case of hive style partitioning). + pub fn written(&self) -> HashMap { + self.written.lock().clone() + } + + /// Create writer properties based upon configuration settings, + /// including partitioning and the inclusion of arrow schema metadata. + fn create_writer_props(&self) -> Result { + let schema = if self.parquet_options.global.allow_single_file_parallelism { + // If parallelizing writes, we may be also be doing hive style partitioning + // into multiple files which impacts the schema per file. + // Refer to `get_writer_schema()` + &get_writer_schema(&self.config) + } else { + self.config.output_schema() + }; + + // TODO: avoid this clone in follow up PR, where the writer properties & schema + // are calculated once on `ParquetSink::new` + let mut parquet_opts = self.parquet_options.clone(); + if !self.parquet_options.global.skip_arrow_metadata { + parquet_opts.arrow_schema(schema); + } + + Ok(WriterPropertiesBuilder::try_from(&parquet_opts)?.build()) + } + + /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore + /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized + async fn create_async_arrow_writer( + &self, + location: &Path, + object_store: Arc, + parquet_props: WriterProperties, + ) -> Result> { + let buf_writer = BufWriter::new(object_store, location.clone()); + let options = ArrowWriterOptions::new() + .with_properties(parquet_props) + .with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata); + + let writer = AsyncArrowWriter::try_new_with_options( + buf_writer, + get_writer_schema(&self.config), + options, + )?; + Ok(writer) + } + + /// Parquet options + pub fn parquet_options(&self) -> &TableParquetOptions { + &self.parquet_options + } +} + +#[async_trait] +impl FileSink for ParquetSink { + fn config(&self) -> &FileSinkConfig { + &self.config + } + + async fn spawn_writer_tasks_and_join( + &self, + context: &Arc, + demux_task: SpawnedTask>, + mut file_stream_rx: DemuxedStreamReceiver, + object_store: Arc, + ) -> Result { + let parquet_opts = &self.parquet_options; + let allow_single_file_parallelism = + parquet_opts.global.allow_single_file_parallelism; + + let mut file_write_tasks: JoinSet< + std::result::Result<(Path, FileMetaData), DataFusionError>, + > = JoinSet::new(); + + let parquet_props = self.create_writer_props()?; + let parallel_options = ParallelParquetWriterOptions { + max_parallel_row_groups: parquet_opts + .global + .maximum_parallel_row_group_writers, + max_buffered_record_batches_per_stream: parquet_opts + .global + .maximum_buffered_record_batches_per_stream, + }; + + while let Some((path, mut rx)) = file_stream_rx.recv().await { + if !allow_single_file_parallelism { + let mut writer = self + .create_async_arrow_writer( + &path, + Arc::clone(&object_store), + parquet_props.clone(), + ) + .await?; + let mut reservation = + MemoryConsumer::new(format!("ParquetSink[{}]", path)) + .register(context.memory_pool()); + file_write_tasks.spawn(async move { + while let Some(batch) = rx.recv().await { + writer.write(&batch).await?; + reservation.try_resize(writer.memory_size())?; + } + let file_metadata = writer + .close() + .await + .map_err(DataFusionError::ParquetError)?; + Ok((path, file_metadata)) + }); + } else { + let writer = create_writer( + // Parquet files as a whole are never compressed, since they + // manage compressed blocks themselves. + FileCompressionType::UNCOMPRESSED, + &path, + Arc::clone(&object_store), + ) + .await?; + let schema = get_writer_schema(&self.config); + let props = parquet_props.clone(); + let parallel_options_clone = parallel_options.clone(); + let pool = Arc::clone(context.memory_pool()); + file_write_tasks.spawn(async move { + let file_metadata = output_single_parquet_file_parallelized( + writer, + rx, + schema, + &props, + parallel_options_clone, + pool, + ) + .await?; + Ok((path, file_metadata)) + }); + } + } + + let mut row_count = 0; + while let Some(result) = file_write_tasks.join_next().await { + match result { + Ok(r) => { + let (path, file_metadata) = r?; + row_count += file_metadata.num_rows; + let mut written_files = self.written.lock(); + written_files + .try_insert(path.clone(), file_metadata) + .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; + drop(written_files); + } + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } + } + } + } + + demux_task + .join_unwind() + .await + .map_err(DataFusionError::ExecutionJoin)??; + + Ok(row_count as u64) + } +} + +#[async_trait] +impl DataSink for ParquetSink { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> &SchemaRef { + self.config.output_schema() + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + context: &Arc, + ) -> Result { + FileSink::write_all(self, data, context).await + } +} + +/// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] +/// Once the channel is exhausted, returns the ArrowColumnWriter. +async fn column_serializer_task( + mut rx: Receiver, + mut writer: ArrowColumnWriter, + mut reservation: MemoryReservation, +) -> Result<(ArrowColumnWriter, MemoryReservation)> { + while let Some(col) = rx.recv().await { + writer.write(&col)?; + reservation.try_resize(writer.memory_size())?; + } + Ok((writer, reservation)) +} + +type ColumnWriterTask = SpawnedTask>; +type ColSender = Sender; + +/// Spawns a parallel serialization task for each column +/// Returns join handles for each columns serialization task along with a send channel +/// to send arrow arrays to each serialization task. +fn spawn_column_parallel_row_group_writer( + schema: Arc, + parquet_props: Arc, + max_buffer_size: usize, + pool: &Arc, +) -> Result<(Vec, Vec)> { + let schema_desc = ArrowSchemaConverter::new().convert(&schema)?; + let col_writers = get_column_writers(&schema_desc, &parquet_props, &schema)?; + let num_columns = col_writers.len(); + + let mut col_writer_tasks = Vec::with_capacity(num_columns); + let mut col_array_channels = Vec::with_capacity(num_columns); + for writer in col_writers.into_iter() { + // Buffer size of this channel limits the number of arrays queued up for column level serialization + let (send_array, receive_array) = + mpsc::channel::(max_buffer_size); + col_array_channels.push(send_array); + + let reservation = + MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(pool); + let task = SpawnedTask::spawn(column_serializer_task( + receive_array, + writer, + reservation, + )); + col_writer_tasks.push(task); + } + + Ok((col_writer_tasks, col_array_channels)) +} + +/// Settings related to writing parquet files in parallel +#[derive(Clone)] +struct ParallelParquetWriterOptions { + max_parallel_row_groups: usize, + max_buffered_record_batches_per_stream: usize, +} + +/// This is the return type of calling [ArrowColumnWriter].close() on each column +/// i.e. the Vec of encoded columns which can be appended to a row group +type RBStreamSerializeResult = Result<(Vec, MemoryReservation, usize)>; + +/// Sends the ArrowArrays in passed [RecordBatch] through the channels to their respective +/// parallel column serializers. +async fn send_arrays_to_col_writers( + col_array_channels: &[ColSender], + rb: &RecordBatch, + schema: Arc, +) -> Result<()> { + // Each leaf column has its own channel, increment next_channel for each leaf column sent. + let mut next_channel = 0; + for (array, field) in rb.columns().iter().zip(schema.fields()) { + for c in compute_leaves(field, array)? { + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if col_array_channels[next_channel].send(c).await.is_err() { + return Ok(()); + } + + next_channel += 1; + } + } + + Ok(()) +} + +/// Spawns a tokio task which joins the parallel column writer tasks, +/// and finalizes the row group +fn spawn_rg_join_and_finalize_task( + column_writer_tasks: Vec, + rg_rows: usize, + pool: &Arc, +) -> SpawnedTask { + let mut rg_reservation = + MemoryConsumer::new("ParquetSink(SerializedRowGroupWriter)").register(pool); + + SpawnedTask::spawn(async move { + let num_cols = column_writer_tasks.len(); + let mut finalized_rg = Vec::with_capacity(num_cols); + for task in column_writer_tasks.into_iter() { + let (writer, _col_reservation) = task + .join_unwind() + .await + .map_err(DataFusionError::ExecutionJoin)??; + let encoded_size = writer.get_estimated_total_bytes(); + rg_reservation.grow(encoded_size); + finalized_rg.push(writer.close()?); + } + + Ok((finalized_rg, rg_reservation, rg_rows)) + }) +} + +/// This task coordinates the serialization of a parquet file in parallel. +/// As the query produces RecordBatches, these are written to a RowGroup +/// via parallel [ArrowColumnWriter] tasks. Once the desired max rows per +/// row group is reached, the parallel tasks are joined on another separate task +/// and sent to a concatenation task. This task immediately continues to work +/// on the next row group in parallel. So, parquet serialization is parallelized +/// across both columns and row_groups, with a theoretical max number of parallel tasks +/// given by n_columns * num_row_groups. +fn spawn_parquet_parallel_serialization_task( + mut data: Receiver, + serialize_tx: Sender>, + schema: Arc, + writer_props: Arc, + parallel_options: ParallelParquetWriterOptions, + pool: Arc, +) -> SpawnedTask> { + SpawnedTask::spawn(async move { + let max_buffer_rb = parallel_options.max_buffered_record_batches_per_stream; + let max_row_group_rows = writer_props.max_row_group_size(); + let (mut column_writer_handles, mut col_array_channels) = + spawn_column_parallel_row_group_writer( + Arc::clone(&schema), + Arc::clone(&writer_props), + max_buffer_rb, + &pool, + )?; + let mut current_rg_rows = 0; + + while let Some(mut rb) = data.recv().await { + // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case + // when max_row_group_rows < execution.batch_size as an alternative to a recursive async + // function. + loop { + if current_rg_rows + rb.num_rows() < max_row_group_rows { + send_arrays_to_col_writers( + &col_array_channels, + &rb, + Arc::clone(&schema), + ) + .await?; + current_rg_rows += rb.num_rows(); + break; + } else { + let rows_left = max_row_group_rows - current_rg_rows; + let a = rb.slice(0, rows_left); + send_arrays_to_col_writers( + &col_array_channels, + &a, + Arc::clone(&schema), + ) + .await?; + + // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup + // on a separate task, so that we can immediately start on the next RG before waiting + // for the current one to finish. + drop(col_array_channels); + let finalize_rg_task = spawn_rg_join_and_finalize_task( + column_writer_handles, + max_row_group_rows, + &pool, + ); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + + current_rg_rows = 0; + rb = rb.slice(rows_left, rb.num_rows() - rows_left); + + (column_writer_handles, col_array_channels) = + spawn_column_parallel_row_group_writer( + Arc::clone(&schema), + Arc::clone(&writer_props), + max_buffer_rb, + &pool, + )?; + } + } + } + + drop(col_array_channels); + // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows + if current_rg_rows > 0 { + let finalize_rg_task = spawn_rg_join_and_finalize_task( + column_writer_handles, + current_rg_rows, + &pool, + ); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + } + + Ok(()) + }) +} + +/// Consume RowGroups serialized by other parallel tasks and concatenate them in +/// to the final parquet file, while flushing finalized bytes to an [ObjectStore] +async fn concatenate_parallel_row_groups( + mut serialize_rx: Receiver>, + schema: Arc, + writer_props: Arc, + mut object_store_writer: Box, + pool: Arc, +) -> Result { + let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES); + + let mut file_reservation = + MemoryConsumer::new("ParquetSink(SerializedFileWriter)").register(&pool); + + let schema_desc = ArrowSchemaConverter::new().convert(schema.as_ref())?; + let mut parquet_writer = SerializedFileWriter::new( + merged_buff.clone(), + schema_desc.root_schema_ptr(), + writer_props, + )?; + + while let Some(task) = serialize_rx.recv().await { + let result = task.join_unwind().await; + let mut rg_out = parquet_writer.next_row_group()?; + let (serialized_columns, mut rg_reservation, _cnt) = + result.map_err(DataFusionError::ExecutionJoin)??; + for chunk in serialized_columns { + chunk.append_to_row_group(&mut rg_out)?; + rg_reservation.free(); + + let mut buff_to_flush = merged_buff.buffer.try_lock().unwrap(); + file_reservation.try_resize(buff_to_flush.len())?; + + if buff_to_flush.len() > BUFFER_FLUSH_BYTES { + object_store_writer + .write_all(buff_to_flush.as_slice()) + .await?; + buff_to_flush.clear(); + file_reservation.try_resize(buff_to_flush.len())?; // will set to zero + } + } + rg_out.close()?; + } + + let file_metadata = parquet_writer.close()?; + let final_buff = merged_buff.buffer.try_lock().unwrap(); + + object_store_writer.write_all(final_buff.as_slice()).await?; + object_store_writer.shutdown().await?; + file_reservation.free(); + + Ok(file_metadata) +} + +/// Parallelizes the serialization of a single parquet file, by first serializing N +/// independent RecordBatch streams in parallel to RowGroups in memory. Another +/// task then stitches these independent RowGroups together and streams this large +/// single parquet file to an ObjectStore in multiple parts. +async fn output_single_parquet_file_parallelized( + object_store_writer: Box, + data: Receiver, + output_schema: Arc, + parquet_props: &WriterProperties, + parallel_options: ParallelParquetWriterOptions, + pool: Arc, +) -> Result { + let max_rowgroups = parallel_options.max_parallel_row_groups; + // Buffer size of this channel limits maximum number of RowGroups being worked on in parallel + let (serialize_tx, serialize_rx) = + mpsc::channel::>(max_rowgroups); + + let arc_props = Arc::new(parquet_props.clone()); + let launch_serialization_task = spawn_parquet_parallel_serialization_task( + data, + serialize_tx, + Arc::clone(&output_schema), + Arc::clone(&arc_props), + parallel_options, + Arc::clone(&pool), + ); + let file_metadata = concatenate_parallel_row_groups( + serialize_rx, + Arc::clone(&output_schema), + Arc::clone(&arc_props), + object_store_writer, + pool, + ) + .await?; + + launch_serialization_task + .join_unwind() + .await + .map_err(DataFusionError::ExecutionJoin)??; + Ok(file_metadata) +} + +/// Min/max aggregation can take Dictionary encode input but always produces unpacked +/// (aka non Dictionary) output. We need to adjust the output data type to reflect this. +/// The reason min/max aggregate produces unpacked output because there is only one +/// min/max value per group; there is no needs to keep them Dictionary encode +fn min_max_aggregate_data_type(input_type: &DataType) -> &DataType { + if let DataType::Dictionary(_, value_type) = input_type { + value_type.as_ref() + } else { + input_type + } +} + +fn create_max_min_accs( + schema: &Schema, +) -> (Vec>, Vec>) { + let max_values: Vec> = schema + .fields() + .iter() + .map(|field| { + MaxAccumulator::try_new(min_max_aggregate_data_type(field.data_type())).ok() + }) + .collect(); + let min_values: Vec> = schema + .fields() + .iter() + .map(|field| { + MinAccumulator::try_new(min_max_aggregate_data_type(field.data_type())).ok() + }) + .collect(); + (max_values, min_values) +} diff --git a/datafusion/core/src/datasource/physical_plan/parquet/metrics.rs b/datafusion/datasource-parquet/src/metrics.rs similarity index 99% rename from datafusion/core/src/datasource/physical_plan/parquet/metrics.rs rename to datafusion/datasource-parquet/src/metrics.rs index f1b5f71530dc..3213d0201295 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/metrics.rs +++ b/datafusion/datasource-parquet/src/metrics.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::physical_plan::metrics::{ +use datafusion_physical_plan::metrics::{ Count, ExecutionPlanMetricsSet, MetricBuilder, Time, }; diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs new file mode 100644 index 000000000000..fb1f2d55169f --- /dev/null +++ b/datafusion/datasource-parquet/src/mod.rs @@ -0,0 +1,547 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ParquetExec`] FileSource for reading Parquet files + +pub mod access_plan; +pub mod file_format; +mod metrics; +mod opener; +mod page_filter; +mod reader; +mod row_filter; +mod row_group_filter; +pub mod source; +mod writer; + +use std::any::Any; +use std::fmt::Formatter; +use std::sync::Arc; + +pub use access_plan::{ParquetAccessPlan, RowGroupAccess}; +use arrow::datatypes::SchemaRef; +use datafusion_common::config::{ConfigOptions, TableParquetOptions}; +use datafusion_common::Result; +use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::schema_adapter::SchemaAdapterFactory; +use datafusion_datasource::source::DataSourceExec; +use datafusion_datasource::PartitionedFile; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, PhysicalExpr, +}; +use datafusion_physical_optimizer::pruning::PruningPredicate; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::metrics::MetricsSet; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, +}; +pub use file_format::*; +pub use metrics::ParquetFileMetrics; +pub use page_filter::PagePruningAccessPlanFilter; +pub use reader::{DefaultParquetFileReaderFactory, ParquetFileReaderFactory}; +pub use row_filter::build_row_filter; +pub use row_filter::can_expr_be_pushed_down_with_schemas; +pub use row_group_filter::RowGroupAccessPlanFilter; +use source::ParquetSource; +pub use writer::plan_to_parquet; + +use log::debug; + +#[derive(Debug, Clone)] +#[deprecated(since = "46.0.0", note = "use DataSourceExec instead")] +/// Deprecated Execution plan replaced with DataSourceExec +pub struct ParquetExec { + inner: DataSourceExec, + base_config: FileScanConfig, + table_parquet_options: TableParquetOptions, + /// Optional predicate for row filtering during parquet scan + predicate: Option>, + /// Optional predicate for pruning row groups (derived from `predicate`) + pruning_predicate: Option>, + /// Optional user defined parquet file reader factory + parquet_file_reader_factory: Option>, + /// Optional user defined schema adapter + schema_adapter_factory: Option>, +} + +#[allow(unused, deprecated)] +impl From for ParquetExecBuilder { + fn from(exec: ParquetExec) -> Self { + exec.into_builder() + } +} + +/// [`ParquetExecBuilder`], deprecated builder for [`ParquetExec`]. +/// +/// ParquetExec is replaced with `DataSourceExec` and it includes `ParquetSource` +/// +/// See example on [`ParquetSource`]. +#[deprecated( + since = "46.0.0", + note = "use DataSourceExec with ParquetSource instead" +)] +#[allow(unused, deprecated)] +pub struct ParquetExecBuilder { + file_scan_config: FileScanConfig, + predicate: Option>, + metadata_size_hint: Option, + table_parquet_options: TableParquetOptions, + parquet_file_reader_factory: Option>, + schema_adapter_factory: Option>, +} + +#[allow(unused, deprecated)] +impl ParquetExecBuilder { + /// Create a new builder to read the provided file scan configuration + pub fn new(file_scan_config: FileScanConfig) -> Self { + Self::new_with_options(file_scan_config, TableParquetOptions::default()) + } + + /// Create a new builder to read the data specified in the file scan + /// configuration with the provided `TableParquetOptions`. + pub fn new_with_options( + file_scan_config: FileScanConfig, + table_parquet_options: TableParquetOptions, + ) -> Self { + Self { + file_scan_config, + predicate: None, + metadata_size_hint: None, + table_parquet_options, + parquet_file_reader_factory: None, + schema_adapter_factory: None, + } + } + + /// Update the list of files groups to read + pub fn with_file_groups(mut self, file_groups: Vec>) -> Self { + self.file_scan_config.file_groups = file_groups; + self + } + + /// Set the filter predicate when reading. + /// + /// See the "Predicate Pushdown" section of the [`ParquetExec`] documentation + /// for more details. + pub fn with_predicate(mut self, predicate: Arc) -> Self { + self.predicate = Some(predicate); + self + } + + /// Set the metadata size hint + /// + /// This value determines how many bytes at the end of the file the default + /// [`ParquetFileReaderFactory`] will request in the initial IO. If this is + /// too small, the ParquetExec will need to make additional IO requests to + /// read the footer. + pub fn with_metadata_size_hint(mut self, metadata_size_hint: usize) -> Self { + self.metadata_size_hint = Some(metadata_size_hint); + self + } + + /// Set the options for controlling how the ParquetExec reads parquet files. + /// + /// See also [`Self::new_with_options`] + pub fn with_table_parquet_options( + mut self, + table_parquet_options: TableParquetOptions, + ) -> Self { + self.table_parquet_options = table_parquet_options; + self + } + + /// Set optional user defined parquet file reader factory. + /// + /// You can use [`ParquetFileReaderFactory`] to more precisely control how + /// data is read from parquet files (e.g. skip re-reading metadata, coalesce + /// I/O operations, etc). + /// + /// The default reader factory reads directly from an [`ObjectStore`] + /// instance using individual I/O operations for the footer and each page. + /// + /// If a custom `ParquetFileReaderFactory` is provided, then data access + /// operations will be routed to this factory instead of [`ObjectStore`]. + /// + /// [`ObjectStore`]: object_store::ObjectStore + pub fn with_parquet_file_reader_factory( + mut self, + parquet_file_reader_factory: Arc, + ) -> Self { + self.parquet_file_reader_factory = Some(parquet_file_reader_factory); + self + } + + /// Set optional schema adapter factory. + /// + /// [`SchemaAdapterFactory`] allows user to specify how fields from the + /// parquet file get mapped to that of the table schema. The default schema + /// adapter uses arrow's cast library to map the parquet fields to the table + /// schema. + pub fn with_schema_adapter_factory( + mut self, + schema_adapter_factory: Arc, + ) -> Self { + self.schema_adapter_factory = Some(schema_adapter_factory); + self + } + + /// Convenience: build an `Arc`d `ParquetExec` from this builder + pub fn build_arc(self) -> Arc { + Arc::new(self.build()) + } + + /// Build a [`ParquetExec`] + #[must_use] + pub fn build(self) -> ParquetExec { + let Self { + file_scan_config, + predicate, + metadata_size_hint, + table_parquet_options, + parquet_file_reader_factory, + schema_adapter_factory, + } = self; + let mut parquet = ParquetSource::new(table_parquet_options); + if let Some(predicate) = predicate.clone() { + parquet = parquet + .with_predicate(Arc::clone(&file_scan_config.file_schema), predicate); + } + if let Some(metadata_size_hint) = metadata_size_hint { + parquet = parquet.with_metadata_size_hint(metadata_size_hint) + } + if let Some(parquet_reader_factory) = parquet_file_reader_factory { + parquet = parquet.with_parquet_file_reader_factory(parquet_reader_factory) + } + if let Some(schema_factory) = schema_adapter_factory { + parquet = parquet.with_schema_adapter_factory(schema_factory); + } + + let base_config = file_scan_config.with_source(Arc::new(parquet.clone())); + debug!("Creating ParquetExec, files: {:?}, projection {:?}, predicate: {:?}, limit: {:?}", + base_config.file_groups, base_config.projection, predicate, base_config.limit); + + ParquetExec { + inner: DataSourceExec::new(Arc::new(base_config.clone())), + base_config, + predicate, + pruning_predicate: parquet.pruning_predicate, + schema_adapter_factory: parquet.schema_adapter_factory, + parquet_file_reader_factory: parquet.parquet_file_reader_factory, + table_parquet_options: parquet.table_parquet_options, + } + } +} + +#[allow(unused, deprecated)] +impl ParquetExec { + /// Create a new Parquet reader execution plan provided file list and schema. + pub fn new( + base_config: FileScanConfig, + predicate: Option>, + metadata_size_hint: Option, + table_parquet_options: TableParquetOptions, + ) -> Self { + let mut builder = + ParquetExecBuilder::new_with_options(base_config, table_parquet_options); + if let Some(predicate) = predicate { + builder = builder.with_predicate(predicate); + } + if let Some(metadata_size_hint) = metadata_size_hint { + builder = builder.with_metadata_size_hint(metadata_size_hint); + } + builder.build() + } + /// Return a [`ParquetExecBuilder`]. + /// + /// See example on [`ParquetExec`] and [`ParquetExecBuilder`] for specifying + /// parquet table options. + pub fn builder(file_scan_config: FileScanConfig) -> ParquetExecBuilder { + ParquetExecBuilder::new(file_scan_config) + } + + /// Convert this `ParquetExec` into a builder for modification + pub fn into_builder(self) -> ParquetExecBuilder { + // list out fields so it is clear what is being dropped + // (note the fields which are dropped are re-created as part of calling + // `build` on the builder) + let file_scan_config = self.file_scan_config(); + let parquet = self.parquet_source(); + + ParquetExecBuilder { + file_scan_config, + predicate: parquet.predicate, + metadata_size_hint: parquet.metadata_size_hint, + table_parquet_options: parquet.table_parquet_options, + parquet_file_reader_factory: parquet.parquet_file_reader_factory, + schema_adapter_factory: parquet.schema_adapter_factory, + } + } + fn file_scan_config(&self) -> FileScanConfig { + self.inner + .data_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + fn parquet_source(&self) -> ParquetSource { + self.file_scan_config() + .file_source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + /// [`FileScanConfig`] that controls this scan (such as which files to read) + pub fn base_config(&self) -> &FileScanConfig { + &self.base_config + } + /// Options passed to the parquet reader for this scan + pub fn table_parquet_options(&self) -> &TableParquetOptions { + &self.table_parquet_options + } + /// Optional predicate. + pub fn predicate(&self) -> Option<&Arc> { + self.predicate.as_ref() + } + /// Optional reference to this parquet scan's pruning predicate + pub fn pruning_predicate(&self) -> Option<&Arc> { + self.pruning_predicate.as_ref() + } + /// return the optional file reader factory + pub fn parquet_file_reader_factory( + &self, + ) -> Option<&Arc> { + self.parquet_file_reader_factory.as_ref() + } + /// Optional user defined parquet file reader factory. + pub fn with_parquet_file_reader_factory( + mut self, + parquet_file_reader_factory: Arc, + ) -> Self { + let mut parquet = self.parquet_source(); + parquet.parquet_file_reader_factory = + Some(Arc::clone(&parquet_file_reader_factory)); + let file_source = self.file_scan_config(); + self.inner = self + .inner + .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); + self.parquet_file_reader_factory = Some(parquet_file_reader_factory); + self + } + /// return the optional schema adapter factory + pub fn schema_adapter_factory(&self) -> Option<&Arc> { + self.schema_adapter_factory.as_ref() + } + /// Set optional schema adapter factory. + /// + /// [`SchemaAdapterFactory`] allows user to specify how fields from the + /// parquet file get mapped to that of the table schema. The default schema + /// adapter uses arrow's cast library to map the parquet fields to the table + /// schema. + pub fn with_schema_adapter_factory( + mut self, + schema_adapter_factory: Arc, + ) -> Self { + let mut parquet = self.parquet_source(); + parquet.schema_adapter_factory = Some(Arc::clone(&schema_adapter_factory)); + let file_source = self.file_scan_config(); + self.inner = self + .inner + .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); + self.schema_adapter_factory = Some(schema_adapter_factory); + self + } + /// If true, the predicate will be used during the parquet scan. + /// Defaults to false + /// + /// [`Expr`]: datafusion_expr::Expr + pub fn with_pushdown_filters(mut self, pushdown_filters: bool) -> Self { + let mut parquet = self.parquet_source(); + parquet.table_parquet_options.global.pushdown_filters = pushdown_filters; + let file_source = self.file_scan_config(); + self.inner = self + .inner + .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); + self.table_parquet_options.global.pushdown_filters = pushdown_filters; + self + } + + /// Return the value described in [`Self::with_pushdown_filters`] + fn pushdown_filters(&self) -> bool { + self.parquet_source() + .table_parquet_options + .global + .pushdown_filters + } + /// If true, the `RowFilter` made by `pushdown_filters` may try to + /// minimize the cost of filter evaluation by reordering the + /// predicate [`Expr`]s. If false, the predicates are applied in + /// the same order as specified in the query. Defaults to false. + /// + /// [`Expr`]: datafusion_expr::Expr + pub fn with_reorder_filters(mut self, reorder_filters: bool) -> Self { + let mut parquet = self.parquet_source(); + parquet.table_parquet_options.global.reorder_filters = reorder_filters; + let file_source = self.file_scan_config(); + self.inner = self + .inner + .with_data_source(Arc::new(file_source.with_source(Arc::new(parquet)))); + self.table_parquet_options.global.reorder_filters = reorder_filters; + self + } + /// Return the value described in [`Self::with_reorder_filters`] + fn reorder_filters(&self) -> bool { + self.parquet_source() + .table_parquet_options + .global + .reorder_filters + } + /// If enabled, the reader will read the page index + /// This is used to optimize filter pushdown + /// via `RowSelector` and `RowFilter` by + /// eliminating unnecessary IO and decoding + fn bloom_filter_on_read(&self) -> bool { + self.parquet_source() + .table_parquet_options + .global + .bloom_filter_on_read + } + /// Return the value described in [`ParquetSource::with_enable_page_index`] + fn enable_page_index(&self) -> bool { + self.parquet_source() + .table_parquet_options + .global + .enable_page_index + } + + fn output_partitioning_helper(file_config: &FileScanConfig) -> Partitioning { + Partitioning::UnknownPartitioning(file_config.file_groups.len()) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: &[LexOrdering], + constraints: Constraints, + file_config: &FileScanConfig, + ) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, orderings) + .with_constraints(constraints), + Self::output_partitioning_helper(file_config), // Output Partitioning + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + + /// Updates the file groups to read and recalculates the output partitioning + /// + /// Note this function does not update statistics or other properties + /// that depend on the file groups. + fn with_file_groups_and_update_partitioning( + mut self, + file_groups: Vec>, + ) -> Self { + let mut config = self.file_scan_config(); + config.file_groups = file_groups; + self.inner = self.inner.with_data_source(Arc::new(config)); + self + } +} + +#[allow(unused, deprecated)] +impl DisplayAs for ParquetExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(unused, deprecated)] +impl ExecutionPlan for ParquetExec { + fn name(&self) -> &'static str { + "ParquetExec" + } + + /// Return a reference to Any that can be used for downcasting + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &PlanProperties { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + // this is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + /// Redistribute files across partitions according to their size + /// See comments on `FileGroupPartitioner` for more detail. + fn repartitioned( + &self, + target_partitions: usize, + config: &ConfigOptions, + ) -> Result>> { + self.inner.repartitioned(target_partitions, config) + } + + fn execute( + &self, + partition_index: usize, + ctx: Arc, + ) -> Result { + self.inner.execute(partition_index, ctx) + } + fn metrics(&self) -> Option { + self.inner.metrics() + } + fn statistics(&self) -> Result { + self.inner.statistics() + } + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn with_fetch(&self, limit: Option) -> Option> { + self.inner.with_fetch(limit) + } +} + +fn should_enable_page_index( + enable_page_index: bool, + page_pruning_predicate: &Option>, +) -> bool { + enable_page_index + && page_pruning_predicate.is_some() + && page_pruning_predicate + .as_ref() + .map(|p| p.filter_number() > 0) + .unwrap_or(false) +} diff --git a/datafusion/core/src/datasource/physical_plan/parquet/opener.rs b/datafusion/datasource-parquet/src/opener.rs similarity index 95% rename from datafusion/core/src/datasource/physical_plan/parquet/opener.rs rename to datafusion/datasource-parquet/src/opener.rs index 4230a1bdce38..3c623f558e43 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/opener.rs +++ b/datafusion/datasource-parquet/src/opener.rs @@ -19,18 +19,18 @@ use std::sync::Arc; -use crate::datasource::file_format::parquet::{ +use crate::file_format::{ coerce_file_schema_to_string_type, coerce_file_schema_to_view_type, }; -use crate::datasource::physical_plan::parquet::page_filter::PagePruningAccessPlanFilter; -use crate::datasource::physical_plan::parquet::row_group_filter::RowGroupAccessPlanFilter; -use crate::datasource::physical_plan::parquet::{ - row_filter, should_enable_page_index, ParquetAccessPlan, +use crate::page_filter::PagePruningAccessPlanFilter; +use crate::row_group_filter::RowGroupAccessPlanFilter; +use crate::{ + row_filter, should_enable_page_index, ParquetAccessPlan, ParquetFileMetrics, + ParquetFileReaderFactory, }; -use crate::datasource::physical_plan::{ - FileMeta, FileOpenFuture, FileOpener, ParquetFileMetrics, ParquetFileReaderFactory, -}; -use crate::datasource::schema_adapter::SchemaAdapterFactory; +use datafusion_datasource::file_meta::FileMeta; +use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; +use datafusion_datasource::schema_adapter::SchemaAdapterFactory; use arrow::datatypes::SchemaRef; use arrow::error::ArrowError; diff --git a/datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs similarity index 99% rename from datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs rename to datafusion/datasource-parquet/src/page_filter.rs index 02329effb09a..ef832d808647 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -21,7 +21,7 @@ use std::collections::HashSet; use std::sync::Arc; use super::metrics::ParquetFileMetrics; -use crate::datasource::physical_plan::parquet::ParquetAccessPlan; +use crate::ParquetAccessPlan; use arrow::array::BooleanArray; use arrow::{ diff --git a/datafusion/core/src/datasource/physical_plan/parquet/reader.rs b/datafusion/datasource-parquet/src/reader.rs similarity index 98% rename from datafusion/core/src/datasource/physical_plan/parquet/reader.rs rename to datafusion/datasource-parquet/src/reader.rs index 8a4ba136fc96..5924a5b5038f 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/reader.rs +++ b/datafusion/datasource-parquet/src/reader.rs @@ -18,8 +18,8 @@ //! [`ParquetFileReaderFactory`] and [`DefaultParquetFileReaderFactory`] for //! low level control of parquet file readers -use crate::datasource::physical_plan::{FileMeta, ParquetFileMetrics}; use bytes::Bytes; +use datafusion_datasource::file_meta::FileMeta; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use futures::future::BoxFuture; use object_store::ObjectStore; @@ -29,6 +29,8 @@ use std::fmt::Debug; use std::ops::Range; use std::sync::Arc; +use crate::ParquetFileMetrics; + /// Interface for reading parquet files. /// /// The combined implementations of [`ParquetFileReaderFactory`] and diff --git a/datafusion/core/src/datasource/physical_plan/parquet/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs similarity index 98% rename from datafusion/core/src/datasource/physical_plan/parquet/row_filter.rs rename to datafusion/datasource-parquet/src/row_filter.rs index ac6eaf2c8f63..39fcecf37c6d 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/row_filter.rs +++ b/datafusion/datasource-parquet/src/row_filter.rs @@ -71,17 +71,17 @@ use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter}; use parquet::arrow::ProjectionMask; use parquet::file::metadata::ParquetMetaData; -use crate::datasource::schema_adapter::SchemaMapper; use datafusion_common::cast::as_boolean_array; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; use datafusion_common::{arrow_datafusion_err, DataFusionError, Result, ScalarValue}; +use datafusion_datasource::schema_adapter::SchemaMapper; use datafusion_physical_expr::expressions::{Column, Literal}; use datafusion_physical_expr::utils::reassign_predicate_columns; use datafusion_physical_expr::{split_conjunction, PhysicalExpr}; -use crate::physical_plan::metrics; +use datafusion_physical_plan::metrics; use super::ParquetFileMetrics; @@ -584,7 +584,7 @@ pub fn build_row_filter( #[cfg(test)] mod test { use super::*; - use crate::datasource::schema_adapter::{ + use datafusion_datasource::schema_adapter::{ DefaultSchemaAdapterFactory, SchemaAdapterFactory, }; @@ -601,7 +601,7 @@ mod test { // We should ignore predicate that read non-primitive columns #[test] fn test_filter_candidate_builder_ignore_complex_types() { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let file = std::fs::File::open(format!("{testdata}/list_columns.parquet")) .expect("opening file"); @@ -626,7 +626,7 @@ mod test { // If a column exists in the table schema but not the file schema it should be rewritten to a null expression #[test] fn test_filter_candidate_builder_rewrite_missing_column() { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let file = std::fs::File::open(format!("{testdata}/alltypes_plain.parquet")) .expect("opening file"); @@ -665,7 +665,7 @@ mod test { #[test] fn test_filter_type_coercion() { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let file = std::fs::File::open(format!("{testdata}/alltypes_plain.parquet")) .expect("opening file"); @@ -835,7 +835,7 @@ mod test { } fn get_basic_table_schema() -> Schema { - let testdata = crate::test_util::parquet_test_data(); + let testdata = datafusion_common::test_util::parquet_test_data(); let file = std::fs::File::open(format!("{testdata}/alltypes_plain.parquet")) .expect("opening file"); diff --git a/datafusion/core/src/datasource/physical_plan/parquet/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs similarity index 99% rename from datafusion/core/src/datasource/physical_plan/parquet/row_group_filter.rs rename to datafusion/datasource-parquet/src/row_group_filter.rs index 27bfb26902e5..9d5f9fa16b6e 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -19,13 +19,11 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use super::{ParquetAccessPlan, ParquetFileMetrics}; -use crate::datasource::listing::FileRange; - use arrow::array::{ArrayRef, BooleanArray}; use arrow::datatypes::Schema; use datafusion_common::{Column, Result, ScalarValue}; +use datafusion_datasource::FileRange; use datafusion_physical_optimizer::pruning::{PruningPredicate, PruningStatistics}; - use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::arrow::parquet_column; use parquet::basic::Type; @@ -435,15 +433,14 @@ mod tests { use std::sync::Arc; use super::*; - use crate::datasource::physical_plan::parquet::reader::ParquetFileReader; - use crate::physical_plan::metrics::ExecutionPlanMetricsSet; + use crate::reader::ParquetFileReader; use arrow::datatypes::DataType::Decimal128; use arrow::datatypes::{DataType, Field}; use datafusion_common::Result; use datafusion_expr::{cast, col, lit, Expr}; use datafusion_physical_expr::planner::logical2physical; - + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use parquet::arrow::async_reader::ParquetObjectReader; use parquet::arrow::ArrowSchemaConverter; use parquet::basic::LogicalType; diff --git a/datafusion/core/src/datasource/physical_plan/parquet/source.rs b/datafusion/datasource-parquet/src/source.rs similarity index 94% rename from datafusion/core/src/datasource/physical_plan/parquet/source.rs rename to datafusion/datasource-parquet/src/source.rs index 142725524f1b..683d62a1df49 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -20,11 +20,12 @@ use std::any::Any; use std::fmt::Formatter; use std::sync::Arc; -use crate::datasource::physical_plan::parquet::opener::ParquetOpener; -use crate::datasource::physical_plan::parquet::page_filter::PagePruningAccessPlanFilter; -use crate::datasource::physical_plan::parquet::DefaultParquetFileReaderFactory; -use crate::datasource::physical_plan::{FileOpener, ParquetFileReaderFactory}; -use crate::datasource::schema_adapter::{ +use crate::opener::ParquetOpener; +use crate::page_filter::PagePruningAccessPlanFilter; +use crate::DefaultParquetFileReaderFactory; +use crate::ParquetFileReaderFactory; +use datafusion_datasource::file_stream::FileOpener; +use datafusion_datasource::schema_adapter::{ DefaultSchemaAdapterFactory, SchemaAdapterFactory, }; @@ -75,12 +76,12 @@ use object_store::ObjectStore; /// ``` /// # use std::sync::Arc; /// # use arrow::datatypes::Schema; -/// # use datafusion::datasource::physical_plan::FileScanConfig; -/// # use datafusion::datasource::physical_plan::parquet::source::ParquetSource; -/// # use datafusion::datasource::listing::PartitionedFile; +/// # use datafusion_datasource::file_scan_config::FileScanConfig; +/// # use datafusion_datasource_parquet::source::ParquetSource; +/// # use datafusion_datasource::PartitionedFile; /// # use datafusion_execution::object_store::ObjectStoreUrl; /// # use datafusion_physical_expr::expressions::lit; -/// # use datafusion::datasource::source::DataSourceExec; +/// # use datafusion_datasource::source::DataSourceExec; /// # use datafusion_common::config::TableParquetOptions; /// /// # let file_schema = Arc::new(Schema::empty()); @@ -157,9 +158,9 @@ use object_store::ObjectStore; /// ```no_run /// # use std::sync::Arc; /// # use arrow::datatypes::Schema; -/// # use datafusion::datasource::physical_plan::FileScanConfig; -/// # use datafusion::datasource::listing::PartitionedFile; -/// # use datafusion::datasource::source::DataSourceExec; +/// # use datafusion_datasource::file_scan_config::FileScanConfig; +/// # use datafusion_datasource::PartitionedFile; +/// # use datafusion_datasource::source::DataSourceExec; /// /// # fn parquet_exec() -> DataSourceExec { unimplemented!() } /// // Split a single DataSourceExec into multiple DataSourceExecs, one for each file @@ -196,12 +197,12 @@ use object_store::ObjectStore; /// ``` /// # use std::sync::Arc; /// # use arrow::datatypes::{Schema, SchemaRef}; -/// # use datafusion::datasource::listing::PartitionedFile; -/// # use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; -/// # use datafusion::datasource::physical_plan::FileScanConfig; -/// # use datafusion::datasource::physical_plan::parquet::source::ParquetSource; +/// # use datafusion_datasource::PartitionedFile; +/// # use datafusion_datasource_parquet::ParquetAccessPlan; +/// # use datafusion_datasource::file_scan_config::FileScanConfig; +/// # use datafusion_datasource_parquet::source::ParquetSource; /// # use datafusion_execution::object_store::ObjectStoreUrl; -/// # use datafusion::datasource::source::DataSourceExec; +/// # use datafusion_datasource::source::DataSourceExec; /// /// # fn schema() -> SchemaRef { /// # Arc::new(Schema::empty()) @@ -247,7 +248,7 @@ use object_store::ObjectStore; /// filled with nulls, but this can be customized via [`SchemaAdapterFactory`]. /// /// [`RecordBatch`]: arrow::record_batch::RecordBatch -/// [`SchemaAdapter`]: crate::datasource::schema_adapter::SchemaAdapter +/// [`SchemaAdapter`]: datafusion_datasource::schema_adapter::SchemaAdapter /// [`ParquetMetadata`]: parquet::file::metadata::ParquetMetaData #[derive(Clone, Default, Debug)] pub struct ParquetSource { @@ -558,7 +559,6 @@ impl FileSource for ParquetSource { .predicate() .map(|p| format!(", predicate={p}")) .unwrap_or_default(); - let pruning_predicate_string = self .pruning_predicate() .map(|pre| { @@ -578,6 +578,12 @@ impl FileSource for ParquetSource { write!(f, "{}{}", predicate_string, pruning_predicate_string) } + DisplayFormatType::TreeRender => { + if let Some(predicate) = self.predicate() { + writeln!(f, "predicate={predicate}")?; + } + Ok(()) + } } } } diff --git a/datafusion/core/src/datasource/physical_plan/parquet/writer.rs b/datafusion/datasource-parquet/src/writer.rs similarity index 98% rename from datafusion/core/src/datasource/physical_plan/parquet/writer.rs rename to datafusion/datasource-parquet/src/writer.rs index 00926dc2330b..cfdb057a4bc4 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/writer.rs +++ b/datafusion/datasource-parquet/src/writer.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::datasource::listing::ListingTableUrl; use datafusion_common::DataFusionError; +use datafusion_datasource::ListingTableUrl; use datafusion_execution::TaskContext; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use futures::StreamExt; diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index d75b27c2f685..473800c7779f 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -31,14 +31,13 @@ version.workspace = true all-features = true [features] -# Temporary feature while I move things around -avro = [] +parquet = ["dep:parquet", "tempfile"] compression = ["async-compression", "xz2", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] [dependencies] arrow = { workspace = true } -async-compression = { version = "0.4.0", features = [ +async-compression = { version = "0.4.19", features = [ "bzip2", "gzip", "xz", @@ -63,7 +62,9 @@ glob = "0.3.0" itertools = { workspace = true } log = { workspace = true } object_store = { workspace = true } +parquet = { workspace = true, optional = true } rand = { workspace = true } +tempfile = { workspace = true, optional = true } tokio = { workspace = true } tokio-util = { version = "0.7.4", features = ["io"], optional = true } url = { workspace = true } diff --git a/datafusion/datasource/src/decoder.rs b/datafusion/datasource/src/decoder.rs new file mode 100644 index 000000000000..654569f74113 --- /dev/null +++ b/datafusion/datasource/src/decoder.rs @@ -0,0 +1,191 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Module containing helper methods for the various file formats +//! See write.rs for write related helper methods + +use ::arrow::array::RecordBatch; + +use arrow::error::ArrowError; +use bytes::Buf; +use bytes::Bytes; +use datafusion_common::Result; +use futures::stream::BoxStream; +use futures::StreamExt as _; +use futures::{ready, Stream}; +use std::collections::VecDeque; +use std::fmt; +use std::task::Poll; + +/// Possible outputs of a [`BatchDeserializer`]. +#[derive(Debug, PartialEq)] +pub enum DeserializerOutput { + /// A successfully deserialized [`RecordBatch`]. + RecordBatch(RecordBatch), + /// The deserializer requires more data to make progress. + RequiresMoreData, + /// The input data has been exhausted. + InputExhausted, +} + +/// Trait defining a scheme for deserializing byte streams into structured data. +/// Implementors of this trait are responsible for converting raw bytes into +/// `RecordBatch` objects. +pub trait BatchDeserializer: Send + fmt::Debug { + /// Feeds a message for deserialization, updating the internal state of + /// this `BatchDeserializer`. Note that one can call this function multiple + /// times before calling `next`, which will queue multiple messages for + /// deserialization. Returns the number of bytes consumed. + fn digest(&mut self, message: T) -> usize; + + /// Attempts to deserialize any pending messages and returns a + /// `DeserializerOutput` to indicate progress. + fn next(&mut self) -> Result; + + /// Informs the deserializer that no more messages will be provided for + /// deserialization. + fn finish(&mut self); +} + +/// A general interface for decoders such as [`arrow::json::reader::Decoder`] and +/// [`arrow::csv::reader::Decoder`]. Defines an interface similar to +/// [`Decoder::decode`] and [`Decoder::flush`] methods, but also includes +/// a method to check if the decoder can flush early. Intended to be used in +/// conjunction with [`DecoderDeserializer`]. +/// +/// [`arrow::json::reader::Decoder`]: ::arrow::json::reader::Decoder +/// [`arrow::csv::reader::Decoder`]: ::arrow::csv::reader::Decoder +/// [`Decoder::decode`]: ::arrow::json::reader::Decoder::decode +/// [`Decoder::flush`]: ::arrow::json::reader::Decoder::flush +pub trait Decoder: Send + fmt::Debug { + /// See [`arrow::json::reader::Decoder::decode`]. + /// + /// [`arrow::json::reader::Decoder::decode`]: ::arrow::json::reader::Decoder::decode + fn decode(&mut self, buf: &[u8]) -> Result; + + /// See [`arrow::json::reader::Decoder::flush`]. + /// + /// [`arrow::json::reader::Decoder::flush`]: ::arrow::json::reader::Decoder::flush + fn flush(&mut self) -> Result, ArrowError>; + + /// Whether the decoder can flush early in its current state. + fn can_flush_early(&self) -> bool; +} + +impl fmt::Debug for DecoderDeserializer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Deserializer") + .field("buffered_queue", &self.buffered_queue) + .field("finalized", &self.finalized) + .finish() + } +} + +impl BatchDeserializer for DecoderDeserializer { + fn digest(&mut self, message: Bytes) -> usize { + if message.is_empty() { + return 0; + } + + let consumed = message.len(); + self.buffered_queue.push_back(message); + consumed + } + + fn next(&mut self) -> Result { + while let Some(buffered) = self.buffered_queue.front_mut() { + let decoded = self.decoder.decode(buffered)?; + buffered.advance(decoded); + + if buffered.is_empty() { + self.buffered_queue.pop_front(); + } + + // Flush when the stream ends or batch size is reached + // Certain implementations can flush early + if decoded == 0 || self.decoder.can_flush_early() { + return match self.decoder.flush() { + Ok(Some(batch)) => Ok(DeserializerOutput::RecordBatch(batch)), + Ok(None) => continue, + Err(e) => Err(e), + }; + } + } + if self.finalized { + Ok(DeserializerOutput::InputExhausted) + } else { + Ok(DeserializerOutput::RequiresMoreData) + } + } + + fn finish(&mut self) { + self.finalized = true; + // Ensure the decoder is flushed: + self.buffered_queue.push_back(Bytes::new()); + } +} + +/// A generic, decoder-based deserialization scheme for processing encoded data. +/// +/// This struct is responsible for converting a stream of bytes, which represent +/// encoded data, into a stream of `RecordBatch` objects, following the specified +/// schema and formatting options. It also handles any buffering necessary to satisfy +/// the `Decoder` interface. +pub struct DecoderDeserializer { + /// The underlying decoder used for deserialization + pub(crate) decoder: T, + /// The buffer used to store the remaining bytes to be decoded + pub(crate) buffered_queue: VecDeque, + /// Whether the input stream has been fully consumed + pub(crate) finalized: bool, +} + +impl DecoderDeserializer { + /// Creates a new `DecoderDeserializer` with the provided decoder. + pub fn new(decoder: T) -> Self { + DecoderDeserializer { + decoder, + buffered_queue: VecDeque::new(), + finalized: false, + } + } +} + +/// Deserializes a stream of bytes into a stream of [`RecordBatch`] objects using the +/// provided deserializer. +/// +/// Returns a boxed stream of `Result`. The stream yields [`RecordBatch`] +/// objects as they are produced by the deserializer, or an [`ArrowError`] if an error +/// occurs while polling the input or deserializing. +pub fn deserialize_stream<'a>( + mut input: impl Stream> + Unpin + Send + 'a, + mut deserializer: impl BatchDeserializer + 'a, +) -> BoxStream<'a, Result> { + futures::stream::poll_fn(move |cx| loop { + match ready!(input.poll_next_unpin(cx)).transpose()? { + Some(b) => _ = deserializer.digest(b), + None => deserializer.finish(), + }; + + return match deserializer.next()? { + DeserializerOutput::RecordBatch(rb) => Poll::Ready(Some(Ok(rb))), + DeserializerOutput::InputExhausted => Poll::Ready(None), + DeserializerOutput::RequiresMoreData => continue, + }; + }) + .boxed() +} diff --git a/datafusion/datasource/src/display.rs b/datafusion/datasource/src/display.rs index 58fc27bb8010..7ab8d407be52 100644 --- a/datafusion/datasource/src/display.rs +++ b/datafusion/datasource/src/display.rs @@ -36,7 +36,7 @@ impl DisplayAs for FileGroupsDisplay<'_> { let groups = if n_groups == 1 { "group" } else { "groups" }; write!(f, "{{{n_groups} {groups}: [")?; match t { - DisplayFormatType::Default => { + DisplayFormatType::Default | DisplayFormatType::TreeRender => { // To avoid showing too many partitions let max_groups = 5; fmt_up_to_n_elements(self.0, max_groups, f, |group, f| { @@ -66,7 +66,7 @@ impl DisplayAs for FileGroupDisplay<'_> { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult { write!(f, "[")?; match t { - DisplayFormatType::Default => { + DisplayFormatType::Default | DisplayFormatType::TreeRender => { // To avoid showing too many files let max_files = 5; fmt_up_to_n_elements(self.0, max_files, f, |pf, f| { diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 8d8cbbc67b9a..0066f39801a1 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -33,9 +33,9 @@ use datafusion_physical_plan::DisplayFormatType; use object_store::ObjectStore; -/// Common behaviors that every file format needs to implement. +/// Common file format behaviors needs to implement. /// -/// See initialization examples on `ParquetSource`, `CsvSource` +/// See implementation examples such as `ParquetSource`, `CsvSource` pub trait FileSource: Send + Sync { /// Creates a `dyn FileOpener` based on given parameters fn create_file_opener( diff --git a/datafusion/datasource/src/file_scan_config.rs b/datafusion/datasource/src/file_scan_config.rs index 79279b5c8231..91b5f0157739 100644 --- a/datafusion/datasource/src/file_scan_config.rs +++ b/datafusion/datasource/src/file_scan_config.rs @@ -194,26 +194,37 @@ impl DataSource for FileScanConfig { } fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> FmtResult { - let (schema, _, _, orderings) = self.project(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let (schema, _, _, orderings) = self.project(); - write!(f, "file_groups=")?; - FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?; + write!(f, "file_groups=")?; + FileGroupsDisplay(&self.file_groups).fmt_as(t, f)?; - if !schema.fields().is_empty() { - write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?; - } + if !schema.fields().is_empty() { + write!(f, ", projection={}", ProjectSchemaDisplay(&schema))?; + } - if let Some(limit) = self.limit { - write!(f, ", limit={limit}")?; - } + if let Some(limit) = self.limit { + write!(f, ", limit={limit}")?; + } - display_orderings(f, &orderings)?; + display_orderings(f, &orderings)?; - if !self.constraints.is_empty() { - write!(f, ", {}", self.constraints)?; - } + if !self.constraints.is_empty() { + write!(f, ", {}", self.constraints)?; + } - self.fmt_file_source(t, f) + self.fmt_file_source(t, f) + } + DisplayFormatType::TreeRender => { + writeln!(f, "format={}", self.file_source.file_type())?; + self.file_source.fmt_extra(t, f)?; + let num_files = self.file_groups.iter().map(Vec::len).sum::(); + writeln!(f, "files={num_files}")?; + Ok(()) + } + } } /// If supported by the underlying [`FileSource`], redistribute files across partitions according to their size. @@ -264,9 +275,20 @@ impl DataSource for FileScanConfig { &self, projection: &ProjectionExec, ) -> Result>> { - // If there is any non-column or alias-carrier expression, Projection should not be removed. // This process can be moved into CsvExec, but it would be an overlap of their responsibility. - Ok(all_alias_free_columns(projection.expr()).then(|| { + + // Must be all column references, with no table partition columns (which can not be projected) + let partitioned_columns_in_proj = projection.expr().iter().any(|(expr, _)| { + expr.as_any() + .downcast_ref::() + .map(|expr| expr.index() >= self.file_schema.fields().len()) + .unwrap_or(false) + }); + + // If there is any non-column or alias-carrier expression, Projection should not be removed. + let no_aliases = all_alias_free_columns(projection.expr()); + + Ok((no_aliases && !partitioned_columns_in_proj).then(|| { let file_scan = self.clone(); let source = Arc::clone(&file_scan.file_source); let new_projections = new_projections_for_columns( @@ -501,7 +523,6 @@ impl FileScanConfig { (schema, constraints, stats, output_ordering) } - #[cfg_attr(not(feature = "avro"), allow(unused))] // Only used by avro pub fn projected_file_column_names(&self) -> Option> { self.projection.as_ref().map(|p| { p.iter() diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 363ef4348e91..64fd56971b29 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -412,10 +412,10 @@ impl DataSource for MemorySourceConfig { .map_or(String::new(), |limit| format!(", fetch={}", limit)); if self.show_sizes { write!( - f, - "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", - partition_sizes.len(), - ) + f, + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) } else { write!( f, @@ -424,6 +424,19 @@ impl DataSource for MemorySourceConfig { ) } } + DisplayFormatType::TreeRender => { + let total_rows = self.partitions.iter().map(|b| b.len()).sum::(); + let total_bytes: usize = self + .partitions + .iter() + .flatten() + .map(|batch| batch.get_array_memory_size()) + .sum(); + writeln!(f, "format=memory")?; + writeln!(f, "rows={total_rows}")?; + writeln!(f, "bytes={total_bytes}")?; + Ok(()) + } } } diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 0ed59758476a..240e3c82bbfc 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -24,6 +24,7 @@ //! A table that uses the `ObjectStore` listing capability //! to get the list of files to process. +pub mod decoder; pub mod display; pub mod file; pub mod file_compression_type; @@ -34,17 +35,23 @@ pub mod file_scan_config; pub mod file_sink_config; pub mod file_stream; pub mod memory; +pub mod schema_adapter; pub mod source; mod statistics; + #[cfg(test)] mod test_util; + pub mod url; pub mod write; use chrono::TimeZone; use datafusion_common::Result; use datafusion_common::{ScalarValue, Statistics}; -use futures::Stream; +use file_meta::FileMeta; +use futures::{Stream, StreamExt}; use object_store::{path::Path, ObjectMeta}; +use object_store::{GetOptions, GetRange, ObjectStore}; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; @@ -190,6 +197,110 @@ impl From for PartitionedFile { } } +/// Represents the possible outcomes of a range calculation. +/// +/// This enum is used to encapsulate the result of calculating the range of +/// bytes to read from an object (like a file) in an object store. +/// +/// Variants: +/// - `Range(Option>)`: +/// Represents a range of bytes to be read. It contains an `Option` wrapping a +/// `Range`. `None` signifies that the entire object should be read, +/// while `Some(range)` specifies the exact byte range to read. +/// - `TerminateEarly`: +/// Indicates that the range calculation determined no further action is +/// necessary, possibly because the calculated range is empty or invalid. +pub enum RangeCalculation { + Range(Option>), + TerminateEarly, +} + +/// Calculates an appropriate byte range for reading from an object based on the +/// provided metadata. +/// +/// This asynchronous function examines the `FileMeta` of an object in an object store +/// and determines the range of bytes to be read. The range calculation may adjust +/// the start and end points to align with meaningful data boundaries (like newlines). +/// +/// Returns a `Result` wrapping a `RangeCalculation`, which is either a calculated byte range or an indication to terminate early. +/// +/// Returns an `Error` if any part of the range calculation fails, such as issues in reading from the object store or invalid range boundaries. +pub async fn calculate_range( + file_meta: &FileMeta, + store: &Arc, + terminator: Option, +) -> Result { + let location = file_meta.location(); + let file_size = file_meta.object_meta.size; + let newline = terminator.unwrap_or(b'\n'); + + match file_meta.range { + None => Ok(RangeCalculation::Range(None)), + Some(FileRange { start, end }) => { + let (start, end) = (start as usize, end as usize); + + let start_delta = if start != 0 { + find_first_newline(store, location, start - 1, file_size, newline).await? + } else { + 0 + }; + + let end_delta = if end != file_size { + find_first_newline(store, location, end - 1, file_size, newline).await? + } else { + 0 + }; + + let range = start + start_delta..end + end_delta; + + if range.start == range.end { + return Ok(RangeCalculation::TerminateEarly); + } + + Ok(RangeCalculation::Range(Some(range))) + } + } +} + +/// Asynchronously finds the position of the first newline character in a specified byte range +/// within an object, such as a file, in an object store. +/// +/// This function scans the contents of the object starting from the specified `start` position +/// up to the `end` position, looking for the first occurrence of a newline character. +/// It returns the position of the first newline relative to the start of the range. +/// +/// Returns a `Result` wrapping a `usize` that represents the position of the first newline character found within the specified range. If no newline is found, it returns the length of the scanned data, effectively indicating the end of the range. +/// +/// The function returns an `Error` if any issues arise while reading from the object store or processing the data stream. +/// +async fn find_first_newline( + object_store: &Arc, + location: &Path, + start: usize, + end: usize, + newline: u8, +) -> Result { + let options = GetOptions { + range: Some(GetRange::Bounded(start..end)), + ..Default::default() + }; + + let result = object_store.get_opts(location, options).await?; + let mut result_stream = result.into_stream(); + + let mut index = 0; + + while let Some(chunk) = result_stream.next().await.transpose()? { + if let Some(position) = chunk.iter().position(|&byte| byte == newline) { + return Ok(index + position); + } + + index += chunk.len(); + } + + Ok(index) +} + #[cfg(test)] mod tests { use super::ListingTableUrl; diff --git a/datafusion/core/src/datasource/schema_adapter.rs b/datafusion/datasource/src/schema_adapter.rs similarity index 69% rename from datafusion/core/src/datasource/schema_adapter.rs rename to datafusion/datasource/src/schema_adapter.rs index 8076c114ad16..e3a4ea4918c1 100644 --- a/datafusion/core/src/datasource/schema_adapter.rs +++ b/datafusion/datasource/src/schema_adapter.rs @@ -159,7 +159,7 @@ pub trait SchemaMapper: Debug + Send + Sync { /// ``` /// # use std::sync::Arc; /// # use arrow::datatypes::{DataType, Field, Schema}; -/// # use datafusion::datasource::schema_adapter::{DefaultSchemaAdapterFactory, SchemaAdapterFactory}; +/// # use datafusion_datasource::schema_adapter::{DefaultSchemaAdapterFactory, SchemaAdapterFactory}; /// # use datafusion_common::record_batch; /// // Table has fields "a", "b" and "c" /// let table_schema = Schema::new(vec![ @@ -427,223 +427,3 @@ impl SchemaMapper for SchemaMapping { Ok(record_batch) } } - -#[cfg(test)] -mod tests { - use std::fs; - use std::sync::Arc; - - use crate::assert_batches_sorted_eq; - use arrow::array::{Int32Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use arrow::record_batch::RecordBatch; - use datafusion_datasource::file_scan_config::FileScanConfig; - use object_store::path::Path; - use object_store::ObjectMeta; - - use crate::datasource::listing::PartitionedFile; - use crate::datasource::object_store::ObjectStoreUrl; - use crate::datasource::physical_plan::ParquetSource; - use crate::datasource::schema_adapter::{ - DefaultSchemaAdapterFactory, SchemaAdapter, SchemaAdapterFactory, SchemaMapper, - }; - use crate::physical_plan::collect; - use crate::prelude::SessionContext; - - use datafusion_common::record_batch; - #[cfg(feature = "parquet")] - use parquet::arrow::ArrowWriter; - use tempfile::TempDir; - - #[tokio::test] - async fn can_override_schema_adapter() { - // Test shows that SchemaAdapter can add a column that doesn't existing in the - // record batches returned from parquet. This can be useful for schema evolution - // where older files may not have all columns. - let tmp_dir = TempDir::new().unwrap(); - let table_dir = tmp_dir.path().join("parquet_test"); - fs::DirBuilder::new().create(table_dir.as_path()).unwrap(); - let f1 = Field::new("id", DataType::Int32, true); - - let file_schema = Arc::new(Schema::new(vec![f1.clone()])); - let filename = "part.parquet".to_string(); - let path = table_dir.as_path().join(filename.clone()); - let file = fs::File::create(path.clone()).unwrap(); - let mut writer = ArrowWriter::try_new(file, file_schema.clone(), None).unwrap(); - - let ids = Arc::new(Int32Array::from(vec![1i32])); - let rec_batch = RecordBatch::try_new(file_schema.clone(), vec![ids]).unwrap(); - - writer.write(&rec_batch).unwrap(); - writer.close().unwrap(); - - let location = Path::parse(path.to_str().unwrap()).unwrap(); - let metadata = fs::metadata(path.as_path()).expect("Local file metadata"); - let meta = ObjectMeta { - location, - last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), - size: metadata.len() as usize, - e_tag: None, - version: None, - }; - - let partitioned_file = PartitionedFile { - object_meta: meta, - partition_values: vec![], - range: None, - statistics: None, - extensions: None, - metadata_size_hint: None, - }; - - let f1 = Field::new("id", DataType::Int32, true); - let f2 = Field::new("extra_column", DataType::Utf8, true); - - let schema = Arc::new(Schema::new(vec![f1.clone(), f2.clone()])); - let source = Arc::new( - ParquetSource::default() - .with_schema_adapter_factory(Arc::new(TestSchemaAdapterFactory {})), - ); - let base_conf = - FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, source) - .with_file(partitioned_file); - - let parquet_exec = base_conf.build(); - - let session_ctx = SessionContext::new(); - let task_ctx = session_ctx.task_ctx(); - let read = collect(parquet_exec, task_ctx).await.unwrap(); - - let expected = [ - "+----+--------------+", - "| id | extra_column |", - "+----+--------------+", - "| 1 | foo |", - "+----+--------------+", - ]; - - assert_batches_sorted_eq!(expected, &read); - } - - #[test] - fn default_schema_adapter() { - let table_schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Utf8, true), - ]); - - // file has a subset of the table schema fields and different type - let file_schema = Schema::new(vec![ - Field::new("c", DataType::Float64, true), // not in table schema - Field::new("b", DataType::Float64, true), - ]); - - let adapter = DefaultSchemaAdapterFactory::from_schema(Arc::new(table_schema)); - let (mapper, indices) = adapter.map_schema(&file_schema).unwrap(); - assert_eq!(indices, vec![1]); - - let file_batch = record_batch!(("b", Float64, vec![1.0, 2.0])).unwrap(); - - let mapped_batch = mapper.map_batch(file_batch).unwrap(); - - // the mapped batch has the correct schema and the "b" column has been cast to Utf8 - let expected_batch = record_batch!( - ("a", Int32, vec![None, None]), // missing column filled with nulls - ("b", Utf8, vec!["1.0", "2.0"]) // b was cast to string and order was changed - ) - .unwrap(); - assert_eq!(mapped_batch, expected_batch); - } - - #[test] - fn default_schema_adapter_non_nullable_columns() { - let table_schema = Schema::new(vec![ - Field::new("a", DataType::Int32, false), // "a"" is declared non nullable - Field::new("b", DataType::Utf8, true), - ]); - let file_schema = Schema::new(vec![ - // since file doesn't have "a" it will be filled with nulls - Field::new("b", DataType::Float64, true), - ]); - - let adapter = DefaultSchemaAdapterFactory::from_schema(Arc::new(table_schema)); - let (mapper, indices) = adapter.map_schema(&file_schema).unwrap(); - assert_eq!(indices, vec![0]); - - let file_batch = record_batch!(("b", Float64, vec![1.0, 2.0])).unwrap(); - - // Mapping fails because it tries to fill in a non-nullable column with nulls - let err = mapper.map_batch(file_batch).unwrap_err().to_string(); - assert!(err.contains("Invalid argument error: Column 'a' is declared as non-nullable but contains null values"), "{err}"); - } - - #[derive(Debug)] - struct TestSchemaAdapterFactory; - - impl SchemaAdapterFactory for TestSchemaAdapterFactory { - fn create( - &self, - projected_table_schema: SchemaRef, - _table_schema: SchemaRef, - ) -> Box { - Box::new(TestSchemaAdapter { - table_schema: projected_table_schema, - }) - } - } - - struct TestSchemaAdapter { - /// Schema for the table - table_schema: SchemaRef, - } - - impl SchemaAdapter for TestSchemaAdapter { - fn map_column_index(&self, index: usize, file_schema: &Schema) -> Option { - let field = self.table_schema.field(index); - Some(file_schema.fields.find(field.name())?.0) - } - - fn map_schema( - &self, - file_schema: &Schema, - ) -> datafusion_common::Result<(Arc, Vec)> { - let mut projection = Vec::with_capacity(file_schema.fields().len()); - - for (file_idx, file_field) in file_schema.fields.iter().enumerate() { - if self.table_schema.fields().find(file_field.name()).is_some() { - projection.push(file_idx); - } - } - - Ok((Arc::new(TestSchemaMapping {}), projection)) - } - } - - #[derive(Debug)] - struct TestSchemaMapping {} - - impl SchemaMapper for TestSchemaMapping { - fn map_batch( - &self, - batch: RecordBatch, - ) -> datafusion_common::Result { - let f1 = Field::new("id", DataType::Int32, true); - let f2 = Field::new("extra_column", DataType::Utf8, true); - - let schema = Arc::new(Schema::new(vec![f1, f2])); - - let extra_column = Arc::new(StringArray::from(vec!["foo"])); - let mut new_columns = batch.columns().to_vec(); - new_columns.push(extra_column); - - Ok(RecordBatch::try_new(schema, new_columns).unwrap()) - } - - fn map_partial_batch( - &self, - batch: RecordBatch, - ) -> datafusion_common::Result { - self.map_batch(batch) - } - } -} diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index b3089a6e59fe..6e78df760dc3 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +//! [`DataSource`] and [`DataSourceExec`] + use std::any::Any; use std::fmt; use std::fmt::{Debug, Formatter}; @@ -34,9 +36,15 @@ use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; /// Common behaviors in Data Sources for both from Files and Memory. -/// See `DataSourceExec` for physical plan implementation /// +/// # See Also +/// * [`DataSourceExec`] for physical plan implementation +/// * [`FileSource`] for file format implementations (Parquet, Json, etc) +/// +/// # Notes /// Requires `Debug` to assist debugging +/// +/// [`FileSource`]: crate::file::FileSource pub trait DataSource: Send + Sync + Debug { fn open( &self, @@ -44,6 +52,7 @@ pub trait DataSource: Send + Sync + Debug { context: Arc, ) -> datafusion_common::Result; fn as_any(&self) -> &dyn Any; + /// Format this source for display in explain plans fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result; /// Return a copy of this DataSource with a new partitioning scheme @@ -71,16 +80,32 @@ pub trait DataSource: Send + Sync + Debug { ) -> datafusion_common::Result>>; } -/// Unified data source for file formats like JSON, CSV, AVRO, ARROW, PARQUET +/// [`ExecutionPlan`] handles different file formats like JSON, CSV, AVRO, ARROW, PARQUET +/// +/// `DataSourceExec` implements common functionality such as applying projections, +/// and caching plan properties. +/// +/// The [`DataSource`] trait describes where to find the data for this data +/// source (for example what files or what in memory partitions). Format +/// specifics are implemented with the [`FileSource`] trait. +/// +/// [`FileSource`]: crate::file::FileSource #[derive(Clone, Debug)] pub struct DataSourceExec { + /// The source of the data -- for example, `FileScanConfig` or `MemorySourceConfig` data_source: Arc, + /// Cached plan properties such as sort order cache: PlanProperties, } impl DisplayAs for DataSourceExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { - write!(f, "DataSourceExec: ")?; + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "DataSourceExec: ")?; + } + DisplayFormatType::TreeRender => {} + } self.data_source.fmt_as(t, f) } } diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index ab025069bf76..9a9b98d5041b 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -15,21 +15,20 @@ // specific language governing permissions and limitations // under the License. +use crate::{ + file::FileSource, file_scan_config::FileScanConfig, file_stream::FileOpener, +}; + use std::sync::Arc; use arrow::datatypes::SchemaRef; -use datafusion_common::Statistics; +use datafusion_common::{Result, Statistics}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use object_store::ObjectStore; -use crate::{ - file::FileSource, file_scan_config::FileScanConfig, file_stream::FileOpener, -}; -use datafusion_common::Result; - -/// Minimal [`FileSource`] implementation for use in tests. +/// Minimal [`crate::file::FileSource`] implementation for use in tests. #[derive(Clone, Default)] -pub struct MockSource { +pub(crate) struct MockSource { metrics: ExecutionPlanMetricsSet, projected_statistics: Option, } diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index ba6fadbf7235..063417a254be 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -358,6 +358,8 @@ pub enum ArrayFunctionArgument { /// An argument of type List/LargeList/FixedSizeList. All Array arguments must be coercible /// to the same type. Array, + // A Utf8 argument. + String, } impl Display for ArrayFunctionArgument { @@ -372,6 +374,9 @@ impl Display for ArrayFunctionArgument { ArrayFunctionArgument::Array => { write!(f, "array") } + ArrayFunctionArgument::String => { + write!(f, "string") + } } } } diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index f8baf9c94b3c..56279632251b 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -25,7 +25,6 @@ use std::sync::Arc; use crate::expr_fn::binary_expr; use crate::logical_plan::Subquery; -use crate::utils::expr_to_columns; use crate::Volatility; use crate::{udaf, ExprSchemable, Operator, Signature, WindowFrame, WindowUDF}; @@ -35,7 +34,7 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion, }; use datafusion_common::{ - plan_err, Column, DFSchema, HashMap, Result, ScalarValue, Spans, TableReference, + Column, DFSchema, HashMap, Result, ScalarValue, Spans, TableReference, }; use datafusion_functions_window_common::field::WindowUDFFieldArgs; use sqlparser::ast::{ @@ -311,6 +310,10 @@ pub enum Expr { /// /// This expr has to be resolved to a list of columns before translating logical /// plan into physical plan. + #[deprecated( + since = "46.0.0", + note = "A wildcard needs to be resolved to concrete expressions when constructing the logical plan. See https://github.com/apache/datafusion/issues/7765" + )] Wildcard { qualifier: Option, options: Box, @@ -1086,11 +1089,6 @@ impl PlannedReplaceSelectItem { } impl Expr { - #[deprecated(since = "40.0.0", note = "use schema_name instead")] - pub fn display_name(&self) -> Result { - Ok(self.schema_name().to_string()) - } - /// The name of the column (field) that this `Expr` will produce. /// /// For example, for a projection (e.g. `SELECT `) the resulting arrow @@ -1175,6 +1173,7 @@ impl Expr { Expr::ScalarVariable(..) => "ScalarVariable", Expr::TryCast { .. } => "TryCast", Expr::WindowFunction { .. } => "WindowFunction", + #[expect(deprecated)] Expr::Wildcard { .. } => "Wildcard", Expr::Unnest { .. } => "Unnest", } @@ -1439,15 +1438,6 @@ impl Expr { Box::new(high), )) } - - #[deprecated(since = "39.0.0", note = "use try_as_col instead")] - pub fn try_into_col(&self) -> Result { - match self { - Expr::Column(it) => Ok(it.clone()), - _ => plan_err!("Could not coerce '{self}' into Column!"), - } - } - /// Return a reference to the inner `Column` if any /// /// returns `None` if the expression is not a `Column` @@ -1490,15 +1480,6 @@ impl Expr { } } - /// Return all referenced columns of this expression. - #[deprecated(since = "40.0.0", note = "use Expr::column_refs instead")] - pub fn to_columns(&self) -> Result> { - let mut using_columns = HashSet::new(); - expr_to_columns(self, &mut using_columns)?; - - Ok(using_columns) - } - /// Return all references to columns in this expression. /// /// # Example @@ -1648,6 +1629,8 @@ impl Expr { // Use explicit pattern match instead of a default // implementation, so that in the future if someone adds // new Expr types, they will check here as well + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::AggregateFunction(..) | Expr::Alias(..) | Expr::Between(..) @@ -2229,6 +2212,7 @@ impl HashNode for Expr { Expr::ScalarSubquery(subquery) => { subquery.hash(state); } + #[expect(deprecated)] Expr::Wildcard { qualifier, options } => { qualifier.hash(state); options.hash(state); @@ -2288,6 +2272,8 @@ impl Display for SchemaDisplay<'_> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match self.0 { // The same as Display + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::Column(_) | Expr::Literal(_) | Expr::ScalarVariable(..) @@ -2758,6 +2744,7 @@ impl Display for Expr { write!(f, "{expr} IN ([{}])", expr_vec_fmt!(list)) } } + #[expect(deprecated)] Expr::Wildcard { qualifier, options } => match qualifier { Some(qualifier) => write!(f, "{qualifier}.*{options}"), None => write!(f, "*{options}"), diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs index f47de4a8178f..a8e7fd76d037 100644 --- a/datafusion/expr/src/expr_fn.rs +++ b/datafusion/expr/src/expr_fn.rs @@ -121,6 +121,7 @@ pub fn placeholder(id: impl Into) -> Expr { /// assert_eq!(p.to_string(), "*") /// ``` pub fn wildcard() -> Expr { + #[expect(deprecated)] Expr::Wildcard { qualifier: None, options: Box::new(WildcardOptions::default()), @@ -129,6 +130,7 @@ pub fn wildcard() -> Expr { /// Create an '*' [`Expr::Wildcard`] expression with the wildcard options pub fn wildcard_with_options(options: WildcardOptions) -> Expr { + #[expect(deprecated)] Expr::Wildcard { qualifier: None, options: Box::new(options), @@ -146,6 +148,7 @@ pub fn wildcard_with_options(options: WildcardOptions) -> Expr { /// assert_eq!(p.to_string(), "t.*") /// ``` pub fn qualified_wildcard(qualifier: impl Into) -> Expr { + #[expect(deprecated)] Expr::Wildcard { qualifier: Some(qualifier.into()), options: Box::new(WildcardOptions::default()), @@ -157,6 +160,7 @@ pub fn qualified_wildcard_with_options( qualifier: impl Into, options: WildcardOptions, ) -> Expr { + #[expect(deprecated)] Expr::Wildcard { qualifier: Some(qualifier.into()), options: Box::new(options), diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index c6739c1f48a4..90dcbce46b01 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -253,6 +253,7 @@ fn coerce_exprs_for_schema( Expr::Alias(Alias { expr, name, .. }) => { Ok(expr.cast_to(new_type, src_schema)?.alias(name)) } + #[expect(deprecated)] Expr::Wildcard { .. } => Ok(expr), _ => expr.cast_to(new_type, src_schema), } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index ce1dd2f34c05..0a14cb5c60a0 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -215,6 +215,7 @@ impl ExprSchemable for Expr { Ok(DataType::Null) } } + #[expect(deprecated)] Expr::Wildcard { .. } => Ok(DataType::Null), Expr::GroupingSet(_) => { // Grouping sets do not really have a type and do not appear in projections @@ -329,6 +330,7 @@ impl ExprSchemable for Expr { | Expr::SimilarTo(Like { expr, pattern, .. }) => { Ok(expr.nullable(input_schema)? || pattern.nullable(input_schema)?) } + #[expect(deprecated)] Expr::Wildcard { .. } => Ok(false), Expr::GroupingSet(_) => { // Grouping sets do not really have the concept of nullable and do not appear diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 2bb15da21863..f60bb2f00771 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1675,6 +1675,7 @@ fn project_with_validation( for (e, validate) in expr { let e = e.into(); match e { + #[expect(deprecated)] Expr::Wildcard { .. } => projected_expr.push(e), _ => { if validate { diff --git a/datafusion/expr/src/logical_plan/extension.rs b/datafusion/expr/src/logical_plan/extension.rs index be7153cc4eaa..5bf64a36a654 100644 --- a/datafusion/expr/src/logical_plan/extension.rs +++ b/datafusion/expr/src/logical_plan/extension.rs @@ -82,17 +82,6 @@ pub trait UserDefinedLogicalNode: fmt::Debug + Send + Sync { /// For example: `TopK: k=10` fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result; - #[deprecated(since = "39.0.0", note = "use with_exprs_and_inputs instead")] - #[allow(clippy::wrong_self_convention)] - fn from_template( - &self, - exprs: &[Expr], - inputs: &[LogicalPlan], - ) -> Arc { - self.with_exprs_and_inputs(exprs.to_vec(), inputs.to_vec()) - .unwrap() - } - /// Create a new `UserDefinedLogicalNode` with the specified children /// and expressions. This function is used during optimization /// when the plan is being rewritten and a new instance of the @@ -282,13 +271,6 @@ pub trait UserDefinedLogicalNodeCore: /// For example: `TopK: k=10` fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result; - #[deprecated(since = "39.0.0", note = "use with_exprs_and_inputs instead")] - #[allow(clippy::wrong_self_convention)] - fn from_template(&self, exprs: &[Expr], inputs: &[LogicalPlan]) -> Self { - self.with_exprs_and_inputs(exprs.to_vec(), inputs.to_vec()) - .unwrap() - } - /// Create a new `UserDefinedLogicalNode` with the specified children /// and expressions. This function is used during optimization /// when the plan is being rewritten and a new instance of the diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c6fd95595233..0534b04f5dc3 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -38,9 +38,8 @@ use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; use crate::logical_plan::{DmlStatement, Statement}; use crate::utils::{ - enumerate_grouping_sets, exprlist_len, exprlist_to_fields, find_base_plan, - find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, - split_conjunction, + enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, + grouping_set_expr_count, grouping_set_to_exprlist, split_conjunction, }; use crate::{ build_join_schema, expr_vec_fmt, BinaryExpr, CreateMemoryTable, CreateView, Execute, @@ -903,7 +902,7 @@ impl LogicalPlan { let (left, right) = self.only_two_inputs(inputs)?; let schema = build_join_schema(left.schema(), right.schema(), join_type)?; - let equi_expr_count = on.len(); + let equi_expr_count = on.len() * 2; assert!(expr.len() >= equi_expr_count); // Assume that the last expr, if any, @@ -917,17 +916,16 @@ impl LogicalPlan { // The first part of expr is equi-exprs, // and the struct of each equi-expr is like `left-expr = right-expr`. assert_eq!(expr.len(), equi_expr_count); - let new_on = expr.into_iter().map(|equi_expr| { + let mut new_on = Vec::with_capacity(on.len()); + let mut iter = expr.into_iter(); + while let Some(left) = iter.next() { + let Some(right) = iter.next() else { + internal_err!("Expected a pair of expressions to construct the join on expression")? + }; + // SimplifyExpression rule may add alias to the equi_expr. - let unalias_expr = equi_expr.clone().unalias(); - if let Expr::BinaryExpr(BinaryExpr { left, op: Operator::Eq, right }) = unalias_expr { - Ok((*left, *right)) - } else { - internal_err!( - "The front part expressions should be an binary equality expression, actual:{equi_expr}" - ) - } - }).collect::>>()?; + new_on.push((left.unalias(), right.unalias())); + } Ok(LogicalPlan::Join(Join { left: Arc::new(left), @@ -2141,6 +2139,7 @@ impl Projection { input: Arc, schema: DFSchemaRef, ) -> Result { + #[expect(deprecated)] if !expr.iter().any(|e| matches!(e, Expr::Wildcard { .. })) && expr.len() != schema.fields().len() { @@ -3451,6 +3450,7 @@ fn calc_func_dependencies_for_project( let proj_indices = exprs .iter() .map(|expr| match expr { + #[expect(deprecated)] Expr::Wildcard { qualifier, options } => { let wildcard_fields = exprlist_to_fields( vec![&Expr::Wildcard { @@ -3493,11 +3493,10 @@ fn calc_func_dependencies_for_project( .flatten() .collect::>(); - let len = exprlist_len(exprs, input.schema(), Some(find_base_plan(input).schema()))?; Ok(input .schema() .functional_dependencies() - .project_functional_dependencies(&proj_indices, len)) + .project_functional_dependencies(&proj_indices, exprs.len())) } /// Sorts its input according to a list of sort expressions. @@ -3778,7 +3777,8 @@ mod tests { use crate::builder::LogicalTableSource; use crate::logical_plan::table_scan; use crate::{ - col, exists, in_subquery, lit, placeholder, scalar_subquery, GroupingSet, + binary_expr, col, exists, in_subquery, lit, placeholder, scalar_subquery, + GroupingSet, }; use datafusion_common::tree_node::{ @@ -4630,4 +4630,110 @@ digraph { let parameter_type = params.clone().get(placeholder_value).unwrap().clone(); assert_eq!(parameter_type, None); } + + #[test] + fn test_join_with_new_exprs() -> Result<()> { + fn create_test_join( + on: Vec<(Expr, Expr)>, + filter: Option, + ) -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + + let left_schema = DFSchema::try_from_qualified_schema("t1", &schema)?; + let right_schema = DFSchema::try_from_qualified_schema("t2", &schema)?; + + Ok(LogicalPlan::Join(Join { + left: Arc::new( + table_scan(Some("t1"), left_schema.as_arrow(), None)?.build()?, + ), + right: Arc::new( + table_scan(Some("t2"), right_schema.as_arrow(), None)?.build()?, + ), + on, + filter, + join_type: JoinType::Inner, + join_constraint: JoinConstraint::On, + schema: Arc::new(left_schema.join(&right_schema)?), + null_equals_null: false, + })) + } + + { + let join = create_test_join(vec![(col("t1.a"), (col("t2.a")))], None)?; + let LogicalPlan::Join(join) = join.with_new_exprs( + join.expressions(), + join.inputs().into_iter().cloned().collect(), + )? + else { + unreachable!() + }; + assert_eq!(join.on, vec![(col("t1.a"), (col("t2.a")))]); + assert_eq!(join.filter, None); + } + + { + let join = create_test_join(vec![], Some(col("t1.a").gt(col("t2.a"))))?; + let LogicalPlan::Join(join) = join.with_new_exprs( + join.expressions(), + join.inputs().into_iter().cloned().collect(), + )? + else { + unreachable!() + }; + assert_eq!(join.on, vec![]); + assert_eq!(join.filter, Some(col("t1.a").gt(col("t2.a")))); + } + + { + let join = create_test_join( + vec![(col("t1.a"), (col("t2.a")))], + Some(col("t1.b").gt(col("t2.b"))), + )?; + let LogicalPlan::Join(join) = join.with_new_exprs( + join.expressions(), + join.inputs().into_iter().cloned().collect(), + )? + else { + unreachable!() + }; + assert_eq!(join.on, vec![(col("t1.a"), (col("t2.a")))]); + assert_eq!(join.filter, Some(col("t1.b").gt(col("t2.b")))); + } + + { + let join = create_test_join( + vec![(col("t1.a"), (col("t2.a"))), (col("t1.b"), (col("t2.b")))], + None, + )?; + let LogicalPlan::Join(join) = join.with_new_exprs( + vec![ + binary_expr(col("t1.a"), Operator::Plus, lit(1)), + binary_expr(col("t2.a"), Operator::Plus, lit(2)), + col("t1.b"), + col("t2.b"), + lit(true), + ], + join.inputs().into_iter().cloned().collect(), + )? + else { + unreachable!() + }; + assert_eq!( + join.on, + vec![ + ( + binary_expr(col("t1.a"), Operator::Plus, lit(1)), + binary_expr(col("t2.a"), Operator::Plus, lit(2)) + ), + (col("t1.b"), (col("t2.b"))) + ] + ); + assert_eq!(join.filter, Some(lit(true))); + } + + Ok(()) + } } diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs index 50af62060346..49cc79c60a27 100644 --- a/datafusion/expr/src/tree_node.rs +++ b/datafusion/expr/src/tree_node.rs @@ -67,6 +67,8 @@ impl TreeNode for Expr { Expr::GroupingSet(GroupingSet::GroupingSets(lists_of_exprs)) => { lists_of_exprs.apply_elements(f) } + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::Column(_) // Treat OuterReferenceColumn as a leaf expression | Expr::OuterReferenceColumn(_, _) @@ -113,6 +115,8 @@ impl TreeNode for Expr { mut f: F, ) -> Result> { Ok(match self { + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::Column(_) | Expr::Wildcard { .. } | Expr::Placeholder(Placeholder { .. }) diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index b471feca043f..0ec017bdc27f 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -19,7 +19,7 @@ use super::binary::{binary_numeric_coercion, comparison_coercion}; use crate::{AggregateUDF, ScalarUDF, Signature, TypeSignature, WindowUDF}; use arrow::{ compute::can_cast_types, - datatypes::{DataType, TimeUnit}, + datatypes::{DataType, Field, TimeUnit}, }; use datafusion_common::types::LogicalType; use datafusion_common::utils::{coerced_fixed_size_list_to_list, ListCoercion}; @@ -387,7 +387,7 @@ fn get_valid_types( new_base_type = coerce_array_types(function_name, current_type, &new_base_type)?; } - ArrayFunctionArgument::Index => {} + ArrayFunctionArgument::Index | ArrayFunctionArgument::String => {} } } let new_array_type = datafusion_common::utils::coerced_type_with_base_type_only( @@ -408,6 +408,7 @@ fn get_valid_types( let valid_type = match argument_type { ArrayFunctionArgument::Element => new_elem_type.clone(), ArrayFunctionArgument::Index => DataType::Int64, + ArrayFunctionArgument::String => DataType::Utf8, ArrayFunctionArgument::Array => { let Some(current_type) = array(current_type) else { return Ok(vec![vec![]]); @@ -435,6 +436,10 @@ fn get_valid_types( match array_type { DataType::List(_) | DataType::LargeList(_) => Some(array_type.clone()), DataType::FixedSizeList(field, _) => Some(DataType::List(Arc::clone(field))), + DataType::Null => Some(DataType::List(Arc::new(Field::new_list_field( + DataType::Int64, + true, + )))), _ => None, } } diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 8215b671a379..86bc5852b830 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -172,7 +172,7 @@ impl ScalarUDF { /// /// # Notes /// - /// If a function implement [`ScalarUDFImpl::return_type_from_exprs`], + /// If a function implement [`ScalarUDFImpl::return_type_from_args`], /// its [`ScalarUDFImpl::return_type`] should raise an error. /// /// See [`ScalarUDFImpl::return_type`] for more details. @@ -180,22 +180,6 @@ impl ScalarUDF { self.inner.return_type(arg_types) } - /// The datatype this function returns given the input argument input types. - /// This function is used when the input arguments are [`Expr`]s. - /// - /// - /// See [`ScalarUDFImpl::return_type_from_exprs`] for more details. - #[allow(deprecated)] - pub fn return_type_from_exprs( - &self, - args: &[Expr], - schema: &dyn ExprSchema, - arg_types: &[DataType], - ) -> Result { - // If the implementation provides a return_type_from_exprs, use it - self.inner.return_type_from_exprs(args, schema, arg_types) - } - /// Return the datatype this function returns given the input argument types. /// /// See [`ScalarUDFImpl::return_type_from_args`] for more details. @@ -225,11 +209,13 @@ impl ScalarUDF { self.inner.is_nullable(args, schema) } + #[deprecated(since = "46.0.0", note = "Use `invoke_with_args` instead")] pub fn invoke_batch( &self, args: &[ColumnarValue], number_rows: usize, ) -> Result { + #[allow(deprecated)] self.inner.invoke_batch(args, number_rows) } @@ -244,7 +230,7 @@ impl ScalarUDF { /// /// Note: This method is deprecated and will be removed in future releases. /// User defined functions should implement [`Self::invoke_with_args`] instead. - #[deprecated(since = "42.1.0", note = "Use `invoke_batch` instead")] + #[deprecated(since = "42.1.0", note = "Use `invoke_with_args` instead")] pub fn invoke_no_args(&self, number_rows: usize) -> Result { #[allow(deprecated)] self.inner.invoke_no_args(number_rows) @@ -252,7 +238,7 @@ impl ScalarUDF { /// Returns a `ScalarFunctionImplementation` that can invoke the function /// during execution - #[deprecated(since = "42.0.0", note = "Use `invoke_batch` instead")] + #[deprecated(since = "42.0.0", note = "Use `invoke_with_args` instead")] pub fn fun(&self) -> ScalarFunctionImplementation { let captured = Arc::clone(&self.inner); #[allow(deprecated)] @@ -349,7 +335,7 @@ pub struct ScalarFunctionArgs<'a> { pub args: Vec, /// The number of rows in record batch being evaluated pub number_rows: usize, - /// The return type of the scalar function returned (from `return_type` or `return_type_from_exprs`) + /// The return type of the scalar function returned (from `return_type` or `return_type_from_args`) /// when creating the physical expression from the logical expression pub return_type: &'a DataType, } @@ -538,16 +524,6 @@ pub trait ScalarUDFImpl: Debug + Send + Sync { /// [`DataFusionError::Internal`]: datafusion_common::DataFusionError::Internal fn return_type(&self, arg_types: &[DataType]) -> Result; - #[deprecated(since = "45.0.0", note = "Use `return_type_from_args` instead")] - fn return_type_from_exprs( - &self, - _args: &[Expr], - _schema: &dyn ExprSchema, - arg_types: &[DataType], - ) -> Result { - self.return_type(arg_types) - } - /// What type will be returned by this function, given the arguments? /// /// By default, this function calls [`Self::return_type`] with the @@ -613,6 +589,7 @@ pub trait ScalarUDFImpl: Debug + Send + Sync { /// User defined functions should implement [`Self::invoke_with_args`] instead. /// /// See for more details. + #[deprecated(since = "46.0.0", note = "Use `invoke_with_args` instead")] fn invoke_batch( &self, args: &[ColumnarValue], @@ -643,6 +620,7 @@ pub trait ScalarUDFImpl: Debug + Send + Sync { /// [`ColumnarValue::values_to_arrays`] can be used to convert the arguments /// to arrays, which will likely be simpler code, but be slower. fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + #[allow(deprecated)] self.invoke_batch(&args.args, args.number_rows) } @@ -885,16 +863,6 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { &self.aliases } - #[allow(deprecated)] - fn return_type_from_exprs( - &self, - args: &[Expr], - schema: &dyn ExprSchema, - arg_types: &[DataType], - ) -> Result { - self.inner.return_type_from_exprs(args, schema, arg_types) - } - fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result { self.inner.return_type_from_args(args) } diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 56c1e64554a9..9d7def76a11e 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -19,7 +19,6 @@ use std::cmp::Ordering; use std::collections::{BTreeSet, HashSet}; -use std::ops::Deref; use std::sync::Arc; use crate::expr::{Alias, Sort, WildcardOptions, WindowFunction, WindowFunctionParams}; @@ -48,16 +47,6 @@ pub use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; /// `COUNT()` expressions pub use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; -/// Recursively walk a list of expression trees, collecting the unique set of columns -/// referenced in the expression -#[deprecated(since = "40.0.0", note = "Expr::add_column_refs instead")] -pub fn exprlist_to_columns(expr: &[Expr], accum: &mut HashSet) -> Result<()> { - for e in expr { - expr_to_columns(e, accum)?; - } - Ok(()) -} - /// Count the number of distinct exprs in a list of group by expressions. If the /// first element is a `GroupingSet` expression then it must be the only expr. pub fn grouping_set_expr_count(group_expr: &[Expr]) -> Result { @@ -282,6 +271,8 @@ pub fn expr_to_columns(expr: &Expr, accum: &mut HashSet) -> Result<()> { // Use explicit pattern match instead of a default // implementation, so that in the future if someone adds // new Expr types, they will check here as well + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::Unnest(_) | Expr::ScalarVariable(_, _) | Expr::Alias(_) @@ -704,162 +695,11 @@ pub fn exprlist_to_fields<'a>( plan: &LogicalPlan, ) -> Result, Arc)>> { // Look for exact match in plan's output schema - let wildcard_schema = find_base_plan(plan).schema(); let input_schema = plan.schema(); - let result = exprs - .into_iter() - .map(|e| match e { - Expr::Wildcard { qualifier, options } => match qualifier { - None => { - let mut excluded = exclude_using_columns(plan)?; - excluded.extend(get_excluded_columns( - options.exclude.as_ref(), - options.except.as_ref(), - wildcard_schema, - None, - )?); - Ok(wildcard_schema - .iter() - .filter(|(q, f)| { - !excluded.contains(&Column::new(q.cloned(), f.name())) - }) - .map(|(q, f)| (q.cloned(), Arc::clone(f))) - .collect::>()) - } - Some(qualifier) => { - let excluded: Vec = get_excluded_columns( - options.exclude.as_ref(), - options.except.as_ref(), - wildcard_schema, - Some(qualifier), - )? - .into_iter() - .map(|c| c.flat_name()) - .collect(); - Ok(wildcard_schema - .fields_with_qualified(qualifier) - .into_iter() - .filter_map(|field| { - let flat_name = format!("{}.{}", qualifier, field.name()); - if excluded.contains(&flat_name) { - None - } else { - Some(( - Some(qualifier.clone()), - Arc::new(field.to_owned()), - )) - } - }) - .collect::>()) - } - }, - _ => Ok(vec![e.to_field(input_schema)?]), - }) - .collect::>>()? - .into_iter() - .flatten() - .collect(); - Ok(result) -} - -/// Find the suitable base plan to expand the wildcard expression recursively. -/// When planning [LogicalPlan::Window] and [LogicalPlan::Aggregate], we will generate -/// an intermediate plan based on the relation plan (e.g. [LogicalPlan::TableScan], [LogicalPlan::Subquery], ...). -/// If we expand a wildcard expression basing the intermediate plan, we could get some duplicate fields. -pub fn find_base_plan(input: &LogicalPlan) -> &LogicalPlan { - match input { - LogicalPlan::Window(window) => find_base_plan(&window.input), - LogicalPlan::Aggregate(agg) => find_base_plan(&agg.input), - // [SqlToRel::try_process_unnest] will convert Expr(Unnest(Expr)) to Projection/Unnest/Projection - // We should expand the wildcard expression based on the input plan of the inner Projection. - LogicalPlan::Unnest(unnest) => { - if let LogicalPlan::Projection(projection) = unnest.input.deref() { - find_base_plan(&projection.input) - } else { - input - } - } - LogicalPlan::Filter(filter) => { - if filter.having { - // If a filter is used for a having clause, its input plan is an aggregation. - // We should expand the wildcard expression based on the aggregation's input plan. - find_base_plan(&filter.input) - } else { - input - } - } - _ => input, - } -} - -/// Count the number of real fields. We should expand the wildcard expression to get the actual number. -pub fn exprlist_len( - exprs: &[Expr], - schema: &DFSchemaRef, - wildcard_schema: Option<&DFSchemaRef>, -) -> Result { exprs - .iter() - .map(|e| match e { - Expr::Wildcard { - qualifier: None, - options, - } => { - let excluded = get_excluded_columns( - options.exclude.as_ref(), - options.except.as_ref(), - wildcard_schema.unwrap_or(schema), - None, - )? - .into_iter() - .collect::>(); - Ok( - get_exprs_except_skipped(wildcard_schema.unwrap_or(schema), excluded) - .len(), - ) - } - Expr::Wildcard { - qualifier: Some(qualifier), - options, - } => { - let related_wildcard_schema = wildcard_schema.as_ref().map_or_else( - || Ok(Arc::clone(schema)), - |schema| { - // Eliminate the fields coming from other tables. - let qualified_fields = schema - .fields() - .iter() - .enumerate() - .filter_map(|(idx, field)| { - let (maybe_table_ref, _) = schema.qualified_field(idx); - if maybe_table_ref.is_none_or(|q| q == qualifier) { - Some((maybe_table_ref.cloned(), Arc::clone(field))) - } else { - None - } - }) - .collect::>(); - let metadata = schema.metadata().clone(); - DFSchema::new_with_metadata(qualified_fields, metadata) - .map(Arc::new) - }, - )?; - let excluded = get_excluded_columns( - options.exclude.as_ref(), - options.except.as_ref(), - related_wildcard_schema.as_ref(), - Some(qualifier), - )? - .into_iter() - .collect::>(); - Ok( - get_exprs_except_skipped(related_wildcard_schema.as_ref(), excluded) - .len(), - ) - } - _ => Ok(1), - }) - .sum() + .into_iter() + .map(|e| e.to_field(input_schema)) + .collect() } /// Convert an expression into Column expression if it's already provided as input plan. diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index 97914666688f..5c80c1b04225 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -47,7 +47,7 @@ datafusion-proto = { workspace = true } futures = { workspace = true } log = { workspace = true } prost = { workspace = true } -semver = "1.0.24" +semver = "1.0.26" tokio = { workspace = true } [dev-dependencies] diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 00602474d621..14a0908c4795 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -21,12 +21,12 @@ use abi_stable::{ std_types::{RResult, RString, RVec}, StableAbi, }; -use datafusion::error::Result; use datafusion::{ error::DataFusionError, execution::{SendableRecordBatchStream, TaskContext}, physical_plan::{DisplayAs, ExecutionPlan, PlanProperties}, }; +use datafusion::{error::Result, physical_plan::DisplayFormatType}; use tokio::runtime::Handle; use crate::{ @@ -198,14 +198,22 @@ unsafe impl Sync for ForeignExecutionPlan {} impl DisplayAs for ForeignExecutionPlan { fn fmt_as( &self, - _t: datafusion::physical_plan::DisplayFormatType, + t: DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - write!( - f, - "FFI_ExecutionPlan(number_of_children={})", - self.children.len(), - ) + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "FFI_ExecutionPlan(number_of_children={})", + self.children.len(), + ) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } } } @@ -315,7 +323,7 @@ mod tests { impl DisplayAs for EmptyExec { fn fmt_as( &self, - _t: datafusion::physical_plan::DisplayFormatType, + _t: DisplayFormatType, _f: &mut std::fmt::Formatter, ) -> std::fmt::Result { unimplemented!() diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index a3339f0fceb9..2d995b4a4179 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -17,6 +17,7 @@ use ahash::RandomState; use datafusion_common::stats::Precision; +use datafusion_expr::expr::WindowFunction; use datafusion_functions_aggregate_common::aggregate::count_distinct::BytesViewDistinctCountAccumulator; use datafusion_macros::user_doc; use datafusion_physical_expr::expressions; @@ -51,7 +52,9 @@ use datafusion_expr::{ function::AccumulatorArgs, utils::format_state_name, Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, SetMonotonicity, Signature, Volatility, }; -use datafusion_expr::{Expr, ReversedUDAF, StatisticsArgs, TypeSignature}; +use datafusion_expr::{ + Expr, ReversedUDAF, StatisticsArgs, TypeSignature, WindowFunctionDefinition, +}; use datafusion_functions_aggregate_common::aggregate::count_distinct::{ BytesDistinctCountAccumulator, FloatDistinctCountAccumulator, PrimitiveDistinctCountAccumulator, @@ -79,9 +82,51 @@ pub fn count_distinct(expr: Expr) -> Expr { )) } -/// Creates aggregation to count all rows, equivalent to `COUNT(*)`, `COUNT()`, `COUNT(1)` +/// Creates aggregation to count all rows. +/// +/// In SQL this is `SELECT COUNT(*) ... ` +/// +/// The expression is equivalent to `COUNT(*)`, `COUNT()`, `COUNT(1)`, and is +/// aliased to a column named `"count(*)"` for backward compatibility. +/// +/// Example +/// ``` +/// # use datafusion_functions_aggregate::count::count_all; +/// # use datafusion_expr::col; +/// // create `count(*)` expression +/// let expr = count_all(); +/// assert_eq!(expr.schema_name().to_string(), "count(*)"); +/// // if you need to refer to this column, use the `schema_name` function +/// let expr = col(expr.schema_name().to_string()); +/// ``` pub fn count_all() -> Expr { - count(Expr::Literal(COUNT_STAR_EXPANSION)) + count(Expr::Literal(COUNT_STAR_EXPANSION)).alias("count(*)") +} + +/// Creates window aggregation to count all rows. +/// +/// In SQL this is `SELECT COUNT(*) OVER (..) ... ` +/// +/// The expression is equivalent to `COUNT(*)`, `COUNT()`, `COUNT(1)` +/// +/// Example +/// ``` +/// # use datafusion_functions_aggregate::count::count_all_window; +/// # use datafusion_expr::col; +/// // create `count(*)` OVER ... window function expression +/// let expr = count_all_window(); +/// assert_eq!( +/// expr.schema_name().to_string(), +/// "count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING" +/// ); +/// // if you need to refer to this column, use the `schema_name` function +/// let expr = col(expr.schema_name().to_string()); +/// ``` +pub fn count_all_window() -> Expr { + Expr::WindowFunction(WindowFunction::new( + WindowFunctionDefinition::AggregateUDF(count_udaf()), + vec![Expr::Literal(COUNT_STAR_EXPANSION)], + )) } #[user_doc( diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index 90fb46883de6..83356e2f9fb4 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -573,7 +573,7 @@ fn min_batch(values: &ArrayRef) -> Result { } /// dynamically-typed max(array) -> ScalarValue -fn max_batch(values: &ArrayRef) -> Result { +pub fn max_batch(values: &ArrayRef) -> Result { Ok(match values.data_type() { DataType::Utf8 => { typed_min_max_batch_string!(values, StringArray, Utf8, max_string) diff --git a/datafusion/functions-aggregate/src/planner.rs b/datafusion/functions-aggregate/src/planner.rs index 7c88a8b82624..c8cb84118995 100644 --- a/datafusion/functions-aggregate/src/planner.rs +++ b/datafusion/functions-aggregate/src/planner.rs @@ -82,6 +82,8 @@ impl ExprPlanner for AggregateFunctionPlanner { // handle count() and count(*) case // convert to count(1) as "count()" // or count(1) as "count(*)" + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] if raw_expr.func.name() == "count" && (raw_expr.args.len() == 1 && matches!(raw_expr.args[0], Expr::Wildcard { .. }) diff --git a/datafusion/functions-nested/benches/map.rs b/datafusion/functions-nested/benches/map.rs index 3726cac0752e..2774b24b902a 100644 --- a/datafusion/functions-nested/benches/map.rs +++ b/datafusion/functions-nested/benches/map.rs @@ -28,7 +28,7 @@ use std::sync::Arc; use datafusion_common::ScalarValue; use datafusion_expr::planner::ExprPlanner; -use datafusion_expr::{ColumnarValue, Expr}; +use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionArgs}; use datafusion_functions_nested::map::map_udf; use datafusion_functions_nested::planner::NestedFunctionPlanner; @@ -94,11 +94,18 @@ fn criterion_benchmark(c: &mut Criterion) { let keys = ColumnarValue::Scalar(ScalarValue::List(Arc::new(key_list))); let values = ColumnarValue::Scalar(ScalarValue::List(Arc::new(value_list))); + let return_type = &map_udf() + .return_type(&[DataType::Utf8, DataType::Int32]) + .expect("should get return type"); + b.iter(|| { black_box( - // TODO use invoke_with_args map_udf() - .invoke_batch(&[keys.clone(), values.clone()], 1) + .invoke_with_args(ScalarFunctionArgs { + args: vec![keys.clone(), values.clone()], + number_rows: 1, + return_type, + }) .expect("map should work on valid values"), ); }); diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs index 422b1b612850..321dda55ce09 100644 --- a/datafusion/functions-nested/src/extract.rs +++ b/datafusion/functions-nested/src/extract.rs @@ -166,6 +166,7 @@ impl ScalarUDFImpl for ArrayElement { List(field) | LargeList(field) | FixedSizeList(field, _) => Ok(field.data_type().clone()), + DataType::Null => Ok(List(Arc::new(Field::new_list_field(DataType::Int64, true)))), _ => plan_err!( "ArrayElement can only accept List, LargeList or FixedSizeList as the first argument" ), @@ -1000,9 +1001,9 @@ where mod tests { use super::array_element_udf; use arrow::datatypes::{DataType, Field}; - use datafusion_common::{Column, DFSchema, ScalarValue}; + use datafusion_common::{Column, DFSchema}; use datafusion_expr::expr::ScalarFunction; - use datafusion_expr::{cast, Expr, ExprSchemable}; + use datafusion_expr::{Expr, ExprSchemable}; use std::collections::HashMap; // Regression test for https://github.com/apache/datafusion/issues/13755 @@ -1036,34 +1037,6 @@ mod tests { fixed_size_list_type ); - // ScalarUDFImpl::return_type_from_exprs with typed exprs - assert_eq!( - udf.return_type_from_exprs( - &[ - cast(Expr::Literal(ScalarValue::Null), array_type.clone()), - cast(Expr::Literal(ScalarValue::Null), index_type.clone()), - ], - &schema, - &[array_type.clone(), index_type.clone()] - ) - .unwrap(), - fixed_size_list_type - ); - - // ScalarUDFImpl::return_type_from_exprs with exprs not carrying type - assert_eq!( - udf.return_type_from_exprs( - &[ - Expr::Column(Column::new_unqualified("my_array")), - Expr::Column(Column::new_unqualified("my_index")), - ], - &schema, - &[array_type.clone(), index_type.clone()] - ) - .unwrap(), - fixed_size_list_type - ); - // Via ExprSchemable::get_type (e.g. SimplifyInfo) let udf_expr = Expr::ScalarFunction(ScalarFunction { func: array_element_udf(), diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs index 41ebb4366cff..446cd58865c3 100644 --- a/datafusion/functions-nested/src/lib.rs +++ b/datafusion/functions-nested/src/lib.rs @@ -52,6 +52,7 @@ pub mod map; pub mod map_extract; pub mod map_keys; pub mod map_values; +pub mod max; pub mod planner; pub mod position; pub mod range; @@ -144,6 +145,7 @@ pub fn all_default_nested_functions() -> Vec> { length::array_length_udf(), distance::array_distance_udf(), flatten::flatten_udf(), + max::array_max_udf(), sort::array_sort_udf(), repeat::array_repeat_udf(), resize::array_resize_udf(), diff --git a/datafusion/functions-nested/src/max.rs b/datafusion/functions-nested/src/max.rs new file mode 100644 index 000000000000..22bd14740b5e --- /dev/null +++ b/datafusion/functions-nested/src/max.rs @@ -0,0 +1,137 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_max function. +use crate::utils::make_scalar_function; +use arrow::array::ArrayRef; +use arrow::datatypes::DataType; +use arrow::datatypes::DataType::List; +use datafusion_common::cast::as_list_array; +use datafusion_common::utils::take_function_args; +use datafusion_common::{exec_err, ScalarValue}; +use datafusion_doc::Documentation; +use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility}; +use datafusion_functions_aggregate::min_max; +use datafusion_macros::user_doc; +use itertools::Itertools; +use std::any::Any; + +make_udf_expr_and_func!( + ArrayMax, + array_max, + array, + "returns the maximum value in the array.", + array_max_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the maximum value in the array.", + syntax_example = "array_max(array)", + sql_example = r#"```sql +> select array_max([3,1,4,2]); ++-----------------------------------------+ +| array_max(List([3,1,4,2])) | ++-----------------------------------------+ +| 4 | ++-----------------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug)] +pub struct ArrayMax { + signature: Signature, + aliases: Vec, +} + +impl Default for ArrayMax { + fn default() -> Self { + Self::new() + } +} + +impl ArrayMax { + pub fn new() -> Self { + Self { + signature: Signature::array(Volatility::Immutable), + aliases: vec!["list_max".to_string()], + } + } +} + +impl ScalarUDFImpl for ArrayMax { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "array_max" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result { + match &arg_types[0] { + List(field) => Ok(field.data_type().clone()), + _ => exec_err!("Not reachable, data_type should be List"), + } + } + + fn invoke_batch( + &self, + args: &[ColumnarValue], + _number_rows: usize, + ) -> datafusion_common::Result { + make_scalar_function(array_max_inner)(args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +/// array_max SQL function +/// +/// There is one argument for array_max as the array. +/// `array_max(array)` +/// +/// For example: +/// > array_max(\[1, 3, 2]) -> 3 +pub fn array_max_inner(args: &[ArrayRef]) -> datafusion_common::Result { + let [arg1] = take_function_args("array_max", args)?; + + match arg1.data_type() { + List(_) => { + let input_list_array = as_list_array(&arg1)?; + let result_vec = input_list_array + .iter() + .flat_map(|arr| min_max::max_batch(&arr.unwrap())) + .collect_vec(); + ScalarValue::iter_to_array(result_vec) + } + _ => exec_err!("array_max does not support type: {:?}", arg1.data_type()), + } +} diff --git a/datafusion/functions-nested/src/replace.rs b/datafusion/functions-nested/src/replace.rs index 71bfedb72d1c..3dbe672c5b02 100644 --- a/datafusion/functions-nested/src/replace.rs +++ b/datafusion/functions-nested/src/replace.rs @@ -18,8 +18,8 @@ //! [`ScalarUDFImpl`] definitions for array_replace, array_replace_n and array_replace_all functions. use arrow::array::{ - Array, ArrayRef, AsArray, Capacities, GenericListArray, MutableArrayData, - NullBufferBuilder, OffsetSizeTrait, + new_null_array, Array, ArrayRef, AsArray, Capacities, GenericListArray, + MutableArrayData, NullBufferBuilder, OffsetSizeTrait, }; use arrow::datatypes::{DataType, Field}; @@ -429,6 +429,7 @@ pub(crate) fn array_replace_inner(args: &[ArrayRef]) -> Result { let list_array = array.as_list::(); general_replace::(list_array, from, to, arr_n) } + DataType::Null => Ok(new_null_array(array.data_type(), 1)), array_type => exec_err!("array_replace does not support type '{array_type:?}'."), } } @@ -447,6 +448,7 @@ pub(crate) fn array_replace_n_inner(args: &[ArrayRef]) -> Result { let list_array = array.as_list::(); general_replace::(list_array, from, to, arr_n) } + DataType::Null => Ok(new_null_array(array.data_type(), 1)), array_type => { exec_err!("array_replace_n does not support type '{array_type:?}'.") } @@ -467,6 +469,7 @@ pub(crate) fn array_replace_all_inner(args: &[ArrayRef]) -> Result { let list_array = array.as_list::(); general_replace::(list_array, from, to, arr_n) } + DataType::Null => Ok(new_null_array(array.data_type(), 1)), array_type => { exec_err!("array_replace_all does not support type '{array_type:?}'.") } diff --git a/datafusion/functions-nested/src/resize.rs b/datafusion/functions-nested/src/resize.rs index 6c0b91a678e7..145d7e80043b 100644 --- a/datafusion/functions-nested/src/resize.rs +++ b/datafusion/functions-nested/src/resize.rs @@ -23,16 +23,18 @@ use arrow::array::{ MutableArrayData, NullBufferBuilder, OffsetSizeTrait, }; use arrow::buffer::OffsetBuffer; -use arrow::datatypes::ArrowNativeType; use arrow::datatypes::DataType; +use arrow::datatypes::{ArrowNativeType, Field}; use arrow::datatypes::{ DataType::{FixedSizeList, LargeList, List}, FieldRef, }; use datafusion_common::cast::{as_int64_array, as_large_list_array, as_list_array}; +use datafusion_common::utils::ListCoercion; use datafusion_common::{exec_err, internal_datafusion_err, Result, ScalarValue}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, + ScalarUDFImpl, Signature, TypeSignature, Volatility, }; use datafusion_macros::user_doc; use std::any::Any; @@ -83,7 +85,26 @@ impl Default for ArrayResize { impl ArrayResize { pub fn new() -> Self { Self { - signature: Signature::variadic_any(Volatility::Immutable), + signature: Signature::one_of( + vec![ + TypeSignature::ArraySignature(ArrayFunctionSignature::Array { + arguments: vec![ + ArrayFunctionArgument::Array, + ArrayFunctionArgument::Index, + ], + array_coercion: Some(ListCoercion::FixedSizedListToList), + }), + TypeSignature::ArraySignature(ArrayFunctionSignature::Array { + arguments: vec![ + ArrayFunctionArgument::Array, + ArrayFunctionArgument::Index, + ArrayFunctionArgument::Element, + ], + array_coercion: Some(ListCoercion::FixedSizedListToList), + }), + ], + Volatility::Immutable, + ), aliases: vec!["list_resize".to_string()], } } @@ -106,6 +127,9 @@ impl ScalarUDFImpl for ArrayResize { match &arg_types[0] { List(field) | FixedSizeList(field, _) => Ok(List(Arc::clone(field))), LargeList(field) => Ok(LargeList(Arc::clone(field))), + DataType::Null => { + Ok(List(Arc::new(Field::new_list_field(DataType::Int64, true)))) + } _ => exec_err!( "Not reachable, data_type should be List, LargeList or FixedSizeList" ), @@ -137,7 +161,7 @@ pub(crate) fn array_resize_inner(arg: &[ArrayRef]) -> Result { let array = &arg[0]; // Checks if entire array is null - if array.null_count() == array.len() { + if array.logical_null_count() == array.len() { let return_type = match array.data_type() { List(field) => List(Arc::clone(field)), LargeList(field) => LargeList(Arc::clone(field)), diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs index 7dbf9f2b211e..1db245fe52fe 100644 --- a/datafusion/functions-nested/src/sort.rs +++ b/datafusion/functions-nested/src/sort.rs @@ -18,7 +18,7 @@ //! [`ScalarUDFImpl`] definitions for array_sort function. use crate::utils::make_scalar_function; -use arrow::array::{Array, ArrayRef, ListArray, NullBufferBuilder}; +use arrow::array::{new_null_array, Array, ArrayRef, ListArray, NullBufferBuilder}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::DataType::{FixedSizeList, LargeList, List}; use arrow::datatypes::{DataType, Field}; @@ -26,7 +26,8 @@ use arrow::{compute, compute::SortOptions}; use datafusion_common::cast::{as_list_array, as_string_array}; use datafusion_common::{exec_err, Result}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, + ScalarUDFImpl, Signature, TypeSignature, Volatility, }; use datafusion_macros::user_doc; use std::any::Any; @@ -87,7 +88,30 @@ impl Default for ArraySort { impl ArraySort { pub fn new() -> Self { Self { - signature: Signature::variadic_any(Volatility::Immutable), + signature: Signature::one_of( + vec![ + TypeSignature::ArraySignature(ArrayFunctionSignature::Array { + arguments: vec![ArrayFunctionArgument::Array], + array_coercion: None, + }), + TypeSignature::ArraySignature(ArrayFunctionSignature::Array { + arguments: vec![ + ArrayFunctionArgument::Array, + ArrayFunctionArgument::String, + ], + array_coercion: None, + }), + TypeSignature::ArraySignature(ArrayFunctionSignature::Array { + arguments: vec![ + ArrayFunctionArgument::Array, + ArrayFunctionArgument::String, + ArrayFunctionArgument::String, + ], + array_coercion: None, + }), + ], + Volatility::Immutable, + ), aliases: vec!["list_sort".to_string()], } } @@ -115,6 +139,7 @@ impl ScalarUDFImpl for ArraySort { field.data_type().clone(), true, )))), + DataType::Null => Ok(DataType::Null), _ => exec_err!( "Not reachable, data_type should be List, LargeList or FixedSizeList" ), @@ -143,6 +168,10 @@ pub fn array_sort_inner(args: &[ArrayRef]) -> Result { return exec_err!("array_sort expects one to three arguments"); } + if args[1..].iter().any(|array| array.is_null(0)) { + return Ok(new_null_array(args[0].data_type(), args[0].len())); + } + let sort_option = match args.len() { 1 => None, 2 => { @@ -196,12 +225,16 @@ pub fn array_sort_inner(args: &[ArrayRef]) -> Result { .map(|a| a.as_ref()) .collect::>(); - let list_arr = ListArray::new( - Arc::new(Field::new_list_field(data_type, true)), - OffsetBuffer::from_lengths(array_lengths), - Arc::new(compute::concat(elements.as_slice())?), - buffer, - ); + let list_arr = if elements.is_empty() { + ListArray::new_null(Arc::new(Field::new_list_field(data_type, true)), row_count) + } else { + ListArray::new( + Arc::new(Field::new_list_field(data_type, true)), + OffsetBuffer::from_lengths(array_lengths), + Arc::new(compute::concat(elements.as_slice())?), + buffer, + ) + }; Ok(Arc::new(list_arr)) } diff --git a/datafusion/functions-nested/src/string.rs b/datafusion/functions-nested/src/string.rs index 99af3e95c804..d60d1a6e4de0 100644 --- a/datafusion/functions-nested/src/string.rs +++ b/datafusion/functions-nested/src/string.rs @@ -24,7 +24,6 @@ use arrow::array::{ UInt8Array, }; use arrow::datatypes::{DataType, Field}; -use datafusion_expr::TypeSignature; use datafusion_common::{ internal_datafusion_err, not_impl_err, plan_err, DataFusionError, Result, @@ -44,8 +43,10 @@ use arrow::datatypes::DataType::{ }; use datafusion_common::cast::{as_large_list_array, as_list_array}; use datafusion_common::exec_err; +use datafusion_common::types::logical_string; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature, + TypeSignatureClass, Volatility, }; use datafusion_functions::{downcast_arg, downcast_named_arg}; use datafusion_macros::user_doc; @@ -251,7 +252,17 @@ impl StringToArray { pub fn new() -> Self { Self { signature: Signature::one_of( - vec![TypeSignature::String(2), TypeSignature::String(3)], + vec![ + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ]), + ], Volatility::Immutable, ), aliases: vec![String::from("string_to_list")], diff --git a/datafusion/functions-window/src/planner.rs b/datafusion/functions-window/src/planner.rs index ffaccd9369bc..1ddd8b27c420 100644 --- a/datafusion/functions-window/src/planner.rs +++ b/datafusion/functions-window/src/planner.rs @@ -79,6 +79,8 @@ impl ExprPlanner for WindowFunctionPlanner { null_treatment, }; + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] if raw_expr.func_def.name() == "count" && (raw_expr.args.len() == 1 && matches!(raw_expr.args[0], Expr::Wildcard { .. }) diff --git a/datafusion/functions/benches/character_length.rs b/datafusion/functions/benches/character_length.rs index 3655d8409807..bbcfed021064 100644 --- a/datafusion/functions/benches/character_length.rs +++ b/datafusion/functions/benches/character_length.rs @@ -17,7 +17,9 @@ extern crate criterion; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use datafusion_expr::ScalarFunctionArgs; use helper::gen_string_array; mod helper; @@ -26,6 +28,8 @@ fn criterion_benchmark(c: &mut Criterion) { // All benches are single batch run with 8192 rows let character_length = datafusion_functions::unicode::character_length(); + let return_type = DataType::Utf8; + let n_rows = 8192; for str_len in [8, 32, 128, 4096] { // StringArray ASCII only @@ -34,8 +38,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("character_length_StringArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(character_length.invoke_batch(&args_string_ascii, n_rows)) + black_box(character_length.invoke_with_args(ScalarFunctionArgs { + args: args_string_ascii.clone(), + number_rows: n_rows, + return_type: &return_type, + })) }) }, ); @@ -46,8 +53,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("character_length_StringArray_utf8_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(character_length.invoke_batch(&args_string_utf8, n_rows)) + black_box(character_length.invoke_with_args(ScalarFunctionArgs { + args: args_string_utf8.clone(), + number_rows: n_rows, + return_type: &return_type, + })) }) }, ); @@ -58,10 +68,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("character_length_StringViewArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box( - character_length.invoke_batch(&args_string_view_ascii, n_rows), - ) + black_box(character_length.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_ascii.clone(), + number_rows: n_rows, + return_type: &return_type, + })) }) }, ); @@ -72,10 +83,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("character_length_StringViewArray_utf8_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box( - character_length.invoke_batch(&args_string_view_utf8, n_rows), - ) + black_box(character_length.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_utf8.clone(), + number_rows: n_rows, + return_type: &return_type, + })) }) }, ); diff --git a/datafusion/functions/benches/chr.rs b/datafusion/functions/benches/chr.rs index 58c5ee3d68f6..4750fb466653 100644 --- a/datafusion/functions/benches/chr.rs +++ b/datafusion/functions/benches/chr.rs @@ -19,10 +19,11 @@ extern crate criterion; use arrow::{array::PrimitiveArray, datatypes::Int64Type, util::test_util::seedable_rng}; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::string::chr; use rand::Rng; +use arrow::datatypes::DataType; use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { @@ -44,7 +45,17 @@ fn criterion_benchmark(c: &mut Criterion) { let input = Arc::new(input); let args = vec![ColumnarValue::Array(input)]; c.bench_function("chr", |b| { - b.iter(|| black_box(cot_fn.invoke_batch(&args, size).unwrap())) + b.iter(|| { + black_box( + cot_fn + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) + }) }); } diff --git a/datafusion/functions/benches/cot.rs b/datafusion/functions/benches/cot.rs index bb0585a2de9b..b2a9ca0b9f47 100644 --- a/datafusion/functions/benches/cot.rs +++ b/datafusion/functions/benches/cot.rs @@ -22,9 +22,10 @@ use arrow::{ util::bench_util::create_primitive_array, }; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::cot; +use arrow::datatypes::DataType; use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { @@ -34,16 +35,30 @@ fn criterion_benchmark(c: &mut Criterion) { let f32_args = vec![ColumnarValue::Array(f32_array)]; c.bench_function(&format!("cot f32 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(cot_fn.invoke_batch(&f32_args, size).unwrap()) + black_box( + cot_fn + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + number_rows: size, + return_type: &DataType::Float32, + }) + .unwrap(), + ) }) }); let f64_array = Arc::new(create_primitive_array::(size, 0.2)); let f64_args = vec![ColumnarValue::Array(f64_array)]; c.bench_function(&format!("cot f64 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(cot_fn.invoke_batch(&f64_args, size).unwrap()) + black_box( + cot_fn + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + number_rows: size, + return_type: &DataType::Float64, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/date_bin.rs b/datafusion/functions/benches/date_bin.rs index aa7c7710617d..7ea5fdcb2be2 100644 --- a/datafusion/functions/benches/date_bin.rs +++ b/datafusion/functions/benches/date_bin.rs @@ -25,7 +25,7 @@ use datafusion_common::ScalarValue; use rand::rngs::ThreadRng; use rand::Rng; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::date_bin; fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { @@ -45,12 +45,18 @@ fn criterion_benchmark(c: &mut Criterion) { let interval = ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1_000_000)); let timestamps = ColumnarValue::Array(timestamps_array); let udf = date_bin(); + let return_type = udf + .return_type(&[interval.data_type(), timestamps.data_type()]) + .unwrap(); b.iter(|| { - // TODO use invoke_with_args black_box( - udf.invoke_batch(&[interval.clone(), timestamps.clone()], batch_len) - .expect("date_bin should work on valid values"), + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![interval.clone(), timestamps.clone()], + number_rows: batch_len, + return_type: &return_type, + }) + .expect("date_bin should work on valid values"), ) }) }); diff --git a/datafusion/functions/benches/date_trunc.rs b/datafusion/functions/benches/date_trunc.rs index d420b8f6ac70..e7e96fb7a9fa 100644 --- a/datafusion/functions/benches/date_trunc.rs +++ b/datafusion/functions/benches/date_trunc.rs @@ -25,7 +25,7 @@ use datafusion_common::ScalarValue; use rand::rngs::ThreadRng; use rand::Rng; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::date_trunc; fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { @@ -46,11 +46,18 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Scalar(ScalarValue::Utf8(Some("minute".to_string()))); let timestamps = ColumnarValue::Array(timestamps_array); let udf = date_trunc(); - + let args = vec![precision, timestamps]; + let return_type = &udf + .return_type(&args.iter().map(|arg| arg.data_type()).collect::>()) + .unwrap(); b.iter(|| { black_box( - udf.invoke_batch(&[precision.clone(), timestamps.clone()], batch_len) - .expect("date_trunc should work on valid values"), + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: batch_len, + return_type, + }) + .expect("date_trunc should work on valid values"), ) }) }); diff --git a/datafusion/functions/benches/encoding.rs b/datafusion/functions/benches/encoding.rs index e37842a62b4a..cf8f8d2fd62c 100644 --- a/datafusion/functions/benches/encoding.rs +++ b/datafusion/functions/benches/encoding.rs @@ -17,9 +17,10 @@ extern crate criterion; +use arrow::datatypes::DataType; use arrow::util::bench_util::create_string_array_with_len; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::encoding; use std::sync::Arc; @@ -29,35 +30,49 @@ fn criterion_benchmark(c: &mut Criterion) { let str_array = Arc::new(create_string_array_with_len::(size, 0.2, 32)); c.bench_function(&format!("base64_decode/{size}"), |b| { let method = ColumnarValue::Scalar("base64".into()); - // TODO: use invoke_with_args let encoded = encoding::encode() - .invoke_batch( - &[ColumnarValue::Array(str_array.clone()), method.clone()], - size, - ) + .invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(str_array.clone()), method.clone()], + number_rows: size, + return_type: &DataType::Utf8, + }) .unwrap(); let args = vec![encoded, method]; b.iter(|| { - // TODO use invoke_with_args - black_box(decode.invoke_batch(&args, size).unwrap()) + black_box( + decode + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); c.bench_function(&format!("hex_decode/{size}"), |b| { let method = ColumnarValue::Scalar("hex".into()); - // TODO use invoke_with_args let encoded = encoding::encode() - .invoke_batch( - &[ColumnarValue::Array(str_array.clone()), method.clone()], - size, - ) + .invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(str_array.clone()), method.clone()], + number_rows: size, + return_type: &DataType::Utf8, + }) .unwrap(); let args = vec![encoded, method]; b.iter(|| { - // TODO use invoke_with_args - black_box(decode.invoke_batch(&args, size).unwrap()) + black_box( + decode + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/isnan.rs b/datafusion/functions/benches/isnan.rs index 605a520715f4..42004cc24f69 100644 --- a/datafusion/functions/benches/isnan.rs +++ b/datafusion/functions/benches/isnan.rs @@ -17,12 +17,13 @@ extern crate criterion; +use arrow::datatypes::DataType; use arrow::{ datatypes::{Float32Type, Float64Type}, util::bench_util::create_primitive_array, }; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::isnan; use std::sync::Arc; @@ -33,16 +34,30 @@ fn criterion_benchmark(c: &mut Criterion) { let f32_args = vec![ColumnarValue::Array(f32_array)]; c.bench_function(&format!("isnan f32 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(isnan.invoke_batch(&f32_args, size).unwrap()) + black_box( + isnan + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + number_rows: size, + return_type: &DataType::Boolean, + }) + .unwrap(), + ) }) }); let f64_array = Arc::new(create_primitive_array::(size, 0.2)); let f64_args = vec![ColumnarValue::Array(f64_array)]; c.bench_function(&format!("isnan f64 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(isnan.invoke_batch(&f64_args, size).unwrap()) + black_box( + isnan + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + number_rows: size, + return_type: &DataType::Boolean, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/iszero.rs b/datafusion/functions/benches/iszero.rs index 48fb6fbed9c3..9e5f6a84804b 100644 --- a/datafusion/functions/benches/iszero.rs +++ b/datafusion/functions/benches/iszero.rs @@ -17,12 +17,13 @@ extern crate criterion; +use arrow::datatypes::DataType; use arrow::{ datatypes::{Float32Type, Float64Type}, util::bench_util::create_primitive_array, }; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::iszero; use std::sync::Arc; @@ -34,8 +35,15 @@ fn criterion_benchmark(c: &mut Criterion) { let f32_args = vec![ColumnarValue::Array(f32_array)]; c.bench_function(&format!("iszero f32 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(iszero.invoke_batch(&f32_args, batch_len).unwrap()) + black_box( + iszero + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + number_rows: batch_len, + return_type: &DataType::Boolean, + }) + .unwrap(), + ) }) }); let f64_array = Arc::new(create_primitive_array::(size, 0.2)); @@ -43,8 +51,15 @@ fn criterion_benchmark(c: &mut Criterion) { let f64_args = vec![ColumnarValue::Array(f64_array)]; c.bench_function(&format!("iszero f64 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(iszero.invoke_batch(&f64_args, batch_len).unwrap()) + black_box( + iszero + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + number_rows: batch_len, + return_type: &DataType::Boolean, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/make_date.rs b/datafusion/functions/benches/make_date.rs index d9309bcd3db2..8dd7a7a59773 100644 --- a/datafusion/functions/benches/make_date.rs +++ b/datafusion/functions/benches/make_date.rs @@ -20,12 +20,13 @@ extern crate criterion; use std::sync::Arc; use arrow::array::{Array, ArrayRef, Int32Array}; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use rand::rngs::ThreadRng; use rand::Rng; use datafusion_common::ScalarValue; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::make_date; fn years(rng: &mut ThreadRng) -> Int32Array { @@ -64,13 +65,13 @@ fn criterion_benchmark(c: &mut Criterion) { let days = ColumnarValue::Array(Arc::new(days(&mut rng)) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( make_date() - .invoke_batch( - &[years.clone(), months.clone(), days.clone()], - batch_len, - ) + .invoke_with_args(ScalarFunctionArgs { + args: vec![years.clone(), months.clone(), days.clone()], + number_rows: batch_len, + return_type: &DataType::Date32, + }) .expect("make_date should work on valid values"), ) }) @@ -85,13 +86,13 @@ fn criterion_benchmark(c: &mut Criterion) { let days = ColumnarValue::Array(Arc::new(days(&mut rng)) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( make_date() - .invoke_batch( - &[year.clone(), months.clone(), days.clone()], - batch_len, - ) + .invoke_with_args(ScalarFunctionArgs { + args: vec![year.clone(), months.clone(), days.clone()], + number_rows: batch_len, + return_type: &DataType::Date32, + }) .expect("make_date should work on valid values"), ) }) @@ -106,10 +107,13 @@ fn criterion_benchmark(c: &mut Criterion) { let days = ColumnarValue::Array(day_arr); b.iter(|| { - // TODO use invoke_with_args black_box( make_date() - .invoke_batch(&[year.clone(), month.clone(), days.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![year.clone(), month.clone(), days.clone()], + number_rows: batch_len, + return_type: &DataType::Date32, + }) .expect("make_date should work on valid values"), ) }) @@ -121,10 +125,13 @@ fn criterion_benchmark(c: &mut Criterion) { let day = ColumnarValue::Scalar(ScalarValue::Int32(Some(26))); b.iter(|| { - // TODO use invoke_with_args black_box( make_date() - .invoke_batch(&[year.clone(), month.clone(), day.clone()], 1) + .invoke_with_args(ScalarFunctionArgs { + args: vec![year.clone(), month.clone(), day.clone()], + number_rows: 1, + return_type: &DataType::Date32, + }) .expect("make_date should work on valid values"), ) }) diff --git a/datafusion/functions/benches/nullif.rs b/datafusion/functions/benches/nullif.rs index e29fd03aa819..9096c976bf31 100644 --- a/datafusion/functions/benches/nullif.rs +++ b/datafusion/functions/benches/nullif.rs @@ -17,10 +17,11 @@ extern crate criterion; +use arrow::datatypes::DataType; use arrow::util::bench_util::create_string_array_with_len; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use datafusion_common::ScalarValue; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::core::nullif; use std::sync::Arc; @@ -34,8 +35,15 @@ fn criterion_benchmark(c: &mut Criterion) { ]; c.bench_function(&format!("nullif scalar array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(nullif.invoke_batch(&args, size).unwrap()) + black_box( + nullif + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/pad.rs b/datafusion/functions/benches/pad.rs index 6f267b350a35..f78a53fbee19 100644 --- a/datafusion/functions/benches/pad.rs +++ b/datafusion/functions/benches/pad.rs @@ -16,12 +16,12 @@ // under the License. use arrow::array::{ArrayRef, ArrowPrimitiveType, OffsetSizeTrait, PrimitiveArray}; -use arrow::datatypes::Int64Type; +use arrow::datatypes::{DataType, Int64Type}; use arrow::util::bench_util::{ create_string_array_with_len, create_string_view_array_with_len, }; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::unicode::{lpad, rpad}; use rand::distributions::{Distribution, Uniform}; use rand::Rng; @@ -102,24 +102,45 @@ fn criterion_benchmark(c: &mut Criterion) { let args = create_args::(size, 32, false); group.bench_function(BenchmarkId::new("utf8 type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(lpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + lpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); let args = create_args::(size, 32, false); group.bench_function(BenchmarkId::new("largeutf8 type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(lpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + lpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::LargeUtf8, + }) + .unwrap(), + ) }) }); let args = create_args::(size, 32, true); group.bench_function(BenchmarkId::new("stringview type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(lpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + lpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); @@ -130,16 +151,30 @@ fn criterion_benchmark(c: &mut Criterion) { let args = create_args::(size, 32, false); group.bench_function(BenchmarkId::new("utf8 type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(rpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + rpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); let args = create_args::(size, 32, false); group.bench_function(BenchmarkId::new("largeutf8 type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(rpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + rpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::LargeUtf8, + }) + .unwrap(), + ) }) }); @@ -147,8 +182,15 @@ fn criterion_benchmark(c: &mut Criterion) { let args = create_args::(size, 32, true); group.bench_function(BenchmarkId::new("stringview type", size), |b| { b.iter(|| { - // TODO use invoke_with_args - criterion::black_box(rpad().invoke_batch(&args, size).unwrap()) + criterion::black_box( + rpad() + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8, + }) + .unwrap(), + ) }) }); diff --git a/datafusion/functions/benches/random.rs b/datafusion/functions/benches/random.rs index bc20e0ff11c1..78ebf23e02e0 100644 --- a/datafusion/functions/benches/random.rs +++ b/datafusion/functions/benches/random.rs @@ -17,8 +17,9 @@ extern crate criterion; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ScalarUDFImpl; +use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl}; use datafusion_functions::math::random::RandomFunc; fn criterion_benchmark(c: &mut Criterion) { @@ -29,8 +30,15 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function("random_1M_rows_batch_8192", |b| { b.iter(|| { for _ in 0..iterations { - #[allow(deprecated)] // TODO: migrate to invoke_with_args - black_box(random_func.invoke_batch(&[], 8192).unwrap()); + black_box( + random_func + .invoke_with_args(ScalarFunctionArgs { + args: vec![], + number_rows: 8192, + return_type: &DataType::Float64, + }) + .unwrap(), + ); } }) }); @@ -40,8 +48,15 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function("random_1M_rows_batch_128", |b| { b.iter(|| { for _ in 0..iterations_128 { - #[allow(deprecated)] // TODO: migrate to invoke_with_args - black_box(random_func.invoke_batch(&[], 128).unwrap()); + black_box( + random_func + .invoke_with_args(ScalarFunctionArgs { + args: vec![], + number_rows: 128, + return_type: &DataType::Float64, + }) + .unwrap(), + ); } }) }); diff --git a/datafusion/functions/benches/reverse.rs b/datafusion/functions/benches/reverse.rs index 889ca59e2a14..d61f8fb80517 100644 --- a/datafusion/functions/benches/reverse.rs +++ b/datafusion/functions/benches/reverse.rs @@ -18,7 +18,9 @@ extern crate criterion; mod helper; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use datafusion_expr::ScalarFunctionArgs; use helper::gen_string_array; fn criterion_benchmark(c: &mut Criterion) { @@ -42,8 +44,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("reverse_StringArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(reverse.invoke_batch(&args_string_ascii, N_ROWS)) + black_box(reverse.invoke_with_args(ScalarFunctionArgs { + args: args_string_ascii.clone(), + number_rows: N_ROWS, + return_type: &DataType::Utf8, + })) }) }, ); @@ -58,8 +63,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(reverse.invoke_batch(&args_string_utf8, N_ROWS)) + black_box(reverse.invoke_with_args(ScalarFunctionArgs { + args: args_string_utf8.clone(), + number_rows: N_ROWS, + return_type: &DataType::Utf8, + })) }) }, ); @@ -76,8 +84,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("reverse_StringViewArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(reverse.invoke_batch(&args_string_view_ascii, N_ROWS)) + black_box(reverse.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_ascii.clone(), + number_rows: N_ROWS, + return_type: &DataType::Utf8, + })) }) }, ); @@ -92,8 +103,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(reverse.invoke_batch(&args_string_view_utf8, N_ROWS)) + black_box(reverse.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_utf8.clone(), + number_rows: N_ROWS, + return_type: &DataType::Utf8, + })) }) }, ); diff --git a/datafusion/functions/benches/signum.rs b/datafusion/functions/benches/signum.rs index a51b2ebe5ab7..01939fad5f34 100644 --- a/datafusion/functions/benches/signum.rs +++ b/datafusion/functions/benches/signum.rs @@ -17,12 +17,13 @@ extern crate criterion; +use arrow::datatypes::DataType; use arrow::{ datatypes::{Float32Type, Float64Type}, util::bench_util::create_primitive_array, }; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::signum; use std::sync::Arc; @@ -34,8 +35,15 @@ fn criterion_benchmark(c: &mut Criterion) { let f32_args = vec![ColumnarValue::Array(f32_array)]; c.bench_function(&format!("signum f32 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(signum.invoke_batch(&f32_args, batch_len).unwrap()) + black_box( + signum + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + number_rows: batch_len, + return_type: &DataType::Float32, + }) + .unwrap(), + ) }) }); let f64_array = Arc::new(create_primitive_array::(size, 0.2)); @@ -44,8 +52,15 @@ fn criterion_benchmark(c: &mut Criterion) { let f64_args = vec![ColumnarValue::Array(f64_array)]; c.bench_function(&format!("signum f64 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(signum.invoke_batch(&f64_args, batch_len).unwrap()) + black_box( + signum + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + number_rows: batch_len, + return_type: &DataType::Float64, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/benches/strpos.rs b/datafusion/functions/benches/strpos.rs index f4962380dfbf..df57c229e0ad 100644 --- a/datafusion/functions/benches/strpos.rs +++ b/datafusion/functions/benches/strpos.rs @@ -18,8 +18,9 @@ extern crate criterion; use arrow::array::{StringArray, StringViewArray}; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use rand::distributions::Alphanumeric; use rand::prelude::StdRng; use rand::{Rng, SeedableRng}; @@ -114,8 +115,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("strpos_StringArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(strpos.invoke_batch(&args_string_ascii, n_rows)) + black_box(strpos.invoke_with_args(ScalarFunctionArgs { + args: args_string_ascii.clone(), + number_rows: n_rows, + return_type: &DataType::Int32, + })) }) }, ); @@ -126,8 +130,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("strpos_StringArray_utf8_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(strpos.invoke_batch(&args_string_utf8, n_rows)) + black_box(strpos.invoke_with_args(ScalarFunctionArgs { + args: args_string_utf8.clone(), + number_rows: n_rows, + return_type: &DataType::Int32, + })) }) }, ); @@ -138,8 +145,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("strpos_StringViewArray_ascii_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(strpos.invoke_batch(&args_string_view_ascii, n_rows)) + black_box(strpos.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_ascii.clone(), + number_rows: n_rows, + return_type: &DataType::Int32, + })) }) }, ); @@ -150,8 +160,11 @@ fn criterion_benchmark(c: &mut Criterion) { &format!("strpos_StringViewArray_utf8_str_len_{}", str_len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(strpos.invoke_batch(&args_string_view_utf8, n_rows)) + black_box(strpos.invoke_with_args(ScalarFunctionArgs { + args: args_string_view_utf8.clone(), + number_rows: n_rows, + return_type: &DataType::Int32, + })) }) }, ); diff --git a/datafusion/functions/benches/substr.rs b/datafusion/functions/benches/substr.rs index 8b8e8dbc4279..80ab70ef71b0 100644 --- a/datafusion/functions/benches/substr.rs +++ b/datafusion/functions/benches/substr.rs @@ -18,11 +18,12 @@ extern crate criterion; use arrow::array::{ArrayRef, Int64Array, OffsetSizeTrait}; +use arrow::datatypes::DataType; use arrow::util::bench_util::{ create_string_array_with_len, create_string_view_array_with_len, }; use criterion::{black_box, criterion_group, criterion_main, Criterion, SamplingMode}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::unicode; use std::sync::Arc; @@ -109,8 +110,11 @@ fn criterion_benchmark(c: &mut Criterion) { format!("substr_string_view [size={}, strlen={}]", size, len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -120,8 +124,11 @@ fn criterion_benchmark(c: &mut Criterion) { format!("substr_string [size={}, strlen={}]", size, len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -131,8 +138,11 @@ fn criterion_benchmark(c: &mut Criterion) { format!("substr_large_string [size={}, strlen={}]", size, len), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -154,8 +164,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -168,8 +181,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -182,8 +198,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -205,8 +224,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -219,8 +241,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); @@ -233,8 +258,11 @@ fn criterion_benchmark(c: &mut Criterion) { ), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(substr.invoke_batch(&args, size)) + black_box(substr.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: size, + return_type: &DataType::Utf8View, + })) }) }, ); diff --git a/datafusion/functions/benches/substr_index.rs b/datafusion/functions/benches/substr_index.rs index 1ea8e2606f0d..b1c1c3c34a95 100644 --- a/datafusion/functions/benches/substr_index.rs +++ b/datafusion/functions/benches/substr_index.rs @@ -20,12 +20,13 @@ extern crate criterion; use std::sync::Arc; use arrow::array::{ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::DataType; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use rand::distributions::{Alphanumeric, Uniform}; use rand::prelude::Distribution; use rand::Rng; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::unicode::substr_index; struct Filter { @@ -89,12 +90,15 @@ fn criterion_benchmark(c: &mut Criterion) { let delimiters = ColumnarValue::Array(Arc::new(delimiters) as ArrayRef); let counts = ColumnarValue::Array(Arc::new(counts) as ArrayRef); - let args = [strings, delimiters, counts]; + let args = vec![strings, delimiters, counts]; b.iter(|| { - #[allow(deprecated)] // TODO: invoke_with_args black_box( substr_index() - .invoke_batch(&args, batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: batch_len, + return_type: &DataType::Utf8, + }) .expect("substr_index should work on valid values"), ) }) diff --git a/datafusion/functions/benches/to_char.rs b/datafusion/functions/benches/to_char.rs index 72eae45b1e1b..6f20a20dc219 100644 --- a/datafusion/functions/benches/to_char.rs +++ b/datafusion/functions/benches/to_char.rs @@ -20,6 +20,7 @@ extern crate criterion; use std::sync::Arc; use arrow::array::{ArrayRef, Date32Array, StringArray}; +use arrow::datatypes::DataType; use chrono::prelude::*; use chrono::TimeDelta; use criterion::{black_box, criterion_group, criterion_main, Criterion}; @@ -29,7 +30,7 @@ use rand::Rng; use datafusion_common::ScalarValue; use datafusion_common::ScalarValue::TimestampNanosecond; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::to_char; fn random_date_in_range( @@ -88,10 +89,13 @@ fn criterion_benchmark(c: &mut Criterion) { let patterns = ColumnarValue::Array(Arc::new(patterns(&mut rng)) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( to_char() - .invoke_batch(&[data.clone(), patterns.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![data.clone(), patterns.clone()], + number_rows: batch_len, + return_type: &DataType::Utf8, + }) .expect("to_char should work on valid values"), ) }) @@ -106,10 +110,13 @@ fn criterion_benchmark(c: &mut Criterion) { ColumnarValue::Scalar(ScalarValue::Utf8(Some("%Y-%m-%d".to_string()))); b.iter(|| { - // TODO use invoke_with_args black_box( to_char() - .invoke_batch(&[data.clone(), patterns.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![data.clone(), patterns.clone()], + number_rows: batch_len, + return_type: &DataType::Utf8, + }) .expect("to_char should work on valid values"), ) }) @@ -130,10 +137,13 @@ fn criterion_benchmark(c: &mut Criterion) { ))); b.iter(|| { - // TODO use invoke_with_args black_box( to_char() - .invoke_batch(&[data.clone(), pattern.clone()], 1) + .invoke_with_args(ScalarFunctionArgs { + args: vec![data.clone(), pattern.clone()], + number_rows: 1, + return_type: &DataType::Utf8, + }) .expect("to_char should work on valid values"), ) }) diff --git a/datafusion/functions/benches/to_timestamp.rs b/datafusion/functions/benches/to_timestamp.rs index 9f5f6661f998..aec56697691f 100644 --- a/datafusion/functions/benches/to_timestamp.rs +++ b/datafusion/functions/benches/to_timestamp.rs @@ -22,10 +22,10 @@ use std::sync::Arc; use arrow::array::builder::StringBuilder; use arrow::array::{Array, ArrayRef, StringArray}; use arrow::compute::cast; -use arrow::datatypes::DataType; +use arrow::datatypes::{DataType, TimeUnit}; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::to_timestamp; fn data() -> StringArray { @@ -109,16 +109,20 @@ fn data_with_formats() -> (StringArray, StringArray, StringArray, StringArray) { ) } fn criterion_benchmark(c: &mut Criterion) { + let return_type = &DataType::Timestamp(TimeUnit::Nanosecond, None); c.bench_function("to_timestamp_no_formats_utf8", |b| { let arr_data = data(); let batch_len = arr_data.len(); let string_array = ColumnarValue::Array(Arc::new(arr_data) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&[string_array.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![string_array.clone()], + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) @@ -130,10 +134,13 @@ fn criterion_benchmark(c: &mut Criterion) { let string_array = ColumnarValue::Array(Arc::new(data) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&[string_array.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![string_array.clone()], + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) @@ -145,10 +152,13 @@ fn criterion_benchmark(c: &mut Criterion) { let string_array = ColumnarValue::Array(Arc::new(data) as ArrayRef); b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&[string_array.clone()], batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: vec![string_array.clone()], + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) @@ -158,17 +168,20 @@ fn criterion_benchmark(c: &mut Criterion) { let (inputs, format1, format2, format3) = data_with_formats(); let batch_len = inputs.len(); - let args = [ + let args = vec![ ColumnarValue::Array(Arc::new(inputs) as ArrayRef), ColumnarValue::Array(Arc::new(format1) as ArrayRef), ColumnarValue::Array(Arc::new(format2) as ArrayRef), ColumnarValue::Array(Arc::new(format3) as ArrayRef), ]; b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&args.clone(), batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) @@ -178,7 +191,7 @@ fn criterion_benchmark(c: &mut Criterion) { let (inputs, format1, format2, format3) = data_with_formats(); let batch_len = inputs.len(); - let args = [ + let args = vec![ ColumnarValue::Array( Arc::new(cast(&inputs, &DataType::LargeUtf8).unwrap()) as ArrayRef ), @@ -193,10 +206,13 @@ fn criterion_benchmark(c: &mut Criterion) { ), ]; b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&args.clone(), batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) @@ -207,7 +223,7 @@ fn criterion_benchmark(c: &mut Criterion) { let batch_len = inputs.len(); - let args = [ + let args = vec![ ColumnarValue::Array( Arc::new(cast(&inputs, &DataType::Utf8View).unwrap()) as ArrayRef ), @@ -222,10 +238,13 @@ fn criterion_benchmark(c: &mut Criterion) { ), ]; b.iter(|| { - // TODO use invoke_with_args black_box( to_timestamp() - .invoke_batch(&args.clone(), batch_len) + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + number_rows: batch_len, + return_type, + }) .expect("to_timestamp should work on valid values"), ) }) diff --git a/datafusion/functions/benches/trunc.rs b/datafusion/functions/benches/trunc.rs index 83d5b761e809..7fc93921d2e7 100644 --- a/datafusion/functions/benches/trunc.rs +++ b/datafusion/functions/benches/trunc.rs @@ -22,9 +22,10 @@ use arrow::{ util::bench_util::create_primitive_array, }; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_expr::ColumnarValue; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::trunc; +use arrow::datatypes::DataType; use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { @@ -34,16 +35,30 @@ fn criterion_benchmark(c: &mut Criterion) { let f32_args = vec![ColumnarValue::Array(f32_array)]; c.bench_function(&format!("trunc f32 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(trunc.invoke_batch(&f32_args, size).unwrap()) + black_box( + trunc + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + number_rows: size, + return_type: &DataType::Float32, + }) + .unwrap(), + ) }) }); let f64_array = Arc::new(create_primitive_array::(size, 0.2)); let f64_args = vec![ColumnarValue::Array(f64_array)]; c.bench_function(&format!("trunc f64 array: {}", size), |b| { b.iter(|| { - // TODO use invoke_with_args - black_box(trunc.invoke_batch(&f64_args, size).unwrap()) + black_box( + trunc + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + number_rows: size, + return_type: &DataType::Float64, + }) + .unwrap(), + ) }) }); } diff --git a/datafusion/functions/src/core/union_extract.rs b/datafusion/functions/src/core/union_extract.rs index 95814197d8df..420eeed42cc3 100644 --- a/datafusion/functions/src/core/union_extract.rs +++ b/datafusion/functions/src/core/union_extract.rs @@ -82,8 +82,8 @@ impl ScalarUDFImpl for UnionExtractFun { } fn return_type(&self, _: &[DataType]) -> Result { - // should be using return_type_from_exprs and not calling the default implementation - internal_err!("union_extract should return type from exprs") + // should be using return_type_from_args and not calling the default implementation + internal_err!("union_extract should return type from args") } fn return_type_from_args(&self, args: ReturnTypeArgs) -> Result { diff --git a/datafusion/functions/src/datetime/to_char.rs b/datafusion/functions/src/datetime/to_char.rs index 034bbb705070..8b2e5ad87471 100644 --- a/datafusion/functions/src/datetime/to_char.rs +++ b/datafusion/functions/src/datetime/to_char.rs @@ -212,14 +212,6 @@ fn _to_char_scalar( let is_scalar_expression = matches!(&expression, ColumnarValue::Scalar(_)); let array = expression.into_array(1)?; - // fix https://github.com/apache/datafusion/issues/14884 - // If the input date/time is null, return a null Utf8 result. - if array.is_null(0) { - return Ok(match is_scalar_expression { - true => ColumnarValue::Scalar(ScalarValue::Utf8(None)), - false => ColumnarValue::Array(new_null_array(&Utf8, array.len())), - }); - } if format.is_none() { if is_scalar_expression { return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))); @@ -234,15 +226,21 @@ fn _to_char_scalar( }; let formatter = ArrayFormatter::try_new(array.as_ref(), &format_options)?; - let formatted: Result, ArrowError> = (0..array.len()) - .map(|i| formatter.value(i).try_to_string()) + let formatted: Result>, ArrowError> = (0..array.len()) + .map(|i| { + if array.is_null(i) { + Ok(None) + } else { + formatter.value(i).try_to_string().map(Some) + } + }) .collect(); if let Ok(formatted) = formatted { if is_scalar_expression { - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some( - formatted.first().unwrap().to_string(), - )))) + Ok(ColumnarValue::Scalar(ScalarValue::Utf8( + formatted.first().unwrap().clone(), + ))) } else { Ok(ColumnarValue::Array( Arc::new(StringArray::from(formatted)) as ArrayRef @@ -260,13 +258,6 @@ fn _to_char_array(args: &[ColumnarValue]) -> Result { let data_type = arrays[0].data_type(); for idx in 0..arrays[0].len() { - // fix https://github.com/apache/datafusion/issues/14884 - // If the date/time value is null, push None. - if arrays[0].is_null(idx) { - results.push(None); - continue; - } - let format = if format_array.is_null(idx) { None } else { @@ -678,31 +669,4 @@ mod tests { "Execution error: Format for `to_char` must be non-null Utf8, received Timestamp(Nanosecond, None)" ); } - - #[test] - fn test_to_char_input_none_array() { - let date_array = Arc::new(Date32Array::from(vec![Some(18506), None])) as ArrayRef; - let format_array = - StringArray::from(vec!["%Y-%m-%d".to_string(), "%Y-%m-%d".to_string()]); - let args = datafusion_expr::ScalarFunctionArgs { - args: vec![ - ColumnarValue::Array(date_array), - ColumnarValue::Array(Arc::new(format_array) as ArrayRef), - ], - number_rows: 2, - return_type: &DataType::Utf8, - }; - let result = ToCharFunc::new() - .invoke_with_args(args) - .expect("Expected no error"); - if let ColumnarValue::Array(result) = result { - let result = result.as_any().downcast_ref::().unwrap(); - assert_eq!(result.len(), 2); - // The first element is valid, second is null. - assert!(!result.is_null(0)); - assert!(result.is_null(1)); - } else { - panic!("Expected an array value"); - } - } } diff --git a/datafusion/functions/src/regex/regexplike.rs b/datafusion/functions/src/regex/regexplike.rs index 6006309306d5..2080bb9fe818 100644 --- a/datafusion/functions/src/regex/regexplike.rs +++ b/datafusion/functions/src/regex/regexplike.rs @@ -21,12 +21,15 @@ use arrow::array::{Array, ArrayRef, AsArray, GenericStringArray}; use arrow::compute::kernels::regexp; use arrow::datatypes::DataType; use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; -use datafusion_common::exec_err; -use datafusion_common::ScalarValue; -use datafusion_common::{arrow_datafusion_err, plan_err}; -use datafusion_common::{internal_err, DataFusionError, Result}; -use datafusion_expr::{ColumnarValue, Documentation, TypeSignature}; -use datafusion_expr::{ScalarUDFImpl, Signature, Volatility}; +use datafusion_common::types::logical_string; +use datafusion_common::{ + arrow_datafusion_err, exec_err, internal_err, plan_err, DataFusionError, Result, + ScalarValue, +}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; use std::any::Any; @@ -79,7 +82,17 @@ impl RegexpLikeFunc { pub fn new() -> Self { Self { signature: Signature::one_of( - vec![TypeSignature::String(2), TypeSignature::String(3)], + vec![ + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ]), + ], Volatility::Immutable, ), } diff --git a/datafusion/functions/src/string/bit_length.rs b/datafusion/functions/src/string/bit_length.rs index 2a782c59963e..f8740aa4178b 100644 --- a/datafusion/functions/src/string/bit_length.rs +++ b/datafusion/functions/src/string/bit_length.rs @@ -20,9 +20,13 @@ use arrow::datatypes::DataType; use std::any::Any; use crate::utils::utf8_to_int_type; -use datafusion_common::{utils::take_function_args, Result, ScalarValue}; -use datafusion_expr::{ColumnarValue, Documentation, Volatility}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datafusion_common::types::logical_string; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -55,7 +59,12 @@ impl Default for BitLengthFunc { impl BitLengthFunc { pub fn new() -> Self { Self { - signature: Signature::string(1, Volatility::Immutable), + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_string(), + ))], + Volatility::Immutable, + ), } } } diff --git a/datafusion/functions/src/string/contains.rs b/datafusion/functions/src/string/contains.rs index 77774cdb5e1d..05a3edf61c5a 100644 --- a/datafusion/functions/src/string/contains.rs +++ b/datafusion/functions/src/string/contains.rs @@ -20,12 +20,12 @@ use arrow::array::{Array, ArrayRef, AsArray}; use arrow::compute::contains as arrow_contains; use arrow::datatypes::DataType; use arrow::datatypes::DataType::{Boolean, LargeUtf8, Utf8, Utf8View}; -use datafusion_common::exec_err; -use datafusion_common::DataFusionError; -use datafusion_common::Result; +use datafusion_common::types::logical_string; +use datafusion_common::{exec_err, DataFusionError, Result}; +use datafusion_expr::binary::{binary_to_string_coercion, string_coercion}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::any::Any; @@ -60,7 +60,13 @@ impl Default for ContainsFunc { impl ContainsFunc { pub fn new() -> Self { Self { - signature: Signature::string(2, Volatility::Immutable), + signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ], + Volatility::Immutable, + ), } } } @@ -92,29 +98,52 @@ impl ScalarUDFImpl for ContainsFunc { } /// use `arrow::compute::contains` to do the calculation for contains -pub fn contains(args: &[ArrayRef]) -> Result { - match (args[0].data_type(), args[1].data_type()) { - (Utf8View, Utf8View) => { - let mod_str = args[0].as_string_view(); - let match_str = args[1].as_string_view(); - let res = arrow_contains(mod_str, match_str)?; - Ok(Arc::new(res) as ArrayRef) - } - (Utf8, Utf8) => { - let mod_str = args[0].as_string::(); - let match_str = args[1].as_string::(); - let res = arrow_contains(mod_str, match_str)?; - Ok(Arc::new(res) as ArrayRef) - } - (LargeUtf8, LargeUtf8) => { - let mod_str = args[0].as_string::(); - let match_str = args[1].as_string::(); - let res = arrow_contains(mod_str, match_str)?; - Ok(Arc::new(res) as ArrayRef) - } - other => { - exec_err!("Unsupported data type {other:?} for function `contains`.") +fn contains(args: &[ArrayRef]) -> Result { + if let Some(coercion_data_type) = + string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| { + binary_to_string_coercion(args[0].data_type(), args[1].data_type()) + }) + { + let arg0 = if args[0].data_type() == &coercion_data_type { + Arc::clone(&args[0]) + } else { + arrow::compute::kernels::cast::cast(&args[0], &coercion_data_type)? + }; + let arg1 = if args[1].data_type() == &coercion_data_type { + Arc::clone(&args[1]) + } else { + arrow::compute::kernels::cast::cast(&args[1], &coercion_data_type)? + }; + + match coercion_data_type { + Utf8View => { + let mod_str = arg0.as_string_view(); + let match_str = arg1.as_string_view(); + let res = arrow_contains(mod_str, match_str)?; + Ok(Arc::new(res) as ArrayRef) + } + Utf8 => { + let mod_str = arg0.as_string::(); + let match_str = arg1.as_string::(); + let res = arrow_contains(mod_str, match_str)?; + Ok(Arc::new(res) as ArrayRef) + } + LargeUtf8 => { + let mod_str = arg0.as_string::(); + let match_str = arg1.as_string::(); + let res = arrow_contains(mod_str, match_str)?; + Ok(Arc::new(res) as ArrayRef) + } + other => { + exec_err!("Unsupported data type {other:?} for function `contains`.") + } } + } else { + exec_err!( + "Unsupported data type {:?}, {:?} for function `contains`.", + args[0].data_type(), + args[1].data_type() + ) } } diff --git a/datafusion/functions/src/string/ends_with.rs b/datafusion/functions/src/string/ends_with.rs index 5cca79de14ff..eafc310236ee 100644 --- a/datafusion/functions/src/string/ends_with.rs +++ b/datafusion/functions/src/string/ends_with.rs @@ -22,9 +22,13 @@ use arrow::array::ArrayRef; use arrow::datatypes::DataType; use crate::utils::make_scalar_function; +use datafusion_common::types::logical_string; use datafusion_common::{internal_err, Result}; -use datafusion_expr::{ColumnarValue, Documentation, Volatility}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datafusion_expr::binary::{binary_to_string_coercion, string_coercion}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -62,7 +66,13 @@ impl Default for EndsWithFunc { impl EndsWithFunc { pub fn new() -> Self { Self { - signature: Signature::string(2, Volatility::Immutable), + signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ], + Volatility::Immutable, + ), } } } @@ -102,10 +112,29 @@ impl ScalarUDFImpl for EndsWithFunc { /// Returns true if string ends with suffix. /// ends_with('alphabet', 'abet') = 't' -pub fn ends_with(args: &[ArrayRef]) -> Result { - let result = arrow::compute::kernels::comparison::ends_with(&args[0], &args[1])?; - - Ok(Arc::new(result) as ArrayRef) +fn ends_with(args: &[ArrayRef]) -> Result { + if let Some(coercion_data_type) = + string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| { + binary_to_string_coercion(args[0].data_type(), args[1].data_type()) + }) + { + let arg0 = if args[0].data_type() == &coercion_data_type { + Arc::clone(&args[0]) + } else { + arrow::compute::kernels::cast::cast(&args[0], &coercion_data_type)? + }; + let arg1 = if args[1].data_type() == &coercion_data_type { + Arc::clone(&args[1]) + } else { + arrow::compute::kernels::cast::cast(&args[1], &coercion_data_type)? + }; + let result = arrow::compute::kernels::comparison::ends_with(&arg0, &arg1)?; + Ok(Arc::new(result) as ArrayRef) + } else { + internal_err!( + "Unsupported data types for ends_with. Expected Utf8, LargeUtf8 or Utf8View" + ) + } } #[cfg(test)] diff --git a/datafusion/functions/src/string/levenshtein.rs b/datafusion/functions/src/string/levenshtein.rs index a19fcc5b476c..a1a486c7d3cf 100644 --- a/datafusion/functions/src/string/levenshtein.rs +++ b/datafusion/functions/src/string/levenshtein.rs @@ -23,10 +23,17 @@ use arrow::datatypes::DataType; use crate::utils::{make_scalar_function, utf8_to_int_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; +use datafusion_common::types::logical_string; use datafusion_common::utils::datafusion_strsim; -use datafusion_common::{exec_err, utils::take_function_args, Result}; -use datafusion_expr::{ColumnarValue, Documentation}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{exec_err, Result}; +use datafusion_expr::type_coercion::binary::{ + binary_to_string_coercion, string_coercion, +}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -64,7 +71,13 @@ impl Default for LevenshteinFunc { impl LevenshteinFunc { pub fn new() -> Self { Self { - signature: Signature::string(2, Volatility::Immutable), + signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ], + Volatility::Immutable, + ), } } } @@ -83,7 +96,13 @@ impl ScalarUDFImpl for LevenshteinFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "levenshtein") + if let Some(coercion_data_type) = string_coercion(&arg_types[0], &arg_types[1]) + .or_else(|| binary_to_string_coercion(&arg_types[0], &arg_types[1])) + { + utf8_to_int_type(&coercion_data_type, "levenshtein") + } else { + exec_err!("Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View") + } } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -107,60 +126,79 @@ impl ScalarUDFImpl for LevenshteinFunc { ///Returns the Levenshtein distance between the two given strings. /// LEVENSHTEIN('kitten', 'sitting') = 3 -pub fn levenshtein(args: &[ArrayRef]) -> Result { +fn levenshtein(args: &[ArrayRef]) -> Result { let [str1, str2] = take_function_args("levenshtein", args)?; - match str1.data_type() { - DataType::Utf8View => { - let str1_array = as_string_view_array(&str1)?; - let str2_array = as_string_view_array(&str2)?; - let result = str1_array - .iter() - .zip(str2_array.iter()) - .map(|(string1, string2)| match (string1, string2) { - (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein(string1, string2) as i32) - } - _ => None, - }) - .collect::(); - Ok(Arc::new(result) as ArrayRef) - } - DataType::Utf8 => { - let str1_array = as_generic_string_array::(&str1)?; - let str2_array = as_generic_string_array::(&str2)?; - let result = str1_array - .iter() - .zip(str2_array.iter()) - .map(|(string1, string2)| match (string1, string2) { - (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein(string1, string2) as i32) - } - _ => None, - }) - .collect::(); - Ok(Arc::new(result) as ArrayRef) - } - DataType::LargeUtf8 => { - let str1_array = as_generic_string_array::(&str1)?; - let str2_array = as_generic_string_array::(&str2)?; - let result = str1_array - .iter() - .zip(str2_array.iter()) - .map(|(string1, string2)| match (string1, string2) { - (Some(string1), Some(string2)) => { - Some(datafusion_strsim::levenshtein(string1, string2) as i64) - } - _ => None, - }) - .collect::(); - Ok(Arc::new(result) as ArrayRef) - } - other => { - exec_err!( - "levenshtein was called with {other} datatype arguments. It requires Utf8View, Utf8 or LargeUtf8." - ) + if let Some(coercion_data_type) = + string_coercion(args[0].data_type(), args[1].data_type()).or_else(|| { + binary_to_string_coercion(args[0].data_type(), args[1].data_type()) + }) + { + let str1 = if str1.data_type() == &coercion_data_type { + Arc::clone(str1) + } else { + arrow::compute::kernels::cast::cast(&str1, &coercion_data_type)? + }; + let str2 = if str2.data_type() == &coercion_data_type { + Arc::clone(str2) + } else { + arrow::compute::kernels::cast::cast(&str2, &coercion_data_type)? + }; + + match coercion_data_type { + DataType::Utf8View => { + let str1_array = as_string_view_array(&str1)?; + let str2_array = as_string_view_array(&str2)?; + let result = str1_array + .iter() + .zip(str2_array.iter()) + .map(|(string1, string2)| match (string1, string2) { + (Some(string1), Some(string2)) => { + Some(datafusion_strsim::levenshtein(string1, string2) as i32) + } + _ => None, + }) + .collect::(); + Ok(Arc::new(result) as ArrayRef) + } + DataType::Utf8 => { + let str1_array = as_generic_string_array::(&str1)?; + let str2_array = as_generic_string_array::(&str2)?; + let result = str1_array + .iter() + .zip(str2_array.iter()) + .map(|(string1, string2)| match (string1, string2) { + (Some(string1), Some(string2)) => { + Some(datafusion_strsim::levenshtein(string1, string2) as i32) + } + _ => None, + }) + .collect::(); + Ok(Arc::new(result) as ArrayRef) + } + DataType::LargeUtf8 => { + let str1_array = as_generic_string_array::(&str1)?; + let str2_array = as_generic_string_array::(&str2)?; + let result = str1_array + .iter() + .zip(str2_array.iter()) + .map(|(string1, string2)| match (string1, string2) { + (Some(string1), Some(string2)) => { + Some(datafusion_strsim::levenshtein(string1, string2) as i64) + } + _ => None, + }) + .collect::(); + Ok(Arc::new(result) as ArrayRef) + } + other => { + exec_err!( + "levenshtein was called with {other} datatype arguments. It requires Utf8View, Utf8 or LargeUtf8." + ) + } } + } else { + exec_err!("Unsupported data types for levenshtein. Expected Utf8, LargeUtf8 or Utf8View") } } diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 375717e23d6d..226275b13999 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -20,9 +20,12 @@ use std::any::Any; use crate::string::common::to_lower; use crate::utils::utf8_to_str_type; +use datafusion_common::types::logical_string; use datafusion_common::Result; -use datafusion_expr::{ColumnarValue, Documentation}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -55,7 +58,12 @@ impl Default for LowerFunc { impl LowerFunc { pub fn new() -> Self { Self { - signature: Signature::string(1, Volatility::Immutable), + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_string(), + ))], + Volatility::Immutable, + ), } } } diff --git a/datafusion/functions/src/string/octet_length.rs b/datafusion/functions/src/string/octet_length.rs index 46175c96cdc6..17ea2726b071 100644 --- a/datafusion/functions/src/string/octet_length.rs +++ b/datafusion/functions/src/string/octet_length.rs @@ -20,9 +20,13 @@ use arrow::datatypes::DataType; use std::any::Any; use crate::utils::utf8_to_int_type; -use datafusion_common::{utils::take_function_args, Result, ScalarValue}; -use datafusion_expr::{ColumnarValue, Documentation, Volatility}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datafusion_common::types::logical_string; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -55,7 +59,12 @@ impl Default for OctetLengthFunc { impl OctetLengthFunc { pub fn new() -> Self { Self { - signature: Signature::string(1, Volatility::Immutable), + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_string(), + ))], + Volatility::Immutable, + ), } } } diff --git a/datafusion/functions/src/string/replace.rs b/datafusion/functions/src/string/replace.rs index a3488b561fd2..de70215c49c7 100644 --- a/datafusion/functions/src/string/replace.rs +++ b/datafusion/functions/src/string/replace.rs @@ -23,9 +23,15 @@ use arrow::datatypes::DataType; use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; +use datafusion_common::types::logical_string; use datafusion_common::{exec_err, Result}; -use datafusion_expr::{ColumnarValue, Documentation, Volatility}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; +use datafusion_expr::type_coercion::binary::{ + binary_to_string_coercion, string_coercion, +}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; #[user_doc( doc_section(label = "String Functions"), @@ -60,7 +66,14 @@ impl Default for ReplaceFunc { impl ReplaceFunc { pub fn new() -> Self { Self { - signature: Signature::string(3, Volatility::Immutable), + signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ], + Volatility::Immutable, + ), } } } @@ -79,19 +92,64 @@ impl ScalarUDFImpl for ReplaceFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_str_type(&arg_types[0], "replace") + if let Some(coercion_data_type) = string_coercion(&arg_types[0], &arg_types[1]) + .and_then(|dt| string_coercion(&dt, &arg_types[2])) + .or_else(|| { + binary_to_string_coercion(&arg_types[0], &arg_types[1]) + .and_then(|dt| binary_to_string_coercion(&dt, &arg_types[2])) + }) + { + utf8_to_str_type(&coercion_data_type, "replace") + } else { + exec_err!("Unsupported data types for replace. Expected Utf8, LargeUtf8 or Utf8View") + } } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - match args.args[0].data_type() { - DataType::Utf8 => make_scalar_function(replace::, vec![])(&args.args), - DataType::LargeUtf8 => { - make_scalar_function(replace::, vec![])(&args.args) + let data_types = args + .args + .iter() + .map(|arg| arg.data_type()) + .collect::>(); + + if let Some(coercion_type) = string_coercion(&data_types[0], &data_types[1]) + .and_then(|dt| string_coercion(&dt, &data_types[2])) + .or_else(|| { + binary_to_string_coercion(&data_types[0], &data_types[1]) + .and_then(|dt| binary_to_string_coercion(&dt, &data_types[2])) + }) + { + let mut converted_args = Vec::with_capacity(args.args.len()); + for arg in &args.args { + if arg.data_type() == coercion_type { + converted_args.push(arg.clone()); + } else { + let converted = arg.cast_to(&coercion_type, None)?; + converted_args.push(converted); + } } - DataType::Utf8View => make_scalar_function(replace_view, vec![])(&args.args), - other => { - exec_err!("Unsupported data type {other:?} for function replace") + + match coercion_type { + DataType::Utf8 => { + make_scalar_function(replace::, vec![])(&converted_args) + } + DataType::LargeUtf8 => { + make_scalar_function(replace::, vec![])(&converted_args) + } + DataType::Utf8View => { + make_scalar_function(replace_view, vec![])(&converted_args) + } + other => exec_err!( + "Unsupported coercion data type {other:?} for function replace" + ), } + } else { + exec_err!( + "Unsupported data type {:?}, {:?}, {:?} for function replace.", + data_types[0], + data_types[1], + data_types[2] + ) } } @@ -117,6 +175,7 @@ fn replace_view(args: &[ArrayRef]) -> Result { Ok(Arc::new(result) as ArrayRef) } + /// Replaces all occurrences in string of substring from with substring to. /// replace('abcdefabcdef', 'cd', 'XX') = 'abXXefabXXef' fn replace(args: &[ArrayRef]) -> Result { diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index d27b54d29bc6..2fec7305d183 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -18,9 +18,12 @@ use crate::string::common::to_upper; use crate::utils::utf8_to_str_type; use arrow::datatypes::DataType; +use datafusion_common::types::logical_string; use datafusion_common::Result; -use datafusion_expr::{ColumnarValue, Documentation}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_macros::user_doc; use std::any::Any; @@ -54,7 +57,12 @@ impl Default for UpperFunc { impl UpperFunc { pub fn new() -> Self { Self { - signature: Signature::string(1, Volatility::Immutable), + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_string(), + ))], + Volatility::Immutable, + ), } } } diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs index a8a4dd0fa249..c9b0cb77b096 100644 --- a/datafusion/functions/src/unicode/initcap.rs +++ b/datafusion/functions/src/unicode/initcap.rs @@ -25,9 +25,12 @@ use arrow::datatypes::DataType; use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; +use datafusion_common::types::logical_string; use datafusion_common::{exec_err, Result}; -use datafusion_expr::{ColumnarValue, Documentation, Volatility}; -use datafusion_expr::{ScalarUDFImpl, Signature}; +use datafusion_expr::{ + Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass, + Volatility, +}; use datafusion_macros::user_doc; #[user_doc( @@ -61,7 +64,12 @@ impl Default for InitcapFunc { impl InitcapFunc { pub fn new() -> Self { Self { - signature: Signature::string(1, Volatility::Immutable), + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_string(), + ))], + Volatility::Immutable, + ), } } } diff --git a/datafusion/functions/src/unicode/strpos.rs b/datafusion/functions/src/unicode/strpos.rs index a0ad2e3f0b75..b3bc73a29585 100644 --- a/datafusion/functions/src/unicode/strpos.rs +++ b/datafusion/functions/src/unicode/strpos.rs @@ -23,9 +23,11 @@ use arrow::array::{ ArrayRef, ArrowPrimitiveType, AsArray, PrimitiveArray, StringArrayType, }; use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type}; +use datafusion_common::types::logical_string; use datafusion_common::{exec_err, internal_err, Result}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility, + Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass, + Volatility, }; use datafusion_macros::user_doc; @@ -60,7 +62,13 @@ impl Default for StrposFunc { impl StrposFunc { pub fn new() -> Self { Self { - signature: Signature::string(2, Volatility::Immutable), + signature: Signature::coercible( + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + Coercion::new_exact(TypeSignatureClass::Native(logical_string())), + ], + Volatility::Immutable, + ), aliases: vec![String::from("instr"), String::from("position")], } } @@ -115,6 +123,11 @@ fn strpos(args: &[ArrayRef]) -> Result { let substring_array = args[1].as_string::(); calculate_strpos::<_, _, Int32Type>(string_array, substring_array) } + (DataType::Utf8, DataType::Utf8View) => { + let string_array = args[0].as_string::(); + let substring_array = args[1].as_string_view(); + calculate_strpos::<_, _, Int32Type>(string_array, substring_array) + } (DataType::Utf8, DataType::LargeUtf8) => { let string_array = args[0].as_string::(); let substring_array = args[1].as_string::(); @@ -125,6 +138,11 @@ fn strpos(args: &[ArrayRef]) -> Result { let substring_array = args[1].as_string::(); calculate_strpos::<_, _, Int64Type>(string_array, substring_array) } + (DataType::LargeUtf8, DataType::Utf8View) => { + let string_array = args[0].as_string::(); + let substring_array = args[1].as_string_view(); + calculate_strpos::<_, _, Int64Type>(string_array, substring_array) + } (DataType::LargeUtf8, DataType::LargeUtf8) => { let string_array = args[0].as_string::(); let substring_array = args[1].as_string::(); diff --git a/datafusion/macros/Cargo.toml b/datafusion/macros/Cargo.toml index 737d2ed72874..bf412d19784f 100644 --- a/datafusion/macros/Cargo.toml +++ b/datafusion/macros/Cargo.toml @@ -42,4 +42,4 @@ proc-macro = true [dependencies] datafusion-expr = { workspace = true } quote = "1.0.37" -syn = { version = "2.0.79", features = ["full"] } +syn = { version = "2.0.100", features = ["full"] } diff --git a/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs deleted file mode 100644 index 7df4e970ada2..000000000000 --- a/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs +++ /dev/null @@ -1,332 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; - -use crate::AnalyzerRule; -use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult}; -use datafusion_common::{Column, Result}; -use datafusion_expr::builder::validate_unique_names; -use datafusion_expr::expr::PlannedReplaceSelectItem; -use datafusion_expr::utils::{ - expand_qualified_wildcard, expand_wildcard, find_base_plan, -}; -use datafusion_expr::{ - Distinct, DistinctOn, Expr, LogicalPlan, Projection, SubqueryAlias, -}; - -#[derive(Default, Debug)] -pub struct ExpandWildcardRule {} - -impl ExpandWildcardRule { - pub fn new() -> Self { - Self {} - } -} - -impl AnalyzerRule for ExpandWildcardRule { - fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result { - // Because the wildcard expansion is based on the schema of the input plan, - // using `transform_up_with_subqueries` here. - plan.transform_up_with_subqueries(expand_internal).data() - } - - fn name(&self) -> &str { - "expand_wildcard_rule" - } -} - -fn expand_internal(plan: LogicalPlan) -> Result> { - match plan { - LogicalPlan::Projection(Projection { expr, input, .. }) => { - let projected_expr = expand_exprlist(&input, expr)?; - validate_unique_names("Projections", projected_expr.iter())?; - Ok(Transformed::yes( - Projection::try_new(projected_expr, Arc::clone(&input)) - .map(LogicalPlan::Projection)?, - )) - } - // The schema of the plan should also be updated if the child plan is transformed. - LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => { - Ok(Transformed::yes( - SubqueryAlias::try_new(input, alias).map(LogicalPlan::SubqueryAlias)?, - )) - } - LogicalPlan::Distinct(Distinct::On(distinct_on)) => { - let projected_expr = - expand_exprlist(&distinct_on.input, distinct_on.select_expr)?; - validate_unique_names("Distinct", projected_expr.iter())?; - Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::On( - DistinctOn::try_new( - distinct_on.on_expr, - projected_expr, - distinct_on.sort_expr, - distinct_on.input, - )?, - )))) - } - _ => Ok(Transformed::no(plan)), - } -} - -fn expand_exprlist(input: &LogicalPlan, expr: Vec) -> Result> { - let mut projected_expr = vec![]; - let input = find_base_plan(input); - for e in expr { - match e { - Expr::Wildcard { qualifier, options } => { - if let Some(qualifier) = qualifier { - let expanded = expand_qualified_wildcard( - &qualifier, - input.schema(), - Some(&options), - )?; - // If there is a REPLACE statement, replace that column with the given - // replace expression. Column name remains the same. - let replaced = if let Some(replace) = options.replace { - replace_columns(expanded, &replace)? - } else { - expanded - }; - projected_expr.extend(replaced); - } else { - let expanded = - expand_wildcard(input.schema(), input, Some(&options))?; - // If there is a REPLACE statement, replace that column with the given - // replace expression. Column name remains the same. - let replaced = if let Some(replace) = options.replace { - replace_columns(expanded, &replace)? - } else { - expanded - }; - projected_expr.extend(replaced); - } - } - // A workaround to handle the case when the column name is "*". - // We transform the expression to a Expr::Column through [Column::from_name] in many places. - // It would also convert the wildcard expression to a column expression with name "*". - Expr::Column(Column { - ref relation, - ref name, - // TODO Should we use these spans? - spans: _, - }) => { - if name.eq("*") { - if let Some(qualifier) = relation { - projected_expr.extend(expand_qualified_wildcard( - qualifier, - input.schema(), - None, - )?); - } else { - projected_expr.extend(expand_wildcard( - input.schema(), - input, - None, - )?); - } - } else { - projected_expr.push(e.clone()); - } - } - _ => projected_expr.push(e), - } - } - Ok(projected_expr) -} - -/// If there is a REPLACE statement in the projected expression in the form of -/// "REPLACE (some_column_within_an_expr AS some_column)", this function replaces -/// that column with the given replace expression. Column name remains the same. -/// Multiple REPLACEs are also possible with comma separations. -fn replace_columns( - mut exprs: Vec, - replace: &PlannedReplaceSelectItem, -) -> Result> { - for expr in exprs.iter_mut() { - if let Expr::Column(Column { name, .. }) = expr { - if let Some((_, new_expr)) = replace - .items() - .iter() - .zip(replace.expressions().iter()) - .find(|(item, _)| item.column_name.value == *name) - { - *expr = new_expr.clone().alias(name.clone()) - } - } - } - Ok(exprs) -} - -#[cfg(test)] -mod tests { - use arrow::datatypes::{DataType, Field, Schema}; - - use crate::test::{assert_analyzed_plan_eq_display_indent, test_table_scan}; - use crate::Analyzer; - use datafusion_common::{JoinType, TableReference}; - use datafusion_expr::{ - col, in_subquery, qualified_wildcard, table_scan, wildcard, LogicalPlanBuilder, - }; - - use super::*; - - fn assert_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> { - assert_analyzed_plan_eq_display_indent( - Arc::new(ExpandWildcardRule::new()), - plan, - expected, - ) - } - - #[test] - fn test_expand_wildcard() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .project(vec![wildcard()])? - .build()?; - let expected = - "Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_expand_qualified_wildcard() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .project(vec![qualified_wildcard(TableReference::bare("test"))])? - .build()?; - let expected = - "Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_expand_qualified_wildcard_in_subquery() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .project(vec![qualified_wildcard(TableReference::bare("test"))])? - .build()?; - let plan = LogicalPlanBuilder::from(plan) - .project(vec![wildcard()])? - .build()?; - let expected = - "Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32]\ - \n Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_expand_wildcard_in_subquery() -> Result<()> { - let projection_a = LogicalPlanBuilder::from(test_table_scan()?) - .project(vec![col("a")])? - .build()?; - let subquery = LogicalPlanBuilder::from(projection_a) - .project(vec![wildcard()])? - .build()?; - let plan = LogicalPlanBuilder::from(test_table_scan()?) - .filter(in_subquery(col("a"), Arc::new(subquery)))? - .project(vec![wildcard()])? - .build()?; - let expected = "\ - Projection: test.a, test.b, test.c [a:UInt32, b:UInt32, c:UInt32]\ - \n Filter: test.a IN () [a:UInt32, b:UInt32, c:UInt32]\ - \n Subquery: [a:UInt32]\ - \n Projection: test.a [a:UInt32]\ - \n Projection: test.a [a:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_expand_wildcard_in_distinct_on() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .distinct_on(vec![col("a")], vec![wildcard()], None)? - .build()?; - let expected = "\ - DistinctOn: on_expr=[[test.a]], select_expr=[[test.a, test.b, test.c]], sort_expr=[[]] [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_subquery_schema() -> Result<()> { - let analyzer = Analyzer::with_rules(vec![Arc::new(ExpandWildcardRule::new())]); - let options = ConfigOptions::default(); - let subquery = LogicalPlanBuilder::from(test_table_scan()?) - .project(vec![wildcard()])? - .build()?; - let plan = LogicalPlanBuilder::from(subquery) - .alias("sub")? - .project(vec![wildcard()])? - .build()?; - let analyzed_plan = analyzer.execute_and_check(plan, &options, |_, _| {})?; - for x in analyzed_plan.inputs() { - for field in x.schema().fields() { - assert_ne!(field.name(), "*"); - } - } - Ok(()) - } - - fn employee_schema() -> Schema { - Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("first_name", DataType::Utf8, false), - Field::new("last_name", DataType::Utf8, false), - Field::new("state", DataType::Utf8, false), - Field::new("salary", DataType::Int32, false), - ]) - } - - #[test] - fn plan_using_join_wildcard_projection() -> Result<()> { - let t2 = table_scan(Some("t2"), &employee_schema(), None)?.build()?; - - let plan = table_scan(Some("t1"), &employee_schema(), None)? - .join_using(t2, JoinType::Inner, vec!["id"])? - .project(vec![wildcard()])? - .build()?; - - let expected = "Projection: *\ - \n Inner Join: Using t1.id = t2.id\ - \n TableScan: t1\ - \n TableScan: t2"; - - assert_eq!(expected, format!("{plan}")); - - let analyzer = Analyzer::with_rules(vec![Arc::new(ExpandWildcardRule::new())]); - let options = ConfigOptions::default(); - - let analyzed_plan = analyzer.execute_and_check(plan, &options, |_, _| {})?; - - // id column should only show up once in projection - let expected = "Projection: t1.id, t1.first_name, t1.last_name, t1.state, t1.salary, t2.first_name, t2.last_name, t2.state, t2.salary\ - \n Inner Join: Using t1.id = t2.id\ - \n TableScan: t1\ - \n TableScan: t2"; - assert_eq!(expected, format!("{analyzed_plan}")); - - Ok(()) - } -} diff --git a/datafusion/optimizer/src/analyzer/inline_table_scan.rs b/datafusion/optimizer/src/analyzer/inline_table_scan.rs index 95781b395f3c..350e65e1e329 100644 --- a/datafusion/optimizer/src/analyzer/inline_table_scan.rs +++ b/datafusion/optimizer/src/analyzer/inline_table_scan.rs @@ -23,7 +23,8 @@ use crate::analyzer::AnalyzerRule; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Column, Result}; -use datafusion_expr::{logical_plan::LogicalPlan, wildcard, Expr, LogicalPlanBuilder}; +use datafusion_expr::utils::expand_wildcard; +use datafusion_expr::{logical_plan::LogicalPlan, Expr, LogicalPlanBuilder}; /// Analyzed rule that inlines TableScan that provide a [`LogicalPlan`] /// (DataFrame / ViewTable) @@ -92,7 +93,8 @@ fn generate_projection_expr( ))); } } else { - exprs.push(wildcard()); + let expanded = expand_wildcard(sub_plan.schema(), sub_plan, None)?; + exprs.extend(expanded); } Ok(exprs) } @@ -181,7 +183,7 @@ mod tests { let plan = scan.filter(col("x.a").eq(lit(1)))?.build()?; let expected = "Filter: x.a = Int32(1)\ \n SubqueryAlias: x\ - \n Projection: *\ + \n Projection: y.a, y.b\ \n TableScan: y"; assert_analyzed_plan_eq(Arc::new(InlineTableScan::new()), plan, expected) diff --git a/datafusion/optimizer/src/analyzer/mod.rs b/datafusion/optimizer/src/analyzer/mod.rs index c506616d142e..1d199f2faafc 100644 --- a/datafusion/optimizer/src/analyzer/mod.rs +++ b/datafusion/optimizer/src/analyzer/mod.rs @@ -28,7 +28,6 @@ use datafusion_common::Result; use datafusion_expr::expr_rewriter::FunctionRewrite; use datafusion_expr::{InvariantLevel, LogicalPlan}; -use crate::analyzer::expand_wildcard_rule::ExpandWildcardRule; use crate::analyzer::inline_table_scan::InlineTableScan; use crate::analyzer::resolve_grouping_function::ResolveGroupingFunction; use crate::analyzer::type_coercion::TypeCoercion; @@ -36,7 +35,6 @@ use crate::utils::log_plan; use self::function_rewrite::ApplyFunctionRewrites; -pub mod expand_wildcard_rule; pub mod function_rewrite; pub mod inline_table_scan; pub mod resolve_grouping_function; @@ -99,9 +97,6 @@ impl Analyzer { pub fn new() -> Self { let rules: Vec> = vec![ Arc::new(InlineTableScan::new()), - // Every rule that will generate [Expr::Wildcard] should be placed in front of [ExpandWildcardRule]. - Arc::new(ExpandWildcardRule::new()), - // [Expr::Wildcard] should be expanded before [TypeCoercion] Arc::new(ResolveGroupingFunction::new()), Arc::new(TypeCoercion::new()), ]; diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index d1d491cc7a64..538ef98ac7be 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -565,6 +565,8 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { .build()?, )) } + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::Alias(_) | Expr::Column(_) | Expr::ScalarVariable(_, _) @@ -1021,6 +1023,7 @@ fn project_with_column_index( spans: _, }) if name != schema.field(i).name() => Ok(e.alias(schema.field(i).name())), Expr::Alias { .. } | Expr::Column { .. } => Ok(e), + #[expect(deprecated)] Expr::Wildcard { .. } => { plan_err!("Wildcard should be expanded before type coercion") } diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index bfa53a5ce852..5dc1a7e5ac5b 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -678,6 +678,8 @@ impl CSEController for ExprCSEController<'_> { } fn is_ignored(&self, node: &Expr) -> bool { + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] let is_normal_minus_aggregates = matches!( node, Expr::Literal(..) diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index b192f9740483..71ff863b51a1 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -56,10 +56,16 @@ pub struct PullUpCorrelatedExpr { /// Indicates if we encounter any correlated expression that can not be pulled up /// above a aggregation without changing the meaning of the query. can_pull_over_aggregation: bool, - /// Do we need to handle [the Count bug] during the pull up process. - /// TODO this parameter should be removed or renamed semantically + /// Do we need to handle [the count bug] during the pull up process. /// - /// [the Count bug]: https://github.com/apache/datafusion/issues/10553 + /// The "count bug" was described in [Optimization of Nested SQL + /// Queries Revisited](https://dl.acm.org/doi/pdf/10.1145/38714.38723). This bug is + /// not specific to the COUNT function, and it can occur with any aggregate function, + /// such as SUM, AVG, etc. The anomaly arises because aggregates fail to distinguish + /// between an empty set and null values when optimizing a correlated query as a join. + /// Here, we use "the count bug" to refer to all such cases. + /// + /// [the count bug]: https://github.com/apache/datafusion/issues/10553 pub need_handle_count_bug: bool, /// mapping from the plan to its expressions' evaluation result on empty batch pub collected_count_expr_map: HashMap, @@ -88,10 +94,9 @@ impl PullUpCorrelatedExpr { } } - /// Set if we need to handle [the Count bug] during the pull up process - /// TODO this should be removed or renamed semantically + /// Set if we need to handle [the count bug] during the pull up process /// - /// [the Count bug]: https://github.com/apache/datafusion/issues/10553 + /// [the count bug]: https://github.com/apache/datafusion/issues/10553 pub fn with_need_handle_count_bug(mut self, need_handle_count_bug: bool) -> Self { self.need_handle_count_bug = need_handle_count_bug; self diff --git a/datafusion/optimizer/src/eliminate_group_by_constant.rs b/datafusion/optimizer/src/eliminate_group_by_constant.rs index 1213c8ffb368..7e252d6dcea0 100644 --- a/datafusion/optimizer/src/eliminate_group_by_constant.rs +++ b/datafusion/optimizer/src/eliminate_group_by_constant.rs @@ -121,8 +121,8 @@ mod tests { use datafusion_common::Result; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{ - col, lit, ColumnarValue, LogicalPlanBuilder, ScalarUDF, ScalarUDFImpl, Signature, - TypeSignature, + col, lit, ColumnarValue, LogicalPlanBuilder, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, TypeSignature, }; use datafusion_functions_aggregate::expr_fn::count; @@ -155,11 +155,7 @@ mod tests { fn return_type(&self, _args: &[DataType]) -> Result { Ok(DataType::Int32) } - fn invoke_batch( - &self, - _args: &[ColumnarValue], - _number_rows: usize, - ) -> Result { + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { unimplemented!() } } diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs index 61ca9b31cd29..ce198560805a 100644 --- a/datafusion/optimizer/src/lib.rs +++ b/datafusion/optimizer/src/lib.rs @@ -60,7 +60,6 @@ pub mod replace_distinct_aggregate; pub mod scalar_subquery_to_join; pub mod simplify_expressions; pub mod single_distinct_to_groupby; -pub mod unwrap_cast_in_comparison; pub mod utils; #[cfg(test)] @@ -70,8 +69,6 @@ pub use analyzer::{Analyzer, AnalyzerRule}; pub use optimizer::{ ApplyOrder, Optimizer, OptimizerConfig, OptimizerContext, OptimizerRule, }; -#[allow(deprecated)] -pub use utils::optimize_children; pub(crate) mod join_key_set; mod plan_signature; diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index 49bce3c1ce82..3a69bd91e749 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -54,7 +54,6 @@ use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate; use crate::scalar_subquery_to_join::ScalarSubqueryToJoin; use crate::simplify_expressions::SimplifyExpressions; use crate::single_distinct_to_groupby::SingleDistinctToGroupBy; -use crate::unwrap_cast_in_comparison::UnwrapCastInComparison; use crate::utils::log_plan; /// `OptimizerRule`s transforms one [`LogicalPlan`] into another which @@ -70,24 +69,6 @@ use crate::utils::log_plan; /// [`AnalyzerRule`]: crate::analyzer::AnalyzerRule /// [`SessionState::add_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_optimizer_rule pub trait OptimizerRule: Debug { - /// Try and rewrite `plan` to an optimized form, returning None if the plan - /// cannot be optimized by this rule. - /// - /// Note this API will be deprecated in the future as it requires `clone`ing - /// the input plan, which can be expensive. OptimizerRules should implement - /// [`Self::rewrite`] instead. - #[deprecated( - since = "40.0.0", - note = "please implement supports_rewrite and rewrite instead" - )] - fn try_optimize( - &self, - _plan: &LogicalPlan, - _config: &dyn OptimizerConfig, - ) -> Result> { - internal_err!("Should have called rewrite") - } - /// A human readable name for this optimizer rule fn name(&self) -> &str; @@ -100,15 +81,13 @@ pub trait OptimizerRule: Debug { } /// Does this rule support rewriting owned plans (rather than by reference)? + #[deprecated(since = "47.0.0", note = "This method is no longer used")] fn supports_rewrite(&self) -> bool { true } /// Try to rewrite `plan` to an optimized form, returning `Transformed::yes` /// if the plan was rewritten and `Transformed::no` if it was not. - /// - /// Note: this function is only called if [`Self::supports_rewrite`] returns - /// true. Otherwise the Optimizer calls [`Self::try_optimize`] fn rewrite( &self, _plan: LogicalPlan, @@ -243,7 +222,6 @@ impl Optimizer { let rules: Vec> = vec![ Arc::new(EliminateNestedUnion::new()), Arc::new(SimplifyExpressions::new()), - Arc::new(UnwrapCastInComparison::new()), Arc::new(ReplaceDistinctWithAggregate::new()), Arc::new(EliminateJoin::new()), Arc::new(DecorrelatePredicateSubquery::new()), @@ -266,7 +244,6 @@ impl Optimizer { // The previous optimizations added expressions and projections, // that might benefit from the following rules Arc::new(SimplifyExpressions::new()), - Arc::new(UnwrapCastInComparison::new()), Arc::new(CommonSubexprEliminate::new()), Arc::new(EliminateGroupByConstant::new()), Arc::new(OptimizeProjections::new()), @@ -307,7 +284,9 @@ impl TreeNodeRewriter for Rewriter<'_> { fn f_down(&mut self, node: LogicalPlan) -> Result> { if self.apply_order == ApplyOrder::TopDown { - optimize_plan_node(node, self.rule, self.config) + { + self.rule.rewrite(node, self.config) + } } else { Ok(Transformed::no(node)) } @@ -315,35 +294,15 @@ impl TreeNodeRewriter for Rewriter<'_> { fn f_up(&mut self, node: LogicalPlan) -> Result> { if self.apply_order == ApplyOrder::BottomUp { - optimize_plan_node(node, self.rule, self.config) + { + self.rule.rewrite(node, self.config) + } } else { Ok(Transformed::no(node)) } } } -/// Invokes the Optimizer rule to rewrite the LogicalPlan in place. -fn optimize_plan_node( - plan: LogicalPlan, - rule: &dyn OptimizerRule, - config: &dyn OptimizerConfig, -) -> Result> { - if rule.supports_rewrite() { - return rule.rewrite(plan, config); - } - - #[allow(deprecated)] - rule.try_optimize(&plan, config).map(|maybe_plan| { - match maybe_plan { - Some(new_plan) => { - // if the node was rewritten by the optimizer, replace the node - Transformed::yes(new_plan) - } - None => Transformed::no(plan), - } - }) -} - impl Optimizer { /// Optimizes the logical plan by applying optimizer rules, and /// invoking observer function after each call @@ -389,7 +348,9 @@ impl Optimizer { &mut Rewriter::new(apply_order, rule.as_ref(), config), ), // rule handles recursion itself - None => optimize_plan_node(new_plan, rule.as_ref(), config), + None => { + rule.rewrite(new_plan, config) + }, } .and_then(|tnr| { // run checks optimizer invariant checks, per optimizer rule applied diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index c38dd35abd36..0dbb78a2680e 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -285,6 +285,8 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result { | Expr::TryCast(_) | Expr::InList { .. } | Expr::ScalarFunction(_) => Ok(TreeNodeRecursion::Continue), + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::AggregateFunction(_) | Expr::WindowFunction(_) | Expr::Wildcard { .. } @@ -1138,6 +1140,12 @@ impl OptimizerRule for PushDownFilter { }) } LogicalPlan::Extension(extension_plan) => { + // This check prevents the Filter from being removed when the extension node has no children, + // so we return the original Filter unchanged. + if extension_plan.node.inputs().is_empty() { + filter.input = Arc::new(LogicalPlan::Extension(extension_plan)); + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } let prevent_cols = extension_plan.node.prevent_predicate_push_down_columns(); @@ -3784,4 +3792,83 @@ Projection: a, b \n TableScan: test"; assert_optimized_plan_eq(plan, expected_after) } + + #[test] + fn test_push_down_filter_to_user_defined_node() -> Result<()> { + // Define a custom user-defined logical node + #[derive(Debug, Hash, Eq, PartialEq)] + struct TestUserNode { + schema: DFSchemaRef, + } + + impl PartialOrd for TestUserNode { + fn partial_cmp(&self, _other: &Self) -> Option { + None + } + } + + impl TestUserNode { + fn new() -> Self { + let schema = Arc::new( + DFSchema::new_with_metadata( + vec![(None, Field::new("a", DataType::Int64, false).into())], + Default::default(), + ) + .unwrap(), + ); + + Self { schema } + } + } + + impl UserDefinedLogicalNodeCore for TestUserNode { + fn name(&self) -> &str { + "test_node" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "TestUserNode") + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + inputs: Vec, + ) -> Result { + assert!(exprs.is_empty()); + assert!(inputs.is_empty()); + Ok(Self { + schema: Arc::clone(&self.schema), + }) + } + } + + // Create a node and build a plan with a filter + let node = LogicalPlan::Extension(Extension { + node: Arc::new(TestUserNode::new()), + }); + + let plan = LogicalPlanBuilder::from(node).filter(lit(false))?.build()?; + + // Check the original plan format (not part of the test assertions) + let expected_before = "Filter: Boolean(false)\ + \n TestUserNode"; + assert_eq!(format!("{plan}"), expected_before); + + // Check that the filter is pushed down to the user-defined node + let expected_after = "Filter: Boolean(false)\n TestUserNode"; + assert_optimized_plan_eq(plan, expected_after) + } } diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index e43e2e704080..d5a1b84e6aff 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -32,7 +32,6 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; use datafusion_common::{internal_err, DFSchema, DataFusionError, Result, ScalarValue}; -use datafusion_expr::simplify::ExprSimplifyResult; use datafusion_expr::{ and, lit, or, BinaryExpr, Case, ColumnarValue, Expr, Like, Operator, Volatility, WindowFunctionDefinition, @@ -42,14 +41,23 @@ use datafusion_expr::{ expr::{InList, InSubquery, WindowFunction}, utils::{iter_conjunction, iter_conjunction_owned}, }; +use datafusion_expr::{simplify::ExprSimplifyResult, Cast, TryCast}; use datafusion_physical_expr::{create_physical_expr, execution_props::ExecutionProps}; use super::inlist_simplifier::ShortenInListSimplifier; use super::utils::*; -use crate::analyzer::type_coercion::TypeCoercionRewriter; use crate::simplify_expressions::guarantees::GuaranteeRewriter; use crate::simplify_expressions::regex::simplify_regex_expr; +use crate::simplify_expressions::unwrap_cast::{ + is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary, + is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist, + unwrap_cast_in_comparison_for_binary, +}; use crate::simplify_expressions::SimplifyInfo; +use crate::{ + analyzer::type_coercion::TypeCoercionRewriter, + simplify_expressions::unwrap_cast::try_cast_literal_to_type, +}; use indexmap::IndexSet; use regex::Regex; @@ -582,6 +590,8 @@ impl<'a> ConstEvaluator<'a> { // added they can be checked for their ability to be evaluated // at plan time match expr { + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] Expr::AggregateFunction { .. } | Expr::ScalarVariable(_, _) | Expr::Column(_) @@ -940,6 +950,27 @@ impl TreeNodeRewriter for Simplifier<'_, S> { op: And, right, }) if is_op_with(Or, &left, &right) => Transformed::yes(*right), + // A >= constant AND constant <= A --> A = constant + Expr::BinaryExpr(BinaryExpr { + left, + op: And, + right, + }) if can_reduce_to_equal_statement(&left, &right) => { + if let Expr::BinaryExpr(BinaryExpr { + left: left_left, + right: left_right, + .. + }) = *left + { + Transformed::yes(Expr::BinaryExpr(BinaryExpr { + left: left_left, + op: Eq, + right: left_right, + })) + } else { + return internal_err!("can_reduce_to_equal_statement should only be called with a BinaryExpr"); + } + } // // Rules for Multiply @@ -1719,6 +1750,86 @@ impl TreeNodeRewriter for Simplifier<'_, S> { } } + // ======================================= + // unwrap_cast_in_comparison + // ======================================= + // + // For case: + // try_cast/cast(expr as data_type) op literal + Expr::BinaryExpr(BinaryExpr { left, op, right }) + if is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary( + info, &left, &right, + ) && op.supports_propagation() => + { + unwrap_cast_in_comparison_for_binary(info, left, right, op)? + } + // literal op try_cast/cast(expr as data_type) + // --> + // try_cast/cast(expr as data_type) op_swap literal + Expr::BinaryExpr(BinaryExpr { left, op, right }) + if is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary( + info, &right, &left, + ) && op.supports_propagation() + && op.swap().is_some() => + { + unwrap_cast_in_comparison_for_binary( + info, + right, + left, + op.swap().unwrap(), + )? + } + // For case: + // try_cast/cast(expr as left_type) in (expr1,expr2,expr3) + Expr::InList(InList { + expr: mut left, + list, + negated, + }) if is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist( + info, &left, &list, + ) => + { + let (Expr::TryCast(TryCast { + expr: left_expr, .. + }) + | Expr::Cast(Cast { + expr: left_expr, .. + })) = left.as_mut() + else { + return internal_err!("Expect cast expr, but got {:?}", left)?; + }; + + let expr_type = info.get_data_type(left_expr)?; + let right_exprs = list + .into_iter() + .map(|right| { + match right { + Expr::Literal(right_lit_value) => { + // if the right_lit_value can be casted to the type of internal_left_expr + // we need to unwrap the cast for cast/try_cast expr, and add cast to the literal + let Some(value) = try_cast_literal_to_type(&right_lit_value, &expr_type) else { + internal_err!( + "Can't cast the list expr {:?} to type {:?}", + right_lit_value, &expr_type + )? + }; + Ok(lit(value)) + } + other_expr => internal_err!( + "Only support literal expr to optimize, but the expr is {:?}", + &other_expr + ), + } + }) + .collect::>>()?; + + Transformed::yes(Expr::InList(InList { + expr: std::mem::take(left_expr), + list: right_exprs, + negated, + })) + } + // no additional rewrites possible expr => Transformed::no(expr), }) diff --git a/datafusion/optimizer/src/simplify_expressions/mod.rs b/datafusion/optimizer/src/simplify_expressions/mod.rs index 46c066c11c0f..5fbee02e3909 100644 --- a/datafusion/optimizer/src/simplify_expressions/mod.rs +++ b/datafusion/optimizer/src/simplify_expressions/mod.rs @@ -23,6 +23,7 @@ mod guarantees; mod inlist_simplifier; mod regex; pub mod simplify_exprs; +mod unwrap_cast; mod utils; // backwards compatibility diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 6a56c1753328..709d8f79c3d9 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -62,8 +62,6 @@ impl OptimizerRule for SimplifyExpressions { true } - /// if supports_owned returns true, the Optimizer calls - /// [`Self::rewrite`] instead of [`Self::try_optimize`] fn rewrite( &self, plan: LogicalPlan, diff --git a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs similarity index 79% rename from datafusion/optimizer/src/unwrap_cast_in_comparison.rs rename to datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs index e2b8a966cb92..7670bdf98bb4 100644 --- a/datafusion/optimizer/src/unwrap_cast_in_comparison.rs +++ b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs @@ -15,274 +15,176 @@ // specific language governing permissions and limitations // under the License. -//! [`UnwrapCastInComparison`] rewrites `CAST(col) = lit` to `col = CAST(lit)` +//! Unwrap casts in binary comparisons +//! +//! The functions in this module attempt to remove casts from +//! comparisons to literals ([`ScalarValue`]s) by applying the casts +//! to the literals if possible. It is inspired by the optimizer rule +//! `UnwrapCastInBinaryComparison` of Spark. +//! +//! Removing casts often improves performance because: +//! 1. The cast is done once (to the literal) rather than to every value +//! 2. Can enable other optimizations such as predicate pushdown that +//! don't support casting +//! +//! The rule is applied to expressions of the following forms: +//! +//! 1. `cast(left_expr as data_type) comparison_op literal_expr` +//! 2. `literal_expr comparison_op cast(left_expr as data_type)` +//! 3. `cast(literal_expr) IN (expr1, expr2, ...)` +//! 4. `literal_expr IN (cast(expr1) , cast(expr2), ...)` +//! +//! If the expression matches one of the forms above, the rule will +//! ensure the value of `literal` is in range(min, max) of the +//! expr's data_type, and if the scalar is within range, the literal +//! will be casted to the data type of expr on the other side, and the +//! cast will be removed from the other side. +//! +//! # Example +//! +//! If the DataType of c1 is INT32. Given the filter +//! +//! ```text +//! cast(c1 as INT64) > INT64(10)` +//! ``` +//! +//! This rule will remove the cast and rewrite the expression to: +//! +//! ```text +//! c1 > INT32(10) +//! ``` +//! use std::cmp::Ordering; -use std::mem; -use std::sync::Arc; -use crate::optimizer::ApplyOrder; -use crate::{OptimizerConfig, OptimizerRule}; - -use crate::utils::NamePreserver; use arrow::datatypes::{ DataType, TimeUnit, MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION, }; use arrow::temporal_conversions::{MICROSECONDS, MILLISECONDS, NANOSECONDS}; -use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; -use datafusion_common::{internal_err, DFSchema, DFSchemaRef, Result, ScalarValue}; -use datafusion_expr::expr::{BinaryExpr, Cast, InList, TryCast}; -use datafusion_expr::utils::merge_schema; -use datafusion_expr::{lit, Expr, ExprSchemable, LogicalPlan}; - -/// [`UnwrapCastInComparison`] attempts to remove casts from -/// comparisons to literals ([`ScalarValue`]s) by applying the casts -/// to the literals if possible. It is inspired by the optimizer rule -/// `UnwrapCastInBinaryComparison` of Spark. -/// -/// Removing casts often improves performance because: -/// 1. The cast is done once (to the literal) rather than to every value -/// 2. Can enable other optimizations such as predicate pushdown that -/// don't support casting -/// -/// The rule is applied to expressions of the following forms: -/// -/// 1. `cast(left_expr as data_type) comparison_op literal_expr` -/// 2. `literal_expr comparison_op cast(left_expr as data_type)` -/// 3. `cast(literal_expr) IN (expr1, expr2, ...)` -/// 4. `literal_expr IN (cast(expr1) , cast(expr2), ...)` -/// -/// If the expression matches one of the forms above, the rule will -/// ensure the value of `literal` is in range(min, max) of the -/// expr's data_type, and if the scalar is within range, the literal -/// will be casted to the data type of expr on the other side, and the -/// cast will be removed from the other side. -/// -/// # Example -/// -/// If the DataType of c1 is INT32. Given the filter -/// -/// ```text -/// Filter: cast(c1 as INT64) > INT64(10)` -/// ``` -/// -/// This rule will remove the cast and rewrite the expression to: -/// -/// ```text -/// Filter: c1 > INT32(10) -/// ``` -/// -#[derive(Default, Debug)] -pub struct UnwrapCastInComparison {} - -impl UnwrapCastInComparison { - pub fn new() -> Self { - Self::default() +use datafusion_common::{internal_err, tree_node::Transformed}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{lit, BinaryExpr}; +use datafusion_expr::{simplify::SimplifyInfo, Cast, Expr, Operator, TryCast}; + +pub(super) fn unwrap_cast_in_comparison_for_binary( + info: &S, + cast_expr: Box, + literal: Box, + op: Operator, +) -> Result> { + match (*cast_expr, *literal) { + ( + Expr::TryCast(TryCast { expr, .. }) | Expr::Cast(Cast { expr, .. }), + Expr::Literal(lit_value), + ) => { + let Ok(expr_type) = info.get_data_type(&expr) else { + return internal_err!("Can't get the data type of the expr {:?}", &expr); + }; + // if the lit_value can be casted to the type of internal_left_expr + // we need to unwrap the cast for cast/try_cast expr, and add cast to the literal + let Some(value) = try_cast_literal_to_type(&lit_value, &expr_type) else { + return internal_err!( + "Can't cast the literal expr {:?} to type {:?}", + &lit_value, + &expr_type + ); + }; + Ok(Transformed::yes(Expr::BinaryExpr(BinaryExpr { + left: expr, + op, + right: Box::new(lit(value)), + }))) + } + _ => internal_err!("Expect cast expr and literal"), } } -impl OptimizerRule for UnwrapCastInComparison { - fn name(&self) -> &str { - "unwrap_cast_in_comparison" - } - - fn apply_order(&self) -> Option { - Some(ApplyOrder::BottomUp) - } +pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary< + S: SimplifyInfo, +>( + info: &S, + expr: &Expr, + literal: &Expr, +) -> bool { + match (expr, literal) { + ( + Expr::TryCast(TryCast { + expr: left_expr, .. + }) + | Expr::Cast(Cast { + expr: left_expr, .. + }), + Expr::Literal(lit_val), + ) => { + let Ok(expr_type) = info.get_data_type(left_expr) else { + return false; + }; - fn supports_rewrite(&self) -> bool { - true - } + let Ok(lit_type) = info.get_data_type(literal) else { + return false; + }; - fn rewrite( - &self, - plan: LogicalPlan, - _config: &dyn OptimizerConfig, - ) -> Result> { - let mut schema = merge_schema(&plan.inputs()); - - if let LogicalPlan::TableScan(ts) = &plan { - let source_schema = DFSchema::try_from_qualified_schema( - ts.table_name.clone(), - &ts.source.schema(), - )?; - schema.merge(&source_schema); + try_cast_literal_to_type(lit_val, &expr_type).is_some() + && is_supported_type(&expr_type) + && is_supported_type(&lit_type) } + _ => false, + } +} - schema.merge(plan.schema()); +pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist< + S: SimplifyInfo, +>( + info: &S, + expr: &Expr, + list: &[Expr], +) -> bool { + let (Expr::TryCast(TryCast { + expr: left_expr, .. + }) + | Expr::Cast(Cast { + expr: left_expr, .. + })) = expr + else { + return false; + }; - let mut expr_rewriter = UnwrapCastExprRewriter { - schema: Arc::new(schema), - }; + let Ok(expr_type) = info.get_data_type(left_expr) else { + return false; + }; - let name_preserver = NamePreserver::new(&plan); - plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr); - expr.rewrite(&mut expr_rewriter) - .map(|transformed| transformed.update_data(|e| original_name.restore(e))) - }) + if !is_supported_type(&expr_type) { + return false; } -} -struct UnwrapCastExprRewriter { - schema: DFSchemaRef, -} + for right in list { + let Ok(right_type) = info.get_data_type(right) else { + return false; + }; -impl TreeNodeRewriter for UnwrapCastExprRewriter { - type Node = Expr; - - fn f_up(&mut self, mut expr: Expr) -> Result> { - match &mut expr { - // For case: - // try_cast/cast(expr as data_type) op literal - // literal op try_cast/cast(expr as data_type) - Expr::BinaryExpr(BinaryExpr { left, op, right }) - if { - let Ok(left_type) = left.get_type(&self.schema) else { - return Ok(Transformed::no(expr)); - }; - let Ok(right_type) = right.get_type(&self.schema) else { - return Ok(Transformed::no(expr)); - }; - is_supported_type(&left_type) - && is_supported_type(&right_type) - && op.supports_propagation() - } => - { - match (left.as_mut(), right.as_mut()) { - ( - Expr::Literal(left_lit_value), - Expr::TryCast(TryCast { - expr: right_expr, .. - }) - | Expr::Cast(Cast { - expr: right_expr, .. - }), - ) => { - // if the left_lit_value can be cast to the type of expr - // we need to unwrap the cast for cast/try_cast expr, and add cast to the literal - let Ok(expr_type) = right_expr.get_type(&self.schema) else { - return Ok(Transformed::no(expr)); - }; - match expr_type { - // https://github.com/apache/datafusion/issues/12180 - DataType::Utf8View => Ok(Transformed::no(expr)), - _ => { - let Some(value) = - try_cast_literal_to_type(left_lit_value, &expr_type) - else { - return Ok(Transformed::no(expr)); - }; - **left = lit(value); - // unwrap the cast/try_cast for the right expr - **right = mem::take(right_expr); - Ok(Transformed::yes(expr)) - } - } - } - ( - Expr::TryCast(TryCast { - expr: left_expr, .. - }) - | Expr::Cast(Cast { - expr: left_expr, .. - }), - Expr::Literal(right_lit_value), - ) => { - // if the right_lit_value can be cast to the type of expr - // we need to unwrap the cast for cast/try_cast expr, and add cast to the literal - let Ok(expr_type) = left_expr.get_type(&self.schema) else { - return Ok(Transformed::no(expr)); - }; - match expr_type { - // https://github.com/apache/datafusion/issues/12180 - DataType::Utf8View => Ok(Transformed::no(expr)), - _ => { - let Some(value) = - try_cast_literal_to_type(right_lit_value, &expr_type) - else { - return Ok(Transformed::no(expr)); - }; - // unwrap the cast/try_cast for the left expr - **left = mem::take(left_expr); - **right = lit(value); - Ok(Transformed::yes(expr)) - } - } - } - _ => Ok(Transformed::no(expr)), - } - } - // For case: - // try_cast/cast(expr as left_type) in (expr1,expr2,expr3) - Expr::InList(InList { - expr: left, list, .. - }) => { - let (Expr::TryCast(TryCast { - expr: left_expr, .. - }) - | Expr::Cast(Cast { - expr: left_expr, .. - })) = left.as_mut() - else { - return Ok(Transformed::no(expr)); - }; - let Ok(expr_type) = left_expr.get_type(&self.schema) else { - return Ok(Transformed::no(expr)); - }; - if !is_supported_type(&expr_type) { - return Ok(Transformed::no(expr)); - } - let Ok(right_exprs) = list - .iter() - .map(|right| { - let right_type = right.get_type(&self.schema)?; - if !is_supported_type(&right_type) { - internal_err!( - "The type of list expr {} is not supported", - &right_type - )?; - } - match right { - Expr::Literal(right_lit_value) => { - // if the right_lit_value can be casted to the type of internal_left_expr - // we need to unwrap the cast for cast/try_cast expr, and add cast to the literal - let Some(value) = try_cast_literal_to_type(right_lit_value, &expr_type) else { - internal_err!( - "Can't cast the list expr {:?} to type {:?}", - right_lit_value, &expr_type - )? - }; - Ok(lit(value)) - } - other_expr => internal_err!( - "Only support literal expr to optimize, but the expr is {:?}", - &other_expr - ), - } - }) - .collect::>>() else { - return Ok(Transformed::no(expr)) - }; - **left = mem::take(left_expr); - *list = right_exprs; - Ok(Transformed::yes(expr)) - } - // TODO: handle other expr type and dfs visit them - _ => Ok(Transformed::no(expr)), + if !is_supported_type(&right_type) { + return false; + } + + match right { + Expr::Literal(lit_val) + if try_cast_literal_to_type(lit_val, &expr_type).is_some() => {} + _ => return false, } } + + true } -/// Returns true if [UnwrapCastExprRewriter] supports this data type +/// Returns true if unwrap_cast_in_comparison supports this data type fn is_supported_type(data_type: &DataType) -> bool { is_supported_numeric_type(data_type) || is_supported_string_type(data_type) || is_supported_dictionary_type(data_type) } -/// Returns true if [[UnwrapCastExprRewriter]] support this numeric type +/// Returns true if unwrap_cast_in_comparison support this numeric type fn is_supported_numeric_type(data_type: &DataType) -> bool { matches!( data_type, @@ -299,7 +201,7 @@ fn is_supported_numeric_type(data_type: &DataType) -> bool { ) } -/// Returns true if [UnwrapCastExprRewriter] supports casting this value as a string +/// Returns true if unwrap_cast_in_comparison supports casting this value as a string fn is_supported_string_type(data_type: &DataType) -> bool { matches!( data_type, @@ -307,14 +209,14 @@ fn is_supported_string_type(data_type: &DataType) -> bool { ) } -/// Returns true if [UnwrapCastExprRewriter] supports casting this value as a dictionary +/// Returns true if unwrap_cast_in_comparison supports casting this value as a dictionary fn is_supported_dictionary_type(data_type: &DataType) -> bool { matches!(data_type, DataType::Dictionary(_, inner) if is_supported_type(inner)) } /// Convert a literal value from one data type to another -fn try_cast_literal_to_type( +pub(super) fn try_cast_literal_to_type( lit_value: &ScalarValue, target_type: &DataType, ) -> Option { @@ -540,13 +442,16 @@ fn cast_between_timestamp(from: &DataType, to: &DataType, value: i128) -> Option #[cfg(test)] mod tests { - use std::collections::HashMap; - use super::*; + use std::collections::HashMap; + use std::sync::Arc; + use crate::simplify_expressions::ExprSimplifier; use arrow::compute::{cast_with_options, CastOptions}; use arrow::datatypes::Field; - use datafusion_common::tree_node::TransformedResult; + use datafusion_common::{DFSchema, DFSchemaRef}; + use datafusion_expr::execution_props::ExecutionProps; + use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::{cast, col, in_list, try_cast}; #[test] @@ -587,9 +492,9 @@ mod tests { let expected = col("c1").lt(null_i32()); assert_eq!(optimize_test(c1_lt_lit_null, &schema), expected); - // cast(INT8(NULL), INT32) < INT32(12) => INT8(NULL) < INT8(12) + // cast(INT8(NULL), INT32) < INT32(12) => INT8(NULL) < INT8(12) => BOOL(NULL) let lit_lt_lit = cast(null_i8(), DataType::Int32).lt(lit(12i32)); - let expected = null_i8().lt(lit(12i8)); + let expected = null_bool(); assert_eq!(optimize_test(lit_lt_lit, &schema), expected); } @@ -623,7 +528,7 @@ mod tests { // Verify reversed argument order // arrow_cast('value', 'Dictionary') = cast(str1 as Dictionary) => Utf8('value1') = str1 let expr_input = lit(dict.clone()).eq(cast(col("str1"), dict.data_type())); - let expected = lit("value").eq(col("str1")); + let expected = col("str1").eq(lit("value")); assert_eq!(optimize_test(expr_input, &schema), expected); } @@ -740,15 +645,27 @@ mod tests { #[test] fn test_unwrap_list_cast_comparison() { let schema = expr_test_schema(); - // INT32(C1) IN (INT32(12),INT64(24)) -> INT32(C1) IN (INT32(12),INT32(24)) - let expr_lt = - cast(col("c1"), DataType::Int64).in_list(vec![lit(12i64), lit(24i64)], false); - let expected = col("c1").in_list(vec![lit(12i32), lit(24i32)], false); + // INT32(C1) IN (INT32(12),INT64(23),INT64(34),INT64(56),INT64(78)) -> + // INT32(C1) IN (INT32(12),INT32(23),INT32(34),INT32(56),INT32(78)) + let expr_lt = cast(col("c1"), DataType::Int64).in_list( + vec![lit(12i64), lit(23i64), lit(34i64), lit(56i64), lit(78i64)], + false, + ); + let expected = col("c1").in_list( + vec![lit(12i32), lit(23i32), lit(34i32), lit(56i32), lit(78i32)], + false, + ); assert_eq!(optimize_test(expr_lt, &schema), expected); - // INT32(C2) IN (INT64(NULL),INT64(24)) -> INT32(C1) IN (INT32(12),INT32(24)) - let expr_lt = - cast(col("c2"), DataType::Int32).in_list(vec![null_i32(), lit(14i32)], false); - let expected = col("c2").in_list(vec![null_i64(), lit(14i64)], false); + // INT32(C2) IN (INT64(NULL),INT64(24),INT64(34),INT64(56),INT64(78)) -> + // INT32(C2) IN (INT32(NULL),INT32(24),INT32(34),INT32(56),INT32(78)) + let expr_lt = cast(col("c2"), DataType::Int32).in_list( + vec![null_i32(), lit(24i32), lit(34i64), lit(56i64), lit(78i64)], + false, + ); + let expected = col("c2").in_list( + vec![null_i64(), lit(24i64), lit(34i64), lit(56i64), lit(78i64)], + false, + ); assert_eq!(optimize_test(expr_lt, &schema), expected); @@ -774,10 +691,14 @@ mod tests { ); assert_eq!(optimize_test(expr_lt, &schema), expected); - // cast(INT32(12), INT64) IN (.....) - let expr_lt = cast(lit(12i32), DataType::Int64) - .in_list(vec![lit(13i64), lit(12i64)], false); - let expected = lit(12i32).in_list(vec![lit(13i32), lit(12i32)], false); + // cast(INT32(12), INT64) IN (.....) => + // INT64(12) IN (INT64(12),INT64(13),INT64(14),INT64(15),INT64(16)) + // => true + let expr_lt = cast(lit(12i32), DataType::Int64).in_list( + vec![lit(12i64), lit(13i64), lit(14i64), lit(15i64), lit(16i64)], + false, + ); + let expected = lit(true); assert_eq!(optimize_test(expr_lt, &schema), expected); } @@ -815,8 +736,12 @@ mod tests { assert_eq!(optimize_test(expr_input.clone(), &schema), expr_input); // inlist for unsupported data type - let expr_input = - in_list(cast(col("c6"), DataType::Float64), vec![lit(0f64)], false); + let expr_input = in_list( + cast(col("c6"), DataType::Float64), + // need more literals to avoid rewriting to binary expr + vec![lit(0f64), lit(1f64), lit(2f64), lit(3f64), lit(4f64)], + false, + ); assert_eq!(optimize_test(expr_input.clone(), &schema), expr_input); } @@ -833,10 +758,12 @@ mod tests { } fn optimize_test(expr: Expr, schema: &DFSchemaRef) -> Expr { - let mut expr_rewriter = UnwrapCastExprRewriter { - schema: Arc::clone(schema), - }; - expr.rewrite(&mut expr_rewriter).data().unwrap() + let props = ExecutionProps::new(); + let simplifier = ExprSimplifier::new( + SimplifyContext::new(&props).with_schema(Arc::clone(schema)), + ); + + simplifier.simplify(expr).unwrap() } fn expr_test_schema() -> DFSchemaRef { @@ -862,6 +789,10 @@ mod tests { ) } + fn null_bool() -> Expr { + lit(ScalarValue::Boolean(None)) + } + fn null_i8() -> Expr { lit(ScalarValue::Int8(None)) } diff --git a/datafusion/optimizer/src/simplify_expressions/utils.rs b/datafusion/optimizer/src/simplify_expressions/utils.rs index c30c3631c193..cf182175e48e 100644 --- a/datafusion/optimizer/src/simplify_expressions/utils.rs +++ b/datafusion/optimizer/src/simplify_expressions/utils.rs @@ -214,6 +214,25 @@ pub fn is_op_with(target_op: Operator, haystack: &Expr, needle: &Expr) -> bool { matches!(haystack, Expr::BinaryExpr(BinaryExpr { left, op, right }) if op == &target_op && (needle == left.as_ref() || needle == right.as_ref()) && !needle.is_volatile()) } +pub fn can_reduce_to_equal_statement(haystack: &Expr, needle: &Expr) -> bool { + match (haystack, needle) { + // a >= constant and constant <= a => a = constant + ( + Expr::BinaryExpr(BinaryExpr { + left, + op: Operator::GtEq, + right, + }), + Expr::BinaryExpr(BinaryExpr { + left: n_left, + op: Operator::LtEq, + right: n_right, + }), + ) if left == n_left && right == n_right => true, + _ => false, + } +} + /// returns true if `not_expr` is !`expr` (not) pub fn is_not_of(not_expr: &Expr, expr: &Expr) -> bool { matches!(not_expr, Expr::Not(inner) if expr == inner.as_ref()) diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index 39f8cf285d17..c734d908f6d6 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -19,8 +19,6 @@ use std::collections::{BTreeSet, HashMap, HashSet}; -use crate::{OptimizerConfig, OptimizerRule}; - use crate::analyzer::type_coercion::TypeCoercionRewriter; use arrow::array::{new_null_array, Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; @@ -38,44 +36,6 @@ use std::sync::Arc; /// as it was initially placed here and then moved elsewhere. pub use datafusion_expr::expr_rewriter::NamePreserver; -/// Convenience rule for writing optimizers: recursively invoke -/// optimize on plan's children and then return a node of the same -/// type. Useful for optimizer rules which want to leave the type -/// of plan unchanged but still apply to the children. -/// This also handles the case when the `plan` is a [`LogicalPlan::Explain`]. -/// -/// Returning `Ok(None)` indicates that the plan can't be optimized by the `optimizer`. -#[deprecated( - since = "40.0.0", - note = "please use OptimizerRule::apply_order with ApplyOrder::BottomUp instead" -)] -pub fn optimize_children( - optimizer: &impl OptimizerRule, - plan: &LogicalPlan, - config: &dyn OptimizerConfig, -) -> Result> { - let mut new_inputs = Vec::with_capacity(plan.inputs().len()); - let mut plan_is_changed = false; - for input in plan.inputs() { - if optimizer.supports_rewrite() { - let new_input = optimizer.rewrite(input.clone(), config)?; - plan_is_changed = plan_is_changed || new_input.transformed; - new_inputs.push(new_input.data); - } else { - #[allow(deprecated)] - let new_input = optimizer.try_optimize(input, config)?; - plan_is_changed = plan_is_changed || new_input.is_some(); - new_inputs.push(new_input.unwrap_or_else(|| input.clone())) - } - } - if plan_is_changed { - let exprs = plan.expressions(); - plan.with_new_exprs(exprs, new_inputs).map(Some) - } else { - Ok(None) - } -} - /// Returns true if `expr` contains all columns in `schema_cols` pub(crate) fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet) -> bool { let column_refs = expr.column_refs(); diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index 66bd6b75123e..5e66c7ec0313 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -22,16 +22,14 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{assert_contains, plan_err, Result, TableReference}; +use datafusion_common::{plan_err, Result, TableReference}; use datafusion_expr::planner::ExprPlanner; -use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; use datafusion_expr::test::function_stub::sum_udaf; use datafusion_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF}; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::planner::AggregateFunctionPlanner; use datafusion_functions_window::planner::WindowFunctionPlanner; -use datafusion_optimizer::analyzer::type_coercion::TypeCoercionRewriter; use datafusion_optimizer::analyzer::Analyzer; use datafusion_optimizer::optimizer::Optimizer; use datafusion_optimizer::{OptimizerConfig, OptimizerContext, OptimizerRule}; @@ -344,16 +342,6 @@ fn test_propagate_empty_relation_inner_join_and_unions() { assert_eq!(expected, format!("{plan}")); } -#[test] -fn select_wildcard_with_repeated_column() { - let sql = "SELECT *, col_int32 FROM test"; - let err = test_sql(sql).expect_err("query should have failed"); - assert_eq!( - "Schema error: Schema contains duplicate qualified field name test.col_int32", - err.strip_backtrace() - ); -} - #[test] fn select_wildcard_with_repeated_column_but_is_aliased() { let sql = "SELECT *, col_int32 as col_32 FROM test"; @@ -390,32 +378,6 @@ fn select_correlated_predicate_subquery_with_uppercase_ident() { assert_eq!(expected, format!("{plan}")); } -// The test should return an error -// because the wildcard didn't be expanded before type coercion -#[test] -fn test_union_coercion_with_wildcard() -> Result<()> { - let dialect = PostgreSqlDialect {}; - let context_provider = MyContextProvider::default(); - let sql = "select * from (SELECT col_int32, col_uint32 FROM test) union all select * from(SELECT col_uint32, col_int32 FROM test)"; - let statements = Parser::parse_sql(&dialect, sql)?; - let sql_to_rel = SqlToRel::new(&context_provider); - let logical_plan = sql_to_rel.sql_statement_to_plan(statements[0].clone())?; - - if let LogicalPlan::Union(union) = logical_plan { - let err = TypeCoercionRewriter::coerce_union(union) - .err() - .unwrap() - .to_string(); - assert_contains!( - err, - "Error during planning: Wildcard should be expanded before type coercion" - ); - } else { - panic!("Expected Union plan"); - } - Ok(()) -} - fn test_sql(sql: &str) -> Result { // parse the SQL let dialect = GenericDialect {}; // or AnsiDialect, or your own dialect ... diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index 07a98340dbe7..34c4e52d517e 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -91,6 +91,9 @@ impl AggregateExprBuilder { } } + /// Constructs an `AggregateFunctionExpr` from the builder + /// + /// Note that an [`Self::alias`] must be provided before calling this method. pub fn build(self) -> Result { let Self { fun, @@ -132,7 +135,11 @@ impl AggregateExprBuilder { let data_type = fun.return_type(&input_exprs_types)?; let is_nullable = fun.is_nullable(); let name = match alias { - None => return internal_err!("alias should be provided"), + None => { + return internal_err!( + "AggregateExprBuilder::alias must be provided prior to calling build" + ) + } Some(alias) => alias, }; @@ -199,6 +206,8 @@ impl AggregateExprBuilder { } /// Physical aggregate expression of a UDAF. +/// +/// Instances are constructed via [`AggregateExprBuilder`]. #[derive(Debug, Clone)] pub struct AggregateFunctionExpr { fun: AggregateUDF, diff --git a/datafusion/physical-expr/src/equivalence/mod.rs b/datafusion/physical-expr/src/equivalence/mod.rs index fcc1c564d8c8..e94d2bad5712 100644 --- a/datafusion/physical-expr/src/equivalence/mod.rs +++ b/datafusion/physical-expr/src/equivalence/mod.rs @@ -79,6 +79,38 @@ mod tests { LexOrdering, PhysicalSortRequirement, }; + /// Converts a string to a physical sort expression + /// + /// # Example + /// * `"a"` -> (`"a"`, `SortOptions::default()`) + /// * `"a ASC"` -> (`"a"`, `SortOptions { descending: false, nulls_first: false }`) + pub fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr { + let mut parts = name.split_whitespace(); + let name = parts.next().expect("empty sort expression"); + let mut sort_expr = PhysicalSortExpr::new( + col(name, schema).expect("invalid column name"), + SortOptions::default(), + ); + + if let Some(options) = parts.next() { + sort_expr = match options { + "ASC" => sort_expr.asc(), + "DESC" => sort_expr.desc(), + _ => panic!( + "unknown sort options. Expected 'ASC' or 'DESC', got {}", + options + ), + } + } + + assert!( + parts.next().is_none(), + "unexpected tokens in column name. Expected 'name' / 'name ASC' / 'name DESC' but got '{name}'" + ); + + sort_expr + } + pub fn output_schema( mapping: &ProjectionMapping, input_schema: &Arc, diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs deleted file mode 100755 index 042256951250..000000000000 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ /dev/null @@ -1,4544 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::fmt::Display; -use std::hash::{Hash, Hasher}; -use std::iter::Peekable; -use std::slice::Iter; -use std::sync::Arc; -use std::{fmt, mem}; - -use crate::equivalence::class::{const_exprs_contains, AcrossPartitions}; -use crate::equivalence::{ - EquivalenceClass, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping, -}; -use crate::expressions::{with_new_schema, CastExpr, Column, Literal}; -use crate::{ - physical_exprs_contains, ConstExpr, LexOrdering, LexRequirement, PhysicalExpr, - PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequirement, -}; - -use arrow::compute::SortOptions; -use arrow::datatypes::SchemaRef; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{ - internal_err, plan_err, Constraint, Constraints, HashMap, JoinSide, JoinType, Result, -}; -use datafusion_expr::interval_arithmetic::Interval; -use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; -use datafusion_physical_expr_common::utils::ExprPropertiesNode; - -use indexmap::{IndexMap, IndexSet}; -use itertools::Itertools; - -/// A `EquivalenceProperties` object stores information known about the output -/// of a plan node, that can be used to optimize the plan. -/// -/// Currently, it keeps track of: -/// - Sort expressions (orderings) -/// - Equivalent expressions: expressions that are known to have same value. -/// - Constants expressions: expressions that are known to contain a single -/// constant value. -/// -/// # Example equivalent sort expressions -/// -/// Consider table below: -/// -/// ```text -/// ┌-------┐ -/// | a | b | -/// |---|---| -/// | 1 | 9 | -/// | 2 | 8 | -/// | 3 | 7 | -/// | 5 | 5 | -/// └---┴---┘ -/// ``` -/// -/// In this case, both `a ASC` and `b DESC` can describe the table ordering. -/// `EquivalenceProperties`, tracks these different valid sort expressions and -/// treat `a ASC` and `b DESC` on an equal footing. For example if the query -/// specifies the output sorted by EITHER `a ASC` or `b DESC`, the sort can be -/// avoided. -/// -/// # Example equivalent expressions -/// -/// Similarly, consider the table below: -/// -/// ```text -/// ┌-------┐ -/// | a | b | -/// |---|---| -/// | 1 | 1 | -/// | 2 | 2 | -/// | 3 | 3 | -/// | 5 | 5 | -/// └---┴---┘ -/// ``` -/// -/// In this case, columns `a` and `b` always have the same value, which can of -/// such equivalences inside this object. With this information, Datafusion can -/// optimize operations such as. For example, if the partition requirement is -/// `Hash(a)` and output partitioning is `Hash(b)`, then DataFusion avoids -/// repartitioning the data as the existing partitioning satisfies the -/// requirement. -/// -/// # Code Example -/// ``` -/// # use std::sync::Arc; -/// # use arrow::datatypes::{Schema, Field, DataType, SchemaRef}; -/// # use datafusion_physical_expr::{ConstExpr, EquivalenceProperties}; -/// # use datafusion_physical_expr::expressions::col; -/// use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; -/// # let schema: SchemaRef = Arc::new(Schema::new(vec![ -/// # Field::new("a", DataType::Int32, false), -/// # Field::new("b", DataType::Int32, false), -/// # Field::new("c", DataType::Int32, false), -/// # ])); -/// # let col_a = col("a", &schema).unwrap(); -/// # let col_b = col("b", &schema).unwrap(); -/// # let col_c = col("c", &schema).unwrap(); -/// // This object represents data that is sorted by a ASC, c DESC -/// // with a single constant value of b -/// let mut eq_properties = EquivalenceProperties::new(schema) -/// .with_constants(vec![ConstExpr::from(col_b)]); -/// eq_properties.add_new_ordering(LexOrdering::new(vec![ -/// PhysicalSortExpr::new_default(col_a).asc(), -/// PhysicalSortExpr::new_default(col_c).desc(), -/// ])); -/// -/// assert_eq!(eq_properties.to_string(), "order: [[a@0 ASC, c@2 DESC]], const: [b@1(heterogeneous)]") -/// ``` -#[derive(Debug, Clone)] -pub struct EquivalenceProperties { - /// Distinct equivalence classes (exprs known to have the same expressions) - eq_group: EquivalenceGroup, - /// Equivalent sort expressions - oeq_class: OrderingEquivalenceClass, - /// Expressions whose values are constant - /// - /// TODO: We do not need to track constants separately, they can be tracked - /// inside `eq_group` as `Literal` expressions. - constants: Vec, - /// Table constraints - constraints: Constraints, - /// Schema associated with this object. - schema: SchemaRef, -} - -impl EquivalenceProperties { - /// Creates an empty `EquivalenceProperties` object. - pub fn new(schema: SchemaRef) -> Self { - Self { - eq_group: EquivalenceGroup::empty(), - oeq_class: OrderingEquivalenceClass::empty(), - constants: vec![], - constraints: Constraints::empty(), - schema, - } - } - - /// Adds constraints to the properties. - pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.constraints = constraints; - self - } - - /// Creates a new `EquivalenceProperties` object with the given orderings. - pub fn new_with_orderings(schema: SchemaRef, orderings: &[LexOrdering]) -> Self { - Self { - eq_group: EquivalenceGroup::empty(), - oeq_class: OrderingEquivalenceClass::new(orderings.to_vec()), - constants: vec![], - constraints: Constraints::empty(), - schema, - } - } - - /// Returns the associated schema. - pub fn schema(&self) -> &SchemaRef { - &self.schema - } - - /// Returns a reference to the ordering equivalence class within. - pub fn oeq_class(&self) -> &OrderingEquivalenceClass { - &self.oeq_class - } - - /// Return the inner OrderingEquivalenceClass, consuming self - pub fn into_oeq_class(self) -> OrderingEquivalenceClass { - self.oeq_class - } - - /// Returns a reference to the equivalence group within. - pub fn eq_group(&self) -> &EquivalenceGroup { - &self.eq_group - } - - /// Returns a reference to the constant expressions - pub fn constants(&self) -> &[ConstExpr] { - &self.constants - } - - pub fn constraints(&self) -> &Constraints { - &self.constraints - } - - /// Returns the output ordering of the properties. - pub fn output_ordering(&self) -> Option { - let constants = self.constants(); - let mut output_ordering = self.oeq_class().output_ordering().unwrap_or_default(); - // Prune out constant expressions - output_ordering - .retain(|sort_expr| !const_exprs_contains(constants, &sort_expr.expr)); - (!output_ordering.is_empty()).then_some(output_ordering) - } - - /// Returns the normalized version of the ordering equivalence class within. - /// Normalization removes constants and duplicates as well as standardizing - /// expressions according to the equivalence group within. - pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass { - OrderingEquivalenceClass::new( - self.oeq_class - .iter() - .map(|ordering| self.normalize_sort_exprs(ordering)) - .collect(), - ) - } - - /// Extends this `EquivalenceProperties` with the `other` object. - pub fn extend(mut self, other: Self) -> Self { - self.eq_group.extend(other.eq_group); - self.oeq_class.extend(other.oeq_class); - self.with_constants(other.constants) - } - - /// Clears (empties) the ordering equivalence class within this object. - /// Call this method when existing orderings are invalidated. - pub fn clear_orderings(&mut self) { - self.oeq_class.clear(); - } - - /// Removes constant expressions that may change across partitions. - /// This method should be used when data from different partitions are merged. - pub fn clear_per_partition_constants(&mut self) { - self.constants.retain(|item| { - matches!(item.across_partitions(), AcrossPartitions::Uniform(_)) - }) - } - - /// Extends this `EquivalenceProperties` by adding the orderings inside the - /// ordering equivalence class `other`. - pub fn add_ordering_equivalence_class(&mut self, other: OrderingEquivalenceClass) { - self.oeq_class.extend(other); - } - - /// Adds new orderings into the existing ordering equivalence class. - pub fn add_new_orderings( - &mut self, - orderings: impl IntoIterator, - ) { - self.oeq_class.add_new_orderings(orderings); - } - - /// Adds a single ordering to the existing ordering equivalence class. - pub fn add_new_ordering(&mut self, ordering: LexOrdering) { - self.add_new_orderings([ordering]); - } - - /// Incorporates the given equivalence group to into the existing - /// equivalence group within. - pub fn add_equivalence_group(&mut self, other_eq_group: EquivalenceGroup) { - self.eq_group.extend(other_eq_group); - } - - /// Adds a new equality condition into the existing equivalence group. - /// If the given equality defines a new equivalence class, adds this new - /// equivalence class to the equivalence group. - pub fn add_equal_conditions( - &mut self, - left: &Arc, - right: &Arc, - ) -> Result<()> { - // Discover new constants in light of new the equality: - if self.is_expr_constant(left) { - // Left expression is constant, add right as constant - if !const_exprs_contains(&self.constants, right) { - let const_expr = ConstExpr::from(right) - .with_across_partitions(self.get_expr_constant_value(left)); - self.constants.push(const_expr); - } - } else if self.is_expr_constant(right) { - // Right expression is constant, add left as constant - if !const_exprs_contains(&self.constants, left) { - let const_expr = ConstExpr::from(left) - .with_across_partitions(self.get_expr_constant_value(right)); - self.constants.push(const_expr); - } - } - - // Add equal expressions to the state - self.eq_group.add_equal_conditions(left, right); - - // Discover any new orderings - self.discover_new_orderings(left)?; - Ok(()) - } - - /// Track/register physical expressions with constant values. - #[deprecated(since = "43.0.0", note = "Use [`with_constants`] instead")] - pub fn add_constants(self, constants: impl IntoIterator) -> Self { - self.with_constants(constants) - } - - /// Remove the specified constant - pub fn remove_constant(mut self, c: &ConstExpr) -> Self { - self.constants.retain(|existing| existing != c); - self - } - - /// Track/register physical expressions with constant values. - pub fn with_constants( - mut self, - constants: impl IntoIterator, - ) -> Self { - let normalized_constants = constants - .into_iter() - .filter_map(|c| { - let across_partitions = c.across_partitions(); - let expr = c.owned_expr(); - let normalized_expr = self.eq_group.normalize_expr(expr); - - if const_exprs_contains(&self.constants, &normalized_expr) { - return None; - } - - let const_expr = ConstExpr::from(normalized_expr) - .with_across_partitions(across_partitions); - - Some(const_expr) - }) - .collect::>(); - - // Add all new normalized constants - self.constants.extend(normalized_constants); - - // Discover any new orderings based on the constants - for ordering in self.normalized_oeq_class().iter() { - if let Err(e) = self.discover_new_orderings(&ordering[0].expr) { - log::debug!("error discovering new orderings: {e}"); - } - } - - self - } - - // Discover new valid orderings in light of a new equality. - // Accepts a single argument (`expr`) which is used to determine - // which orderings should be updated. - // When constants or equivalence classes are changed, there may be new orderings - // that can be discovered with the new equivalence properties. - // For a discussion, see: https://github.com/apache/datafusion/issues/9812 - fn discover_new_orderings(&mut self, expr: &Arc) -> Result<()> { - let normalized_expr = self.eq_group().normalize_expr(Arc::clone(expr)); - let eq_class = self - .eq_group - .iter() - .find_map(|class| { - class - .contains(&normalized_expr) - .then(|| class.clone().into_vec()) - }) - .unwrap_or_else(|| vec![Arc::clone(&normalized_expr)]); - - let mut new_orderings: Vec = vec![]; - for ordering in self.normalized_oeq_class().iter() { - if !ordering[0].expr.eq(&normalized_expr) { - continue; - } - - let leading_ordering_options = ordering[0].options; - - for equivalent_expr in &eq_class { - let children = equivalent_expr.children(); - if children.is_empty() { - continue; - } - - // Check if all children match the next expressions in the ordering - let mut all_children_match = true; - let mut child_properties = vec![]; - - // Build properties for each child based on the next expressions - for (i, child) in children.iter().enumerate() { - if let Some(next) = ordering.get(i + 1) { - if !child.as_ref().eq(next.expr.as_ref()) { - all_children_match = false; - break; - } - child_properties.push(ExprProperties { - sort_properties: SortProperties::Ordered(next.options), - range: Interval::make_unbounded( - &child.data_type(&self.schema)?, - )?, - preserves_lex_ordering: true, - }); - } else { - all_children_match = false; - break; - } - } - - if all_children_match { - // Check if the expression is monotonic in all arguments - if let Ok(expr_properties) = - equivalent_expr.get_properties(&child_properties) - { - if expr_properties.preserves_lex_ordering - && SortProperties::Ordered(leading_ordering_options) - == expr_properties.sort_properties - { - // Assume existing ordering is [c ASC, a ASC, b ASC] - // When equality c = f(a,b) is given, if we know that given ordering `[a ASC, b ASC]`, - // ordering `[f(a,b) ASC]` is valid, then we can deduce that ordering `[a ASC, b ASC]` is also valid. - // Hence, ordering `[a ASC, b ASC]` can be added to the state as a valid ordering. - // (e.g. existing ordering where leading ordering is removed) - new_orderings.push(LexOrdering::new(ordering[1..].to_vec())); - break; - } - } - } - } - } - - self.oeq_class.add_new_orderings(new_orderings); - Ok(()) - } - - /// Updates the ordering equivalence group within assuming that the table - /// is re-sorted according to the argument `sort_exprs`. Note that constants - /// and equivalence classes are unchanged as they are unaffected by a re-sort. - /// If the given ordering is already satisfied, the function does nothing. - pub fn with_reorder(mut self, sort_exprs: LexOrdering) -> Self { - // Filter out constant expressions as they don't affect ordering - let filtered_exprs = LexOrdering::new( - sort_exprs - .into_iter() - .filter(|expr| !self.is_expr_constant(&expr.expr)) - .collect(), - ); - - if filtered_exprs.is_empty() { - return self; - } - - let mut new_orderings = vec![filtered_exprs.clone()]; - - // Preserve valid suffixes from existing orderings - let oeq_class = mem::take(&mut self.oeq_class); - for existing in oeq_class { - if self.is_prefix_of(&filtered_exprs, &existing) { - let mut extended = filtered_exprs.clone(); - extended.extend(existing.into_iter().skip(filtered_exprs.len())); - new_orderings.push(extended); - } - } - - self.oeq_class = OrderingEquivalenceClass::new(new_orderings); - self - } - - /// Checks if the new ordering matches a prefix of the existing ordering - /// (considering expression equivalences) - fn is_prefix_of(&self, new_order: &LexOrdering, existing: &LexOrdering) -> bool { - // Check if new order is longer than existing - can't be a prefix - if new_order.len() > existing.len() { - return false; - } - - // Check if new order matches existing prefix (considering equivalences) - new_order.iter().zip(existing).all(|(new, existing)| { - self.eq_group.exprs_equal(&new.expr, &existing.expr) - && new.options == existing.options - }) - } - - /// Normalizes the given sort expressions (i.e. `sort_exprs`) using the - /// equivalence group and the ordering equivalence class within. - /// - /// Assume that `self.eq_group` states column `a` and `b` are aliases. - /// Also assume that `self.oeq_class` states orderings `d ASC` and `a ASC, c ASC` - /// are equivalent (in the sense that both describe the ordering of the table). - /// If the `sort_exprs` argument were `vec![b ASC, c ASC, a ASC]`, then this - /// function would return `vec![a ASC, c ASC]`. Internally, it would first - /// normalize to `vec![a ASC, c ASC, a ASC]` and end up with the final result - /// after deduplication. - fn normalize_sort_exprs(&self, sort_exprs: &LexOrdering) -> LexOrdering { - // Convert sort expressions to sort requirements: - let sort_reqs = LexRequirement::from(sort_exprs.clone()); - // Normalize the requirements: - let normalized_sort_reqs = self.normalize_sort_requirements(&sort_reqs); - // Convert sort requirements back to sort expressions: - LexOrdering::from(normalized_sort_reqs) - } - - /// Normalizes the given sort requirements (i.e. `sort_reqs`) using the - /// equivalence group and the ordering equivalence class within. It works by: - /// - Removing expressions that have a constant value from the given requirement. - /// - Replacing sections that belong to some equivalence class in the equivalence - /// group with the first entry in the matching equivalence class. - /// - /// Assume that `self.eq_group` states column `a` and `b` are aliases. - /// Also assume that `self.oeq_class` states orderings `d ASC` and `a ASC, c ASC` - /// are equivalent (in the sense that both describe the ordering of the table). - /// If the `sort_reqs` argument were `vec![b ASC, c ASC, a ASC]`, then this - /// function would return `vec![a ASC, c ASC]`. Internally, it would first - /// normalize to `vec![a ASC, c ASC, a ASC]` and end up with the final result - /// after deduplication. - fn normalize_sort_requirements(&self, sort_reqs: &LexRequirement) -> LexRequirement { - let normalized_sort_reqs = self.eq_group.normalize_sort_requirements(sort_reqs); - let mut constant_exprs = vec![]; - constant_exprs.extend( - self.constants - .iter() - .map(|const_expr| Arc::clone(const_expr.expr())), - ); - let constants_normalized = self.eq_group.normalize_exprs(constant_exprs); - // Prune redundant sections in the requirement: - normalized_sort_reqs - .iter() - .filter(|&order| !physical_exprs_contains(&constants_normalized, &order.expr)) - .cloned() - .collect::() - .collapse() - } - - /// Checks whether the given ordering is satisfied by any of the existing - /// orderings. - pub fn ordering_satisfy(&self, given: &LexOrdering) -> bool { - // Convert the given sort expressions to sort requirements: - let sort_requirements = LexRequirement::from(given.clone()); - self.ordering_satisfy_requirement(&sort_requirements) - } - - /// Checks whether the given sort requirements are satisfied by any of the - /// existing orderings. - pub fn ordering_satisfy_requirement(&self, reqs: &LexRequirement) -> bool { - let mut eq_properties = self.clone(); - // First, standardize the given requirement: - let normalized_reqs = eq_properties.normalize_sort_requirements(reqs); - - // Check whether given ordering is satisfied by constraints first - if self.satisfied_by_constraints(&normalized_reqs) { - return true; - } - - for normalized_req in normalized_reqs { - // Check whether given ordering is satisfied - if !eq_properties.ordering_satisfy_single(&normalized_req) { - return false; - } - // Treat satisfied keys as constants in subsequent iterations. We - // can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - // - // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - // From the analysis above, we know that `[a ASC]` is satisfied. Then, - // we add column `a` as constant to the algorithm state. This enables us - // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. - eq_properties = eq_properties - .with_constants(std::iter::once(ConstExpr::from(normalized_req.expr))); - } - true - } - - /// Checks if the sort requirements are satisfied by any of the table constraints (primary key or unique). - /// Returns true if any constraint fully satisfies the requirements. - fn satisfied_by_constraints( - &self, - normalized_reqs: &[PhysicalSortRequirement], - ) -> bool { - self.constraints.iter().any(|constraint| match constraint { - Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => self - .satisfied_by_constraint( - normalized_reqs, - indices, - matches!(constraint, Constraint::Unique(_)), - ), - }) - } - - /// Checks if sort requirements are satisfied by a constraint (primary key or unique). - /// Returns true if the constraint indices form a valid prefix of an existing ordering - /// that matches the requirements. For unique constraints, also verifies nullable columns. - fn satisfied_by_constraint( - &self, - normalized_reqs: &[PhysicalSortRequirement], - indices: &[usize], - check_null: bool, - ) -> bool { - // Requirements must contain indices - if indices.len() > normalized_reqs.len() { - return false; - } - - // Iterate over all orderings - self.oeq_class.iter().any(|ordering| { - if indices.len() > ordering.len() { - return false; - } - - // Build a map of column positions in the ordering - let mut col_positions = HashMap::with_capacity(ordering.len()); - for (pos, req) in ordering.iter().enumerate() { - if let Some(col) = req.expr.as_any().downcast_ref::() { - col_positions.insert( - col.index(), - (pos, col.nullable(&self.schema).unwrap_or(true)), - ); - } - } - - // Check if all constraint indices appear in valid positions - if !indices.iter().all(|&idx| { - col_positions - .get(&idx) - .map(|&(pos, nullable)| { - // For unique constraints, verify column is not nullable if it's first/last - !check_null - || (pos != 0 && pos != ordering.len() - 1) - || !nullable - }) - .unwrap_or(false) - }) { - return false; - } - - // Check if this ordering matches requirements prefix - let ordering_len = ordering.len(); - normalized_reqs.len() >= ordering_len - && normalized_reqs[..ordering_len].iter().zip(ordering).all( - |(req, existing)| { - req.expr.eq(&existing.expr) - && req - .options - .is_none_or(|req_opts| req_opts == existing.options) - }, - ) - }) - } - - /// Determines whether the ordering specified by the given sort requirement - /// is satisfied based on the orderings within, equivalence classes, and - /// constant expressions. - /// - /// # Parameters - /// - /// - `req`: A reference to a `PhysicalSortRequirement` for which the ordering - /// satisfaction check will be done. - /// - /// # Returns - /// - /// Returns `true` if the specified ordering is satisfied, `false` otherwise. - fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool { - let ExprProperties { - sort_properties, .. - } = self.get_expr_properties(Arc::clone(&req.expr)); - match sort_properties { - SortProperties::Ordered(options) => { - let sort_expr = PhysicalSortExpr { - expr: Arc::clone(&req.expr), - options, - }; - sort_expr.satisfy(req, self.schema()) - } - // Singleton expressions satisfies any ordering. - SortProperties::Singleton => true, - SortProperties::Unordered => false, - } - } - - /// Checks whether the `given` sort requirements are equal or more specific - /// than the `reference` sort requirements. - pub fn requirements_compatible( - &self, - given: &LexRequirement, - reference: &LexRequirement, - ) -> bool { - let normalized_given = self.normalize_sort_requirements(given); - let normalized_reference = self.normalize_sort_requirements(reference); - - (normalized_reference.len() <= normalized_given.len()) - && normalized_reference - .into_iter() - .zip(normalized_given) - .all(|(reference, given)| given.compatible(&reference)) - } - - /// Returns the finer ordering among the orderings `lhs` and `rhs`, breaking - /// any ties by choosing `lhs`. - /// - /// The finer ordering is the ordering that satisfies both of the orderings. - /// If the orderings are incomparable, returns `None`. - /// - /// For example, the finer ordering among `[a ASC]` and `[a ASC, b ASC]` is - /// the latter. - pub fn get_finer_ordering( - &self, - lhs: &LexOrdering, - rhs: &LexOrdering, - ) -> Option { - // Convert the given sort expressions to sort requirements: - let lhs = LexRequirement::from(lhs.clone()); - let rhs = LexRequirement::from(rhs.clone()); - let finer = self.get_finer_requirement(&lhs, &rhs); - // Convert the chosen sort requirements back to sort expressions: - finer.map(LexOrdering::from) - } - - /// Returns the finer ordering among the requirements `lhs` and `rhs`, - /// breaking any ties by choosing `lhs`. - /// - /// The finer requirements are the ones that satisfy both of the given - /// requirements. If the requirements are incomparable, returns `None`. - /// - /// For example, the finer requirements among `[a ASC]` and `[a ASC, b ASC]` - /// is the latter. - pub fn get_finer_requirement( - &self, - req1: &LexRequirement, - req2: &LexRequirement, - ) -> Option { - let mut lhs = self.normalize_sort_requirements(req1); - let mut rhs = self.normalize_sort_requirements(req2); - lhs.inner - .iter_mut() - .zip(rhs.inner.iter_mut()) - .all(|(lhs, rhs)| { - lhs.expr.eq(&rhs.expr) - && match (lhs.options, rhs.options) { - (Some(lhs_opt), Some(rhs_opt)) => lhs_opt == rhs_opt, - (Some(options), None) => { - rhs.options = Some(options); - true - } - (None, Some(options)) => { - lhs.options = Some(options); - true - } - (None, None) => true, - } - }) - .then_some(if lhs.len() >= rhs.len() { lhs } else { rhs }) - } - - /// we substitute the ordering according to input expression type, this is a simplified version - /// In this case, we just substitute when the expression satisfy the following condition: - /// I. just have one column and is a CAST expression - /// TODO: Add one-to-ones analysis for monotonic ScalarFunctions. - /// TODO: we could precompute all the scenario that is computable, for example: atan(x + 1000) should also be substituted if - /// x is DESC or ASC - /// After substitution, we may generate more than 1 `LexOrdering`. As an example, - /// `[a ASC, b ASC]` will turn into `[a ASC, b ASC], [CAST(a) ASC, b ASC]` when projection expressions `a, b, CAST(a)` is applied. - pub fn substitute_ordering_component( - &self, - mapping: &ProjectionMapping, - sort_expr: &LexOrdering, - ) -> Result> { - let new_orderings = sort_expr - .iter() - .map(|sort_expr| { - let referring_exprs: Vec<_> = mapping - .iter() - .map(|(source, _target)| source) - .filter(|source| expr_refers(source, &sort_expr.expr)) - .cloned() - .collect(); - let mut res = LexOrdering::new(vec![sort_expr.clone()]); - // TODO: Add one-to-ones analysis for ScalarFunctions. - for r_expr in referring_exprs { - // we check whether this expression is substitutable or not - if let Some(cast_expr) = r_expr.as_any().downcast_ref::() { - // we need to know whether the Cast Expr matches or not - let expr_type = sort_expr.expr.data_type(&self.schema)?; - if cast_expr.expr.eq(&sort_expr.expr) - && cast_expr.is_bigger_cast(expr_type) - { - res.push(PhysicalSortExpr { - expr: Arc::clone(&r_expr), - options: sort_expr.options, - }); - } - } - } - Ok(res) - }) - .collect::>>()?; - // Generate all valid orderings, given substituted expressions. - let res = new_orderings - .into_iter() - .multi_cartesian_product() - .map(LexOrdering::new) - .collect::>(); - Ok(res) - } - - /// In projection, supposed we have a input function 'A DESC B DESC' and the output shares the same expression - /// with A and B, we could surely use the ordering of the original ordering, However, if the A has been changed, - /// for example, A-> Cast(A, Int64) or any other form, it is invalid if we continue using the original ordering - /// Since it would cause bug in dependency constructions, we should substitute the input order in order to get correct - /// dependency map, happen in issue 8838: - pub fn substitute_oeq_class(&mut self, mapping: &ProjectionMapping) -> Result<()> { - let new_order = self - .oeq_class - .iter() - .map(|order| self.substitute_ordering_component(mapping, order)) - .collect::>>()?; - let new_order = new_order.into_iter().flatten().collect(); - self.oeq_class = OrderingEquivalenceClass::new(new_order); - Ok(()) - } - /// Projects argument `expr` according to `projection_mapping`, taking - /// equivalences into account. - /// - /// For example, assume that columns `a` and `c` are always equal, and that - /// `projection_mapping` encodes following mapping: - /// - /// ```text - /// a -> a1 - /// b -> b1 - /// ``` - /// - /// Then, this function projects `a + b` to `Some(a1 + b1)`, `c + b` to - /// `Some(a1 + b1)` and `d` to `None`, meaning that it cannot be projected. - pub fn project_expr( - &self, - expr: &Arc, - projection_mapping: &ProjectionMapping, - ) -> Option> { - self.eq_group.project_expr(projection_mapping, expr) - } - - /// Constructs a dependency map based on existing orderings referred to in - /// the projection. - /// - /// This function analyzes the orderings in the normalized order-equivalence - /// class and builds a dependency map. The dependency map captures relationships - /// between expressions within the orderings, helping to identify dependencies - /// and construct valid projected orderings during projection operations. - /// - /// # Parameters - /// - /// - `mapping`: A reference to the `ProjectionMapping` that defines the - /// relationship between source and target expressions. - /// - /// # Returns - /// - /// A [`DependencyMap`] representing the dependency map, where each - /// [`DependencyNode`] contains dependencies for the key [`PhysicalSortExpr`]. - /// - /// # Example - /// - /// Assume we have two equivalent orderings: `[a ASC, b ASC]` and `[a ASC, c ASC]`, - /// and the projection mapping is `[a -> a_new, b -> b_new, b + c -> b + c]`. - /// Then, the dependency map will be: - /// - /// ```text - /// a ASC: Node {Some(a_new ASC), HashSet{}} - /// b ASC: Node {Some(b_new ASC), HashSet{a ASC}} - /// c ASC: Node {None, HashSet{a ASC}} - /// ``` - fn construct_dependency_map(&self, mapping: &ProjectionMapping) -> DependencyMap { - let mut dependency_map = DependencyMap::new(); - for ordering in self.normalized_oeq_class().iter() { - for (idx, sort_expr) in ordering.iter().enumerate() { - let target_sort_expr = - self.project_expr(&sort_expr.expr, mapping).map(|expr| { - PhysicalSortExpr { - expr, - options: sort_expr.options, - } - }); - let is_projected = target_sort_expr.is_some(); - if is_projected - || mapping - .iter() - .any(|(source, _)| expr_refers(source, &sort_expr.expr)) - { - // Previous ordering is a dependency. Note that there is no, - // dependency for a leading ordering (i.e. the first sort - // expression). - let dependency = idx.checked_sub(1).map(|a| &ordering[a]); - // Add sort expressions that can be projected or referred to - // by any of the projection expressions to the dependency map: - dependency_map.insert( - sort_expr, - target_sort_expr.as_ref(), - dependency, - ); - } - if !is_projected { - // If we can not project, stop constructing the dependency - // map as remaining dependencies will be invalid after projection. - break; - } - } - } - dependency_map - } - - /// Returns a new `ProjectionMapping` where source expressions are normalized. - /// - /// This normalization ensures that source expressions are transformed into a - /// consistent representation. This is beneficial for algorithms that rely on - /// exact equalities, as it allows for more precise and reliable comparisons. - /// - /// # Parameters - /// - /// - `mapping`: A reference to the original `ProjectionMapping` to be normalized. - /// - /// # Returns - /// - /// A new `ProjectionMapping` with normalized source expressions. - fn normalized_mapping(&self, mapping: &ProjectionMapping) -> ProjectionMapping { - // Construct the mapping where source expressions are normalized. In this way - // In the algorithms below we can work on exact equalities - ProjectionMapping { - map: mapping - .iter() - .map(|(source, target)| { - let normalized_source = - self.eq_group.normalize_expr(Arc::clone(source)); - (normalized_source, Arc::clone(target)) - }) - .collect(), - } - } - - /// Computes projected orderings based on a given projection mapping. - /// - /// This function takes a `ProjectionMapping` and computes the possible - /// orderings for the projected expressions. It considers dependencies - /// between expressions and generates valid orderings according to the - /// specified sort properties. - /// - /// # Parameters - /// - /// - `mapping`: A reference to the `ProjectionMapping` that defines the - /// relationship between source and target expressions. - /// - /// # Returns - /// - /// A vector of `LexOrdering` containing all valid orderings after projection. - fn projected_orderings(&self, mapping: &ProjectionMapping) -> Vec { - let mapping = self.normalized_mapping(mapping); - - // Get dependency map for existing orderings: - let dependency_map = self.construct_dependency_map(&mapping); - let orderings = mapping.iter().flat_map(|(source, target)| { - referred_dependencies(&dependency_map, source) - .into_iter() - .filter_map(|relevant_deps| { - if let Ok(SortProperties::Ordered(options)) = - get_expr_properties(source, &relevant_deps, &self.schema) - .map(|prop| prop.sort_properties) - { - Some((options, relevant_deps)) - } else { - // Do not consider unordered cases - None - } - }) - .flat_map(|(options, relevant_deps)| { - let sort_expr = PhysicalSortExpr { - expr: Arc::clone(target), - options, - }; - // Generate dependent orderings (i.e. prefixes for `sort_expr`): - let mut dependency_orderings = - generate_dependency_orderings(&relevant_deps, &dependency_map); - // Append `sort_expr` to the dependent orderings: - for ordering in dependency_orderings.iter_mut() { - ordering.push(sort_expr.clone()); - } - dependency_orderings - }) - }); - - // Add valid projected orderings. For example, if existing ordering is - // `a + b` and projection is `[a -> a_new, b -> b_new]`, we need to - // preserve `a_new + b_new` as ordered. Please note that `a_new` and - // `b_new` themselves need not be ordered. Such dependencies cannot be - // deduced via the pass above. - let projected_orderings = dependency_map.iter().flat_map(|(sort_expr, node)| { - let mut prefixes = construct_prefix_orderings(sort_expr, &dependency_map); - if prefixes.is_empty() { - // If prefix is empty, there is no dependency. Insert - // empty ordering: - prefixes = vec![LexOrdering::default()]; - } - // Append current ordering on top its dependencies: - for ordering in prefixes.iter_mut() { - if let Some(target) = &node.target_sort_expr { - ordering.push(target.clone()) - } - } - prefixes - }); - - // Simplify each ordering by removing redundant sections: - orderings - .chain(projected_orderings) - .map(|lex_ordering| lex_ordering.collapse()) - .collect() - } - - /// Projects constants based on the provided `ProjectionMapping`. - /// - /// This function takes a `ProjectionMapping` and identifies/projects - /// constants based on the existing constants and the mapping. It ensures - /// that constants are appropriately propagated through the projection. - /// - /// # Parameters - /// - /// - `mapping`: A reference to a `ProjectionMapping` representing the - /// mapping of source expressions to target expressions in the projection. - /// - /// # Returns - /// - /// Returns a `Vec>` containing the projected constants. - fn projected_constants(&self, mapping: &ProjectionMapping) -> Vec { - // First, project existing constants. For example, assume that `a + b` - // is known to be constant. If the projection were `a as a_new`, `b as b_new`, - // then we would project constant `a + b` as `a_new + b_new`. - let mut projected_constants = self - .constants - .iter() - .flat_map(|const_expr| { - const_expr - .map(|expr| self.eq_group.project_expr(mapping, expr)) - .map(|projected_expr| { - projected_expr - .with_across_partitions(const_expr.across_partitions()) - }) - }) - .collect::>(); - - // Add projection expressions that are known to be constant: - for (source, target) in mapping.iter() { - if self.is_expr_constant(source) - && !const_exprs_contains(&projected_constants, target) - { - if self.is_expr_constant_across_partitions(source) { - projected_constants.push( - ConstExpr::from(target) - .with_across_partitions(self.get_expr_constant_value(source)), - ) - } else { - projected_constants.push( - ConstExpr::from(target) - .with_across_partitions(AcrossPartitions::Heterogeneous), - ) - } - } - } - projected_constants - } - - /// Projects constraints according to the given projection mapping. - /// - /// This function takes a projection mapping and extracts the column indices of the target columns. - /// It then projects the constraints to only include relationships between - /// columns that exist in the projected output. - /// - /// # Arguments - /// - /// * `mapping` - A reference to `ProjectionMapping` that defines how expressions are mapped - /// in the projection operation - /// - /// # Returns - /// - /// Returns a new `Constraints` object containing only the constraints - /// that are valid for the projected columns. - fn projected_constraints(&self, mapping: &ProjectionMapping) -> Option { - let indices = mapping - .iter() - .filter_map(|(_, target)| target.as_any().downcast_ref::()) - .map(|col| col.index()) - .collect::>(); - debug_assert_eq!(mapping.map.len(), indices.len()); - self.constraints.project(&indices) - } - - /// Projects the equivalences within according to `mapping` - /// and `output_schema`. - pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self { - let eq_group = self.eq_group.project(mapping); - let oeq_class = OrderingEquivalenceClass::new(self.projected_orderings(mapping)); - let constants = self.projected_constants(mapping); - let constraints = self - .projected_constraints(mapping) - .unwrap_or_else(Constraints::empty); - Self { - schema: output_schema, - eq_group, - oeq_class, - constants, - constraints, - } - } - - /// Returns the longest (potentially partial) permutation satisfying the - /// existing ordering. For example, if we have the equivalent orderings - /// `[a ASC, b ASC]` and `[c DESC]`, with `exprs` containing `[c, b, a, d]`, - /// then this function returns `([a ASC, b ASC, c DESC], [2, 1, 0])`. - /// This means that the specification `[a ASC, b ASC, c DESC]` is satisfied - /// by the existing ordering, and `[a, b, c]` resides at indices: `2, 1, 0` - /// inside the argument `exprs` (respectively). For the mathematical - /// definition of "partial permutation", see: - /// - /// - pub fn find_longest_permutation( - &self, - exprs: &[Arc], - ) -> (LexOrdering, Vec) { - let mut eq_properties = self.clone(); - let mut result = vec![]; - // The algorithm is as follows: - // - Iterate over all the expressions and insert ordered expressions - // into the result. - // - Treat inserted expressions as constants (i.e. add them as constants - // to the state). - // - Continue the above procedure until no expression is inserted; i.e. - // the algorithm reaches a fixed point. - // This algorithm should reach a fixed point in at most `exprs.len()` - // iterations. - let mut search_indices = (0..exprs.len()).collect::>(); - for _idx in 0..exprs.len() { - // Get ordered expressions with their indices. - let ordered_exprs = search_indices - .iter() - .flat_map(|&idx| { - let ExprProperties { - sort_properties, .. - } = eq_properties.get_expr_properties(Arc::clone(&exprs[idx])); - match sort_properties { - SortProperties::Ordered(options) => Some(( - PhysicalSortExpr { - expr: Arc::clone(&exprs[idx]), - options, - }, - idx, - )), - SortProperties::Singleton => { - // Assign default ordering to constant expressions - let options = SortOptions::default(); - Some(( - PhysicalSortExpr { - expr: Arc::clone(&exprs[idx]), - options, - }, - idx, - )) - } - SortProperties::Unordered => None, - } - }) - .collect::>(); - // We reached a fixed point, exit. - if ordered_exprs.is_empty() { - break; - } - // Remove indices that have an ordering from `search_indices`, and - // treat ordered expressions as constants in subsequent iterations. - // We can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs { - eq_properties = - eq_properties.with_constants(std::iter::once(ConstExpr::from(expr))); - search_indices.shift_remove(idx); - } - // Add new ordered section to the state. - result.extend(ordered_exprs); - } - let (left, right) = result.into_iter().unzip(); - (LexOrdering::new(left), right) - } - - /// This function determines whether the provided expression is constant - /// based on the known constants. - /// - /// # Parameters - /// - /// - `expr`: A reference to a `Arc` representing the - /// expression to be checked. - /// - /// # Returns - /// - /// Returns `true` if the expression is constant according to equivalence - /// group, `false` otherwise. - pub fn is_expr_constant(&self, expr: &Arc) -> bool { - // As an example, assume that we know columns `a` and `b` are constant. - // Then, `a`, `b` and `a + b` will all return `true` whereas `c` will - // return `false`. - let const_exprs = self - .constants - .iter() - .map(|const_expr| Arc::clone(const_expr.expr())); - let normalized_constants = self.eq_group.normalize_exprs(const_exprs); - let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); - is_constant_recurse(&normalized_constants, &normalized_expr) - } - - /// This function determines whether the provided expression is constant - /// across partitions based on the known constants. - /// - /// # Parameters - /// - /// - `expr`: A reference to a `Arc` representing the - /// expression to be checked. - /// - /// # Returns - /// - /// Returns `true` if the expression is constant across all partitions according - /// to equivalence group, `false` otherwise - #[deprecated( - since = "45.0.0", - note = "Use [`is_expr_constant_across_partitions`] instead" - )] - pub fn is_expr_constant_accross_partitions( - &self, - expr: &Arc, - ) -> bool { - self.is_expr_constant_across_partitions(expr) - } - - /// This function determines whether the provided expression is constant - /// across partitions based on the known constants. - /// - /// # Parameters - /// - /// - `expr`: A reference to a `Arc` representing the - /// expression to be checked. - /// - /// # Returns - /// - /// Returns `true` if the expression is constant across all partitions according - /// to equivalence group, `false` otherwise. - pub fn is_expr_constant_across_partitions( - &self, - expr: &Arc, - ) -> bool { - // As an example, assume that we know columns `a` and `b` are constant. - // Then, `a`, `b` and `a + b` will all return `true` whereas `c` will - // return `false`. - let const_exprs = self - .constants - .iter() - .filter_map(|const_expr| { - if matches!( - const_expr.across_partitions(), - AcrossPartitions::Uniform { .. } - ) { - Some(Arc::clone(const_expr.expr())) - } else { - None - } - }) - .collect::>(); - let normalized_constants = self.eq_group.normalize_exprs(const_exprs); - let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); - is_constant_recurse(&normalized_constants, &normalized_expr) - } - - /// Retrieves the constant value of a given physical expression, if it exists. - /// - /// Normalizes the input expression and checks if it matches any known constants - /// in the current context. Returns whether the expression has a uniform value, - /// varies across partitions, or is not constant. - /// - /// # Parameters - /// - `expr`: A reference to the physical expression to evaluate. - /// - /// # Returns - /// - `AcrossPartitions::Uniform(value)`: If the expression has the same value across partitions. - /// - `AcrossPartitions::Heterogeneous`: If the expression varies across partitions. - /// - `None`: If the expression is not recognized as constant. - pub fn get_expr_constant_value( - &self, - expr: &Arc, - ) -> AcrossPartitions { - let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); - - if let Some(lit) = normalized_expr.as_any().downcast_ref::() { - return AcrossPartitions::Uniform(Some(lit.value().clone())); - } - - for const_expr in self.constants.iter() { - if normalized_expr.eq(const_expr.expr()) { - return const_expr.across_partitions(); - } - } - - AcrossPartitions::Heterogeneous - } - - /// Retrieves the properties for a given physical expression. - /// - /// This function constructs an [`ExprProperties`] object for the given - /// expression, which encapsulates information about the expression's - /// properties, including its [`SortProperties`] and [`Interval`]. - /// - /// # Parameters - /// - /// - `expr`: An `Arc` representing the physical expression - /// for which ordering information is sought. - /// - /// # Returns - /// - /// Returns an [`ExprProperties`] object containing the ordering and range - /// information for the given expression. - pub fn get_expr_properties(&self, expr: Arc) -> ExprProperties { - ExprPropertiesNode::new_unknown(expr) - .transform_up(|expr| update_properties(expr, self)) - .data() - .map(|node| node.data) - .unwrap_or(ExprProperties::new_unknown()) - } - - /// Transforms this `EquivalenceProperties` into a new `EquivalenceProperties` - /// by mapping columns in the original schema to columns in the new schema - /// by index. - pub fn with_new_schema(self, schema: SchemaRef) -> Result { - // The new schema and the original schema is aligned when they have the - // same number of columns, and fields at the same index have the same - // type in both schemas. - let schemas_aligned = (self.schema.fields.len() == schema.fields.len()) - && self - .schema - .fields - .iter() - .zip(schema.fields.iter()) - .all(|(lhs, rhs)| lhs.data_type().eq(rhs.data_type())); - if !schemas_aligned { - // Rewriting equivalence properties in terms of new schema is not - // safe when schemas are not aligned: - return plan_err!( - "Cannot rewrite old_schema:{:?} with new schema: {:?}", - self.schema, - schema - ); - } - // Rewrite constants according to new schema: - let new_constants = self - .constants - .into_iter() - .map(|const_expr| { - let across_partitions = const_expr.across_partitions(); - let new_const_expr = with_new_schema(const_expr.owned_expr(), &schema)?; - Ok(ConstExpr::new(new_const_expr) - .with_across_partitions(across_partitions)) - }) - .collect::>>()?; - - // Rewrite orderings according to new schema: - let mut new_orderings = vec![]; - for ordering in self.oeq_class { - let new_ordering = ordering - .into_iter() - .map(|mut sort_expr| { - sort_expr.expr = with_new_schema(sort_expr.expr, &schema)?; - Ok(sort_expr) - }) - .collect::>()?; - new_orderings.push(new_ordering); - } - - // Rewrite equivalence classes according to the new schema: - let mut eq_classes = vec![]; - for eq_class in self.eq_group { - let new_eq_exprs = eq_class - .into_vec() - .into_iter() - .map(|expr| with_new_schema(expr, &schema)) - .collect::>()?; - eq_classes.push(EquivalenceClass::new(new_eq_exprs)); - } - - // Construct the resulting equivalence properties: - let mut result = EquivalenceProperties::new(schema); - result.constants = new_constants; - result.add_new_orderings(new_orderings); - result.add_equivalence_group(EquivalenceGroup::new(eq_classes)); - - Ok(result) - } -} - -/// More readable display version of the `EquivalenceProperties`. -/// -/// Format: -/// ```text -/// order: [[a ASC, b ASC], [a ASC, c ASC]], eq: [[a = b], [a = c]], const: [a = 1] -/// ``` -impl Display for EquivalenceProperties { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - if self.eq_group.is_empty() - && self.oeq_class.is_empty() - && self.constants.is_empty() - { - return write!(f, "No properties"); - } - if !self.oeq_class.is_empty() { - write!(f, "order: {}", self.oeq_class)?; - } - if !self.eq_group.is_empty() { - write!(f, ", eq: {}", self.eq_group)?; - } - if !self.constants.is_empty() { - write!(f, ", const: [{}]", ConstExpr::format_list(&self.constants))?; - } - Ok(()) - } -} - -/// Calculates the properties of a given [`ExprPropertiesNode`]. -/// -/// Order information can be retrieved as: -/// - If it is a leaf node, we directly find the order of the node by looking -/// at the given sort expression and equivalence properties if it is a `Column` -/// leaf, or we mark it as unordered. In the case of a `Literal` leaf, we mark -/// it as singleton so that it can cooperate with all ordered columns. -/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` -/// and operator has its own rules on how to propagate the children orderings. -/// However, before we engage in recursion, we check whether this intermediate -/// node directly matches with the sort expression. If there is a match, the -/// sort expression emerges at that node immediately, discarding the recursive -/// result coming from its children. -/// -/// Range information is calculated as: -/// - If it is a `Literal` node, we set the range as a point value. If it is a -/// `Column` node, we set the datatype of the range, but cannot give an interval -/// for the range, yet. -/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` -/// and operator has its own rules on how to propagate the children range. -fn update_properties( - mut node: ExprPropertiesNode, - eq_properties: &EquivalenceProperties, -) -> Result> { - // First, try to gather the information from the children: - if !node.expr.children().is_empty() { - // We have an intermediate (non-leaf) node, account for its children: - let children_props = node.children.iter().map(|c| c.data.clone()).collect_vec(); - node.data = node.expr.get_properties(&children_props)?; - } else if node.expr.as_any().is::() { - // We have a Literal, which is one of the two possible leaf node types: - node.data = node.expr.get_properties(&[])?; - } else if node.expr.as_any().is::() { - // We have a Column, which is the other possible leaf node type: - node.data.range = - Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)? - } - // Now, check what we know about orderings: - let normalized_expr = eq_properties - .eq_group - .normalize_expr(Arc::clone(&node.expr)); - let oeq_class = eq_properties.normalized_oeq_class(); - if eq_properties.is_expr_constant(&normalized_expr) - || oeq_class.is_expr_partial_const(&normalized_expr) - { - node.data.sort_properties = SortProperties::Singleton; - } else if let Some(options) = oeq_class.get_options(&normalized_expr) { - node.data.sort_properties = SortProperties::Ordered(options); - } - Ok(Transformed::yes(node)) -} - -/// This function determines whether the provided expression is constant -/// based on the known constants. -/// -/// # Parameters -/// -/// - `constants`: A `&[Arc]` containing expressions known to -/// be a constant. -/// - `expr`: A reference to a `Arc` representing the expression -/// to check. -/// -/// # Returns -/// -/// Returns `true` if the expression is constant according to equivalence -/// group, `false` otherwise. -fn is_constant_recurse( - constants: &[Arc], - expr: &Arc, -) -> bool { - if physical_exprs_contains(constants, expr) || expr.as_any().is::() { - return true; - } - let children = expr.children(); - !children.is_empty() && children.iter().all(|c| is_constant_recurse(constants, c)) -} - -/// This function examines whether a referring expression directly refers to a -/// given referred expression or if any of its children in the expression tree -/// refer to the specified expression. -/// -/// # Parameters -/// -/// - `referring_expr`: A reference to the referring expression (`Arc`). -/// - `referred_expr`: A reference to the referred expression (`Arc`) -/// -/// # Returns -/// -/// A boolean value indicating whether `referring_expr` refers (needs it to evaluate its result) -/// `referred_expr` or not. -fn expr_refers( - referring_expr: &Arc, - referred_expr: &Arc, -) -> bool { - referring_expr.eq(referred_expr) - || referring_expr - .children() - .iter() - .any(|child| expr_refers(child, referred_expr)) -} - -/// This function analyzes the dependency map to collect referred dependencies for -/// a given source expression. -/// -/// # Parameters -/// -/// - `dependency_map`: A reference to the `DependencyMap` where each -/// `PhysicalSortExpr` is associated with a `DependencyNode`. -/// - `source`: A reference to the source expression (`Arc`) -/// for which relevant dependencies need to be identified. -/// -/// # Returns -/// -/// A `Vec` containing the dependencies for the given source -/// expression. These dependencies are expressions that are referred to by -/// the source expression based on the provided dependency map. -fn referred_dependencies( - dependency_map: &DependencyMap, - source: &Arc, -) -> Vec { - // Associate `PhysicalExpr`s with `PhysicalSortExpr`s that contain them: - let mut expr_to_sort_exprs = IndexMap::::new(); - for sort_expr in dependency_map - .sort_exprs() - .filter(|sort_expr| expr_refers(source, &sort_expr.expr)) - { - let key = ExprWrapper(Arc::clone(&sort_expr.expr)); - expr_to_sort_exprs - .entry(key) - .or_default() - .insert(sort_expr.clone()); - } - - // Generate all valid dependencies for the source. For example, if the source - // is `a + b` and the map is `[a -> (a ASC, a DESC), b -> (b ASC)]`, we get - // `vec![HashSet(a ASC, b ASC), HashSet(a DESC, b ASC)]`. - let dependencies = expr_to_sort_exprs - .into_values() - .map(Dependencies::into_inner) - .collect::>(); - dependencies - .iter() - .multi_cartesian_product() - .map(|referred_deps| { - Dependencies::new_from_iter(referred_deps.into_iter().cloned()) - }) - .collect() -} - -/// This function retrieves the dependencies of the given relevant sort expression -/// from the given dependency map. It then constructs prefix orderings by recursively -/// analyzing the dependencies and include them in the orderings. -/// -/// # Parameters -/// -/// - `relevant_sort_expr`: A reference to the relevant sort expression -/// (`PhysicalSortExpr`) for which prefix orderings are to be constructed. -/// - `dependency_map`: A reference to the `DependencyMap` containing dependencies. -/// -/// # Returns -/// -/// A vector of prefix orderings (`Vec`) based on the given relevant -/// sort expression and its dependencies. -fn construct_prefix_orderings( - relevant_sort_expr: &PhysicalSortExpr, - dependency_map: &DependencyMap, -) -> Vec { - let mut dep_enumerator = DependencyEnumerator::new(); - dependency_map - .get(relevant_sort_expr) - .expect("no relevant sort expr found") - .dependencies - .iter() - .flat_map(|dep| dep_enumerator.construct_orderings(dep, dependency_map)) - .collect() -} - -/// Generates all possible orderings where dependencies are satisfied for the -/// current projection expression. -/// -/// # Example -/// If `dependencies` is `a + b ASC` and the dependency map holds dependencies -/// * `a ASC` --> `[c ASC]` -/// * `b ASC` --> `[d DESC]`, -/// -/// This function generates these two sort orders -/// * `[c ASC, d DESC, a + b ASC]` -/// * `[d DESC, c ASC, a + b ASC]` -/// -/// # Parameters -/// -/// * `dependencies` - Set of relevant expressions. -/// * `dependency_map` - Map of dependencies for expressions that may appear in `dependencies` -/// -/// # Returns -/// -/// A vector of lexical orderings (`Vec`) representing all valid orderings -/// based on the given dependencies. -fn generate_dependency_orderings( - dependencies: &Dependencies, - dependency_map: &DependencyMap, -) -> Vec { - // Construct all the valid prefix orderings for each expression appearing - // in the projection: - let relevant_prefixes = dependencies - .iter() - .flat_map(|dep| { - let prefixes = construct_prefix_orderings(dep, dependency_map); - (!prefixes.is_empty()).then_some(prefixes) - }) - .collect::>(); - - // No dependency, dependent is a leading ordering. - if relevant_prefixes.is_empty() { - // Return an empty ordering: - return vec![LexOrdering::default()]; - } - - relevant_prefixes - .into_iter() - .multi_cartesian_product() - .flat_map(|prefix_orderings| { - prefix_orderings - .iter() - .permutations(prefix_orderings.len()) - .map(|prefixes| { - prefixes - .into_iter() - .flat_map(|ordering| ordering.clone()) - .collect() - }) - .collect::>() - }) - .collect() -} - -/// This function examines the given expression and its properties to determine -/// the ordering properties of the expression. The range knowledge is not utilized -/// yet in the scope of this function. -/// -/// # Parameters -/// -/// - `expr`: A reference to the source expression (`Arc`) for -/// which ordering properties need to be determined. -/// - `dependencies`: A reference to `Dependencies`, containing sort expressions -/// referred to by `expr`. -/// - `schema``: A reference to the schema which the `expr` columns refer. -/// -/// # Returns -/// -/// A `SortProperties` indicating the ordering information of the given expression. -fn get_expr_properties( - expr: &Arc, - dependencies: &Dependencies, - schema: &SchemaRef, -) -> Result { - if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) { - // If exact match is found, return its ordering. - Ok(ExprProperties { - sort_properties: SortProperties::Ordered(column_order.options), - range: Interval::make_unbounded(&expr.data_type(schema)?)?, - preserves_lex_ordering: false, - }) - } else if expr.as_any().downcast_ref::().is_some() { - Ok(ExprProperties { - sort_properties: SortProperties::Unordered, - range: Interval::make_unbounded(&expr.data_type(schema)?)?, - preserves_lex_ordering: false, - }) - } else if let Some(literal) = expr.as_any().downcast_ref::() { - Ok(ExprProperties { - sort_properties: SortProperties::Singleton, - range: Interval::try_new(literal.value().clone(), literal.value().clone())?, - preserves_lex_ordering: true, - }) - } else { - // Find orderings of its children - let child_states = expr - .children() - .iter() - .map(|child| get_expr_properties(child, dependencies, schema)) - .collect::>>()?; - // Calculate expression ordering using ordering of its children. - expr.get_properties(&child_states) - } -} - -/// Represents a node in the dependency map used to construct projected orderings. -/// -/// A `DependencyNode` contains information about a particular sort expression, -/// including its target sort expression and a set of dependencies on other sort -/// expressions. -/// -/// # Fields -/// -/// - `target_sort_expr`: An optional `PhysicalSortExpr` representing the target -/// sort expression associated with the node. It is `None` if the sort expression -/// cannot be projected. -/// - `dependencies`: A [`Dependencies`] containing dependencies on other sort -/// expressions that are referred to by the target sort expression. -#[derive(Debug, Clone, PartialEq, Eq)] -struct DependencyNode { - target_sort_expr: Option, - dependencies: Dependencies, -} - -impl DependencyNode { - /// Insert dependency to the state (if exists). - fn insert_dependency(&mut self, dependency: Option<&PhysicalSortExpr>) { - if let Some(dep) = dependency { - self.dependencies.insert(dep.clone()); - } - } -} - -impl Display for DependencyNode { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - if let Some(target) = &self.target_sort_expr { - write!(f, "(target: {}, ", target)?; - } else { - write!(f, "(")?; - } - write!(f, "dependencies: [{}])", self.dependencies) - } -} - -/// Maps an expression --> DependencyNode -/// -/// # Debugging / deplaying `DependencyMap` -/// -/// This structure implements `Display` to assist debugging. For example: -/// -/// ```text -/// DependencyMap: { -/// a@0 ASC --> (target: a@0 ASC, dependencies: [[]]) -/// b@1 ASC --> (target: b@1 ASC, dependencies: [[a@0 ASC, c@2 ASC]]) -/// c@2 ASC --> (target: c@2 ASC, dependencies: [[b@1 ASC, a@0 ASC]]) -/// d@3 ASC --> (target: d@3 ASC, dependencies: [[c@2 ASC, b@1 ASC]]) -/// } -/// ``` -/// -/// # Note on IndexMap Rationale -/// -/// Using `IndexMap` (which preserves insert order) to ensure consistent results -/// across different executions for the same query. We could have used -/// `HashSet`, `HashMap` in place of them without any loss of functionality. -/// -/// As an example, if existing orderings are -/// 1. `[a ASC, b ASC]` -/// 2. `[c ASC]` for -/// -/// Then both the following output orderings are valid -/// 1. `[a ASC, b ASC, c ASC]` -/// 2. `[c ASC, a ASC, b ASC]` -/// -/// (this are both valid as they are concatenated versions of the alternative -/// orderings). When using `HashSet`, `HashMap` it is not guaranteed to generate -/// consistent result, among the possible 2 results in the example above. -#[derive(Debug)] -struct DependencyMap { - inner: IndexMap, -} - -impl DependencyMap { - fn new() -> Self { - Self { - inner: IndexMap::new(), - } - } - - /// Insert a new dependency `sort_expr` --> `dependency` into the map. - /// - /// If `target_sort_expr` is none, a new entry is created with empty dependencies. - fn insert( - &mut self, - sort_expr: &PhysicalSortExpr, - target_sort_expr: Option<&PhysicalSortExpr>, - dependency: Option<&PhysicalSortExpr>, - ) { - self.inner - .entry(sort_expr.clone()) - .or_insert_with(|| DependencyNode { - target_sort_expr: target_sort_expr.cloned(), - dependencies: Dependencies::new(), - }) - .insert_dependency(dependency) - } - - /// Iterator over (sort_expr, DependencyNode) pairs - fn iter(&self) -> impl Iterator { - self.inner.iter() - } - - /// iterator over all sort exprs - fn sort_exprs(&self) -> impl Iterator { - self.inner.keys() - } - - /// Return the dependency node for the given sort expression, if any - fn get(&self, sort_expr: &PhysicalSortExpr) -> Option<&DependencyNode> { - self.inner.get(sort_expr) - } -} - -impl Display for DependencyMap { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - writeln!(f, "DependencyMap: {{")?; - for (sort_expr, node) in self.inner.iter() { - writeln!(f, " {sort_expr} --> {node}")?; - } - writeln!(f, "}}") - } -} - -/// A list of sort expressions that can be calculated from a known set of -/// dependencies. -#[derive(Debug, Default, Clone, PartialEq, Eq)] -struct Dependencies { - inner: IndexSet, -} - -impl Display for Dependencies { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "[")?; - let mut iter = self.inner.iter(); - if let Some(dep) = iter.next() { - write!(f, "{}", dep)?; - } - for dep in iter { - write!(f, ", {}", dep)?; - } - write!(f, "]") - } -} - -impl Dependencies { - /// Create a new empty `Dependencies` instance. - fn new() -> Self { - Self { - inner: IndexSet::new(), - } - } - - /// Create a new `Dependencies` from an iterator of `PhysicalSortExpr`. - fn new_from_iter(iter: impl IntoIterator) -> Self { - Self { - inner: iter.into_iter().collect(), - } - } - - /// Insert a new dependency into the set. - fn insert(&mut self, sort_expr: PhysicalSortExpr) { - self.inner.insert(sort_expr); - } - - /// Iterator over dependencies in the set - fn iter(&self) -> impl Iterator + Clone { - self.inner.iter() - } - - /// Return the inner set of dependencies - fn into_inner(self) -> IndexSet { - self.inner - } - - /// Returns true if there are no dependencies - fn is_empty(&self) -> bool { - self.inner.is_empty() - } -} - -/// Contains a mapping of all dependencies we have processed for each sort expr -struct DependencyEnumerator<'a> { - /// Maps `expr` --> `[exprs]` that have previously been processed - seen: IndexMap<&'a PhysicalSortExpr, IndexSet<&'a PhysicalSortExpr>>, -} - -impl<'a> DependencyEnumerator<'a> { - fn new() -> Self { - Self { - seen: IndexMap::new(), - } - } - - /// Insert a new dependency, - /// - /// returns false if the dependency was already in the map - /// returns true if the dependency was newly inserted - fn insert( - &mut self, - target: &'a PhysicalSortExpr, - dep: &'a PhysicalSortExpr, - ) -> bool { - self.seen.entry(target).or_default().insert(dep) - } - - /// This function recursively analyzes the dependencies of the given sort - /// expression within the given dependency map to construct lexicographical - /// orderings that include the sort expression and its dependencies. - /// - /// # Parameters - /// - /// - `referred_sort_expr`: A reference to the sort expression (`PhysicalSortExpr`) - /// for which lexicographical orderings satisfying its dependencies are to be - /// constructed. - /// - `dependency_map`: A reference to the `DependencyMap` that contains - /// dependencies for different `PhysicalSortExpr`s. - /// - /// # Returns - /// - /// A vector of lexicographical orderings (`Vec`) based on the given - /// sort expression and its dependencies. - fn construct_orderings( - &mut self, - referred_sort_expr: &'a PhysicalSortExpr, - dependency_map: &'a DependencyMap, - ) -> Vec { - let node = dependency_map - .get(referred_sort_expr) - .expect("`referred_sort_expr` should be inside `dependency_map`"); - // Since we work on intermediate nodes, we are sure `val.target_sort_expr` - // exists. - let target_sort_expr = node.target_sort_expr.as_ref().unwrap(); - // An empty dependency means the referred_sort_expr represents a global ordering. - // Return its projected version, which is the target_expression. - if node.dependencies.is_empty() { - return vec![LexOrdering::new(vec![target_sort_expr.clone()])]; - }; - - node.dependencies - .iter() - .flat_map(|dep| { - let mut orderings = if self.insert(target_sort_expr, dep) { - self.construct_orderings(dep, dependency_map) - } else { - vec![] - }; - - for ordering in orderings.iter_mut() { - ordering.push(target_sort_expr.clone()) - } - orderings - }) - .collect() - } -} - -/// Calculate ordering equivalence properties for the given join operation. -pub fn join_equivalence_properties( - left: EquivalenceProperties, - right: EquivalenceProperties, - join_type: &JoinType, - join_schema: SchemaRef, - maintains_input_order: &[bool], - probe_side: Option, - on: &[(PhysicalExprRef, PhysicalExprRef)], -) -> EquivalenceProperties { - let left_size = left.schema.fields.len(); - let mut result = EquivalenceProperties::new(join_schema); - result.add_equivalence_group(left.eq_group().join( - right.eq_group(), - join_type, - left_size, - on, - )); - - let EquivalenceProperties { - constants: left_constants, - oeq_class: left_oeq_class, - .. - } = left; - let EquivalenceProperties { - constants: right_constants, - oeq_class: mut right_oeq_class, - .. - } = right; - match maintains_input_order { - [true, false] => { - // In this special case, right side ordering can be prefixed with - // the left side ordering. - if let (Some(JoinSide::Left), JoinType::Inner) = (probe_side, join_type) { - updated_right_ordering_equivalence_class( - &mut right_oeq_class, - join_type, - left_size, - ); - - // Right side ordering equivalence properties should be prepended - // with those of the left side while constructing output ordering - // equivalence properties since stream side is the left side. - // - // For example, if the right side ordering equivalences contain - // `b ASC`, and the left side ordering equivalences contain `a ASC`, - // then we should add `a ASC, b ASC` to the ordering equivalences - // of the join output. - let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class); - result.add_ordering_equivalence_class(out_oeq_class); - } else { - result.add_ordering_equivalence_class(left_oeq_class); - } - } - [false, true] => { - updated_right_ordering_equivalence_class( - &mut right_oeq_class, - join_type, - left_size, - ); - // In this special case, left side ordering can be prefixed with - // the right side ordering. - if let (Some(JoinSide::Right), JoinType::Inner) = (probe_side, join_type) { - // Left side ordering equivalence properties should be prepended - // with those of the right side while constructing output ordering - // equivalence properties since stream side is the right side. - // - // For example, if the left side ordering equivalences contain - // `a ASC`, and the right side ordering equivalences contain `b ASC`, - // then we should add `b ASC, a ASC` to the ordering equivalences - // of the join output. - let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class); - result.add_ordering_equivalence_class(out_oeq_class); - } else { - result.add_ordering_equivalence_class(right_oeq_class); - } - } - [false, false] => {} - [true, true] => unreachable!("Cannot maintain ordering of both sides"), - _ => unreachable!("Join operators can not have more than two children"), - } - match join_type { - JoinType::LeftAnti | JoinType::LeftSemi => { - result = result.with_constants(left_constants); - } - JoinType::RightAnti | JoinType::RightSemi => { - result = result.with_constants(right_constants); - } - _ => {} - } - result -} - -/// In the context of a join, update the right side `OrderingEquivalenceClass` -/// so that they point to valid indices in the join output schema. -/// -/// To do so, we increment column indices by the size of the left table when -/// join schema consists of a combination of the left and right schemas. This -/// is the case for `Inner`, `Left`, `Full` and `Right` joins. For other cases, -/// indices do not change. -fn updated_right_ordering_equivalence_class( - right_oeq_class: &mut OrderingEquivalenceClass, - join_type: &JoinType, - left_size: usize, -) { - if matches!( - join_type, - JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right - ) { - right_oeq_class.add_offset(left_size); - } -} - -/// Wrapper struct for `Arc` to use them as keys in a hash map. -#[derive(Debug, Clone)] -struct ExprWrapper(Arc); - -impl PartialEq for ExprWrapper { - fn eq(&self, other: &Self) -> bool { - self.0.eq(&other.0) - } -} - -impl Eq for ExprWrapper {} - -impl Hash for ExprWrapper { - fn hash(&self, state: &mut H) { - self.0.hash(state); - } -} - -/// Calculates the union (in the sense of `UnionExec`) `EquivalenceProperties` -/// of `lhs` and `rhs` according to the schema of `lhs`. -/// -/// Rules: The UnionExec does not interleave its inputs: instead it passes each -/// input partition from the children as its own output. -/// -/// Since the output equivalence properties are properties that are true for -/// *all* output partitions, that is the same as being true for all *input* -/// partitions -fn calculate_union_binary( - lhs: EquivalenceProperties, - mut rhs: EquivalenceProperties, -) -> Result { - // Harmonize the schema of the rhs with the schema of the lhs (which is the accumulator schema): - if !rhs.schema.eq(&lhs.schema) { - rhs = rhs.with_new_schema(Arc::clone(&lhs.schema))?; - } - - // First, calculate valid constants for the union. An expression is constant - // at the output of the union if it is constant in both sides with matching values. - let constants = lhs - .constants() - .iter() - .filter_map(|lhs_const| { - // Find matching constant expression in RHS - rhs.constants() - .iter() - .find(|rhs_const| rhs_const.expr().eq(lhs_const.expr())) - .map(|rhs_const| { - let mut const_expr = ConstExpr::new(Arc::clone(lhs_const.expr())); - - // If both sides have matching constant values, preserve the value and set across_partitions=true - if let ( - AcrossPartitions::Uniform(Some(lhs_val)), - AcrossPartitions::Uniform(Some(rhs_val)), - ) = (lhs_const.across_partitions(), rhs_const.across_partitions()) - { - if lhs_val == rhs_val { - const_expr = const_expr.with_across_partitions( - AcrossPartitions::Uniform(Some(lhs_val)), - ) - } - } - const_expr - }) - }) - .collect::>(); - - // Next, calculate valid orderings for the union by searching for prefixes - // in both sides. - let mut orderings = UnionEquivalentOrderingBuilder::new(); - orderings.add_satisfied_orderings(lhs.normalized_oeq_class(), lhs.constants(), &rhs); - orderings.add_satisfied_orderings(rhs.normalized_oeq_class(), rhs.constants(), &lhs); - let orderings = orderings.build(); - - let mut eq_properties = - EquivalenceProperties::new(lhs.schema).with_constants(constants); - - eq_properties.add_new_orderings(orderings); - Ok(eq_properties) -} - -/// Calculates the union (in the sense of `UnionExec`) `EquivalenceProperties` -/// of the given `EquivalenceProperties` in `eqps` according to the given -/// output `schema` (which need not be the same with those of `lhs` and `rhs` -/// as details such as nullability may be different). -pub fn calculate_union( - eqps: Vec, - schema: SchemaRef, -) -> Result { - // TODO: In some cases, we should be able to preserve some equivalence - // classes. Add support for such cases. - let mut iter = eqps.into_iter(); - let Some(mut acc) = iter.next() else { - return internal_err!( - "Cannot calculate EquivalenceProperties for a union with no inputs" - ); - }; - - // Harmonize the schema of the init with the schema of the union: - if !acc.schema.eq(&schema) { - acc = acc.with_new_schema(schema)?; - } - // Fold in the rest of the EquivalenceProperties: - for props in iter { - acc = calculate_union_binary(acc, props)?; - } - Ok(acc) -} - -#[derive(Debug)] -enum AddedOrdering { - /// The ordering was added to the in progress result - Yes, - /// The ordering was not added - No(LexOrdering), -} - -/// Builds valid output orderings of a `UnionExec` -#[derive(Debug)] -struct UnionEquivalentOrderingBuilder { - orderings: Vec, -} - -impl UnionEquivalentOrderingBuilder { - fn new() -> Self { - Self { orderings: vec![] } - } - - /// Add all orderings from `orderings` that satisfy `properties`, - /// potentially augmented with`constants`. - /// - /// Note: any column that is known to be constant can be inserted into the - /// ordering without changing its meaning - /// - /// For example: - /// * `orderings` contains `[a ASC, c ASC]` and `constants` contains `b` - /// * `properties` has required ordering `[a ASC, b ASC]` - /// - /// Then this will add `[a ASC, b ASC]` to the `orderings` list (as `a` was - /// in the sort order and `b` was a constant). - fn add_satisfied_orderings( - &mut self, - orderings: impl IntoIterator, - constants: &[ConstExpr], - properties: &EquivalenceProperties, - ) { - for mut ordering in orderings.into_iter() { - // Progressively shorten the ordering to search for a satisfied prefix: - loop { - match self.try_add_ordering(ordering, constants, properties) { - AddedOrdering::Yes => break, - AddedOrdering::No(o) => { - ordering = o; - ordering.pop(); - } - } - } - } - } - - /// Adds `ordering`, potentially augmented with constants, if it satisfies - /// the target `properties` properties. - /// - /// Returns - /// - /// * [`AddedOrdering::Yes`] if the ordering was added (either directly or - /// augmented), or was empty. - /// - /// * [`AddedOrdering::No`] if the ordering was not added - fn try_add_ordering( - &mut self, - ordering: LexOrdering, - constants: &[ConstExpr], - properties: &EquivalenceProperties, - ) -> AddedOrdering { - if ordering.is_empty() { - AddedOrdering::Yes - } else if properties.ordering_satisfy(ordering.as_ref()) { - // If the ordering satisfies the target properties, no need to - // augment it with constants. - self.orderings.push(ordering); - AddedOrdering::Yes - } else { - // Did not satisfy target properties, try and augment with constants - // to match the properties - if self.try_find_augmented_ordering(&ordering, constants, properties) { - AddedOrdering::Yes - } else { - AddedOrdering::No(ordering) - } - } - } - - /// Attempts to add `constants` to `ordering` to satisfy the properties. - /// - /// returns true if any orderings were added, false otherwise - fn try_find_augmented_ordering( - &mut self, - ordering: &LexOrdering, - constants: &[ConstExpr], - properties: &EquivalenceProperties, - ) -> bool { - // can't augment if there is nothing to augment with - if constants.is_empty() { - return false; - } - let start_num_orderings = self.orderings.len(); - - // for each equivalent ordering in properties, try and augment - // `ordering` it with the constants to match - for existing_ordering in properties.oeq_class.iter() { - if let Some(augmented_ordering) = self.augment_ordering( - ordering, - constants, - existing_ordering, - &properties.constants, - ) { - if !augmented_ordering.is_empty() { - assert!(properties.ordering_satisfy(augmented_ordering.as_ref())); - self.orderings.push(augmented_ordering); - } - } - } - - self.orderings.len() > start_num_orderings - } - - /// Attempts to augment the ordering with constants to match the - /// `existing_ordering` - /// - /// Returns Some(ordering) if an augmented ordering was found, None otherwise - fn augment_ordering( - &mut self, - ordering: &LexOrdering, - constants: &[ConstExpr], - existing_ordering: &LexOrdering, - existing_constants: &[ConstExpr], - ) -> Option { - let mut augmented_ordering = LexOrdering::default(); - let mut sort_expr_iter = ordering.iter().peekable(); - let mut existing_sort_expr_iter = existing_ordering.iter().peekable(); - - // walk in parallel down the two orderings, trying to match them up - while sort_expr_iter.peek().is_some() || existing_sort_expr_iter.peek().is_some() - { - // If the next expressions are equal, add the next match - // otherwise try and match with a constant - if let Some(expr) = - advance_if_match(&mut sort_expr_iter, &mut existing_sort_expr_iter) - { - augmented_ordering.push(expr); - } else if let Some(expr) = - advance_if_matches_constant(&mut sort_expr_iter, existing_constants) - { - augmented_ordering.push(expr); - } else if let Some(expr) = - advance_if_matches_constant(&mut existing_sort_expr_iter, constants) - { - augmented_ordering.push(expr); - } else { - // no match, can't continue the ordering, return what we have - break; - } - } - - Some(augmented_ordering) - } - - fn build(self) -> Vec { - self.orderings - } -} - -/// Advances two iterators in parallel -/// -/// If the next expressions are equal, the iterators are advanced and returns -/// the matched expression . -/// -/// Otherwise, the iterators are left unchanged and return `None` -fn advance_if_match( - iter1: &mut Peekable>, - iter2: &mut Peekable>, -) -> Option { - if matches!((iter1.peek(), iter2.peek()), (Some(expr1), Some(expr2)) if expr1.eq(expr2)) - { - iter1.next().unwrap(); - iter2.next().cloned() - } else { - None - } -} - -/// Advances the iterator with a constant -/// -/// If the next expression matches one of the constants, advances the iterator -/// returning the matched expression -/// -/// Otherwise, the iterator is left unchanged and returns `None` -fn advance_if_matches_constant( - iter: &mut Peekable>, - constants: &[ConstExpr], -) -> Option { - let expr = iter.peek()?; - let const_expr = constants.iter().find(|c| c.eq_expr(expr))?; - let found_expr = PhysicalSortExpr::new(Arc::clone(const_expr.expr()), expr.options); - iter.next(); - Some(found_expr) -} - -#[cfg(test)] -mod tests { - use std::ops::Not; - - use super::*; - use crate::equivalence::add_offset_to_expr; - use crate::equivalence::tests::{ - convert_to_orderings, convert_to_sort_exprs, convert_to_sort_reqs, - create_test_params, create_test_schema, output_schema, - }; - use crate::expressions::{col, BinaryExpr, Column}; - use crate::ScalarFunctionExpr; - - use arrow::datatypes::{DataType, Field, Fields, Schema, TimeUnit}; - use datafusion_common::{Constraint, ScalarValue}; - use datafusion_expr::Operator; - - use datafusion_functions::string::concat; - - #[test] - fn project_equivalence_properties_test() -> Result<()> { - let input_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, true), - Field::new("b", DataType::Int64, true), - Field::new("c", DataType::Int64, true), - ])); - - let input_properties = EquivalenceProperties::new(Arc::clone(&input_schema)); - let col_a = col("a", &input_schema)?; - - // a as a1, a as a2, a as a3, a as a3 - let proj_exprs = vec![ - (Arc::clone(&col_a), "a1".to_string()), - (Arc::clone(&col_a), "a2".to_string()), - (Arc::clone(&col_a), "a3".to_string()), - (Arc::clone(&col_a), "a4".to_string()), - ]; - let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; - - let out_schema = output_schema(&projection_mapping, &input_schema)?; - // a as a1, a as a2, a as a3, a as a3 - let proj_exprs = vec![ - (Arc::clone(&col_a), "a1".to_string()), - (Arc::clone(&col_a), "a2".to_string()), - (Arc::clone(&col_a), "a3".to_string()), - (Arc::clone(&col_a), "a4".to_string()), - ]; - let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; - - // a as a1, a as a2, a as a3, a as a3 - let col_a1 = &col("a1", &out_schema)?; - let col_a2 = &col("a2", &out_schema)?; - let col_a3 = &col("a3", &out_schema)?; - let col_a4 = &col("a4", &out_schema)?; - let out_properties = input_properties.project(&projection_mapping, out_schema); - - // At the output a1=a2=a3=a4 - assert_eq!(out_properties.eq_group().len(), 1); - let eq_class = out_properties.eq_group().iter().next().unwrap(); - assert_eq!(eq_class.len(), 4); - assert!(eq_class.contains(col_a1)); - assert!(eq_class.contains(col_a2)); - assert!(eq_class.contains(col_a3)); - assert!(eq_class.contains(col_a4)); - - Ok(()) - } - - #[test] - fn project_equivalence_properties_test_multi() -> Result<()> { - // test multiple input orderings with equivalence properties - let input_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, true), - Field::new("b", DataType::Int64, true), - Field::new("c", DataType::Int64, true), - Field::new("d", DataType::Int64, true), - ])); - - let mut input_properties = EquivalenceProperties::new(Arc::clone(&input_schema)); - // add equivalent ordering [a, b, c, d] - input_properties.add_new_ordering(LexOrdering::new(vec![ - parse_sort_expr("a", &input_schema), - parse_sort_expr("b", &input_schema), - parse_sort_expr("c", &input_schema), - parse_sort_expr("d", &input_schema), - ])); - - // add equivalent ordering [a, c, b, d] - input_properties.add_new_ordering(LexOrdering::new(vec![ - parse_sort_expr("a", &input_schema), - parse_sort_expr("c", &input_schema), - parse_sort_expr("b", &input_schema), // NB b and c are swapped - parse_sort_expr("d", &input_schema), - ])); - - // simply project all the columns in order - let proj_exprs = vec![ - (col("a", &input_schema)?, "a".to_string()), - (col("b", &input_schema)?, "b".to_string()), - (col("c", &input_schema)?, "c".to_string()), - (col("d", &input_schema)?, "d".to_string()), - ]; - let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; - let out_properties = input_properties.project(&projection_mapping, input_schema); - - assert_eq!( - out_properties.to_string(), - "order: [[a@0 ASC, c@2 ASC, b@1 ASC, d@3 ASC], [a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC]]" - ); - - Ok(()) - } - - #[test] - fn test_join_equivalence_properties() -> Result<()> { - let schema = create_test_schema()?; - let col_a = &col("a", &schema)?; - let col_b = &col("b", &schema)?; - let col_c = &col("c", &schema)?; - let offset = schema.fields.len(); - let col_a2 = &add_offset_to_expr(Arc::clone(col_a), offset); - let col_b2 = &add_offset_to_expr(Arc::clone(col_b), offset); - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - let test_cases = vec![ - // ------- TEST CASE 1 -------- - // [a ASC], [b ASC] - ( - // [a ASC], [b ASC] - vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], - // [a ASC], [b ASC] - vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], - // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC] - vec![ - vec![(col_a, option_asc), (col_a2, option_asc)], - vec![(col_a, option_asc), (col_b2, option_asc)], - vec![(col_b, option_asc), (col_a2, option_asc)], - vec![(col_b, option_asc), (col_b2, option_asc)], - ], - ), - // ------- TEST CASE 2 -------- - // [a ASC], [b ASC] - ( - // [a ASC], [b ASC], [c ASC] - vec![ - vec![(col_a, option_asc)], - vec![(col_b, option_asc)], - vec![(col_c, option_asc)], - ], - // [a ASC], [b ASC] - vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], - // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC], [c ASC, a2 ASC], [c ASC, b2 ASC] - vec![ - vec![(col_a, option_asc), (col_a2, option_asc)], - vec![(col_a, option_asc), (col_b2, option_asc)], - vec![(col_b, option_asc), (col_a2, option_asc)], - vec![(col_b, option_asc), (col_b2, option_asc)], - vec![(col_c, option_asc), (col_a2, option_asc)], - vec![(col_c, option_asc), (col_b2, option_asc)], - ], - ), - ]; - for (left_orderings, right_orderings, expected) in test_cases { - let mut left_eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - let mut right_eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - let left_orderings = convert_to_orderings(&left_orderings); - let right_orderings = convert_to_orderings(&right_orderings); - let expected = convert_to_orderings(&expected); - left_eq_properties.add_new_orderings(left_orderings); - right_eq_properties.add_new_orderings(right_orderings); - let join_eq = join_equivalence_properties( - left_eq_properties, - right_eq_properties, - &JoinType::Inner, - Arc::new(Schema::empty()), - &[true, false], - Some(JoinSide::Left), - &[], - ); - let err_msg = - format!("expected: {:?}, actual:{:?}", expected, &join_eq.oeq_class); - assert_eq!(join_eq.oeq_class.len(), expected.len(), "{}", err_msg); - for ordering in join_eq.oeq_class { - assert!( - expected.contains(&ordering), - "{}, ordering: {:?}", - err_msg, - ordering - ); - } - } - Ok(()) - } - - #[test] - fn test_expr_consists_of_constants() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - Field::new("d", DataType::Int32, true), - Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true), - ])); - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_d = col("d", &schema)?; - let b_plus_d = Arc::new(BinaryExpr::new( - Arc::clone(&col_b), - Operator::Plus, - Arc::clone(&col_d), - )) as Arc; - - let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b)]; - let expr = Arc::clone(&b_plus_d); - assert!(!is_constant_recurse(&constants, &expr)); - - let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_d)]; - let expr = Arc::clone(&b_plus_d); - assert!(is_constant_recurse(&constants, &expr)); - Ok(()) - } - - #[test] - fn test_get_updated_right_ordering_equivalence_properties() -> Result<()> { - let join_type = JoinType::Inner; - // Join right child schema - let child_fields: Fields = ["x", "y", "z", "w"] - .into_iter() - .map(|name| Field::new(name, DataType::Int32, true)) - .collect(); - let child_schema = Schema::new(child_fields); - let col_x = &col("x", &child_schema)?; - let col_y = &col("y", &child_schema)?; - let col_z = &col("z", &child_schema)?; - let col_w = &col("w", &child_schema)?; - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - // [x ASC, y ASC], [z ASC, w ASC] - let orderings = vec![ - vec![(col_x, option_asc), (col_y, option_asc)], - vec![(col_z, option_asc), (col_w, option_asc)], - ]; - let orderings = convert_to_orderings(&orderings); - // Right child ordering equivalences - let mut right_oeq_class = OrderingEquivalenceClass::new(orderings); - - let left_columns_len = 4; - - let fields: Fields = ["a", "b", "c", "d", "x", "y", "z", "w"] - .into_iter() - .map(|name| Field::new(name, DataType::Int32, true)) - .collect(); - - // Join Schema - let schema = Schema::new(fields); - let col_a = &col("a", &schema)?; - let col_d = &col("d", &schema)?; - let col_x = &col("x", &schema)?; - let col_y = &col("y", &schema)?; - let col_z = &col("z", &schema)?; - let col_w = &col("w", &schema)?; - - let mut join_eq_properties = EquivalenceProperties::new(Arc::new(schema)); - // a=x and d=w - join_eq_properties.add_equal_conditions(col_a, col_x)?; - join_eq_properties.add_equal_conditions(col_d, col_w)?; - - updated_right_ordering_equivalence_class( - &mut right_oeq_class, - &join_type, - left_columns_len, - ); - join_eq_properties.add_ordering_equivalence_class(right_oeq_class); - let result = join_eq_properties.oeq_class().clone(); - - // [x ASC, y ASC], [z ASC, w ASC] - let orderings = vec![ - vec![(col_x, option_asc), (col_y, option_asc)], - vec![(col_z, option_asc), (col_w, option_asc)], - ]; - let orderings = convert_to_orderings(&orderings); - let expected = OrderingEquivalenceClass::new(orderings); - - assert_eq!(result, expected); - - Ok(()) - } - - #[test] - fn test_normalize_ordering_equivalence_classes() -> Result<()> { - let sort_options = SortOptions::default(); - - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - ]); - let col_a_expr = col("a", &schema)?; - let col_b_expr = col("b", &schema)?; - let col_c_expr = col("c", &schema)?; - let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone())); - - eq_properties.add_equal_conditions(&col_a_expr, &col_c_expr)?; - let others = vec![ - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_b_expr), - options: sort_options, - }]), - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_c_expr), - options: sort_options, - }]), - ]; - eq_properties.add_new_orderings(others); - - let mut expected_eqs = EquivalenceProperties::new(Arc::new(schema)); - expected_eqs.add_new_orderings([ - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_b_expr), - options: sort_options, - }]), - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_c_expr), - options: sort_options, - }]), - ]); - - let oeq_class = eq_properties.oeq_class().clone(); - let expected = expected_eqs.oeq_class(); - assert!(oeq_class.eq(expected)); - - Ok(()) - } - - #[test] - fn test_get_indices_of_matching_sort_exprs_with_order_eq() -> Result<()> { - let sort_options = SortOptions::default(); - let sort_options_not = SortOptions::default().not(); - - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - ]); - let col_a = &col("a", &schema)?; - let col_b = &col("b", &schema)?; - let required_columns = [Arc::clone(col_b), Arc::clone(col_a)]; - let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::new(Column::new("b", 1)), - options: sort_options_not, - }, - PhysicalSortExpr { - expr: Arc::new(Column::new("a", 0)), - options: sort_options, - }, - ])]); - let (result, idxs) = eq_properties.find_longest_permutation(&required_columns); - assert_eq!(idxs, vec![0, 1]); - assert_eq!( - result, - LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(col_b), - options: sort_options_not - }, - PhysicalSortExpr { - expr: Arc::clone(col_a), - options: sort_options - } - ]) - ); - - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - ]); - let col_a = &col("a", &schema)?; - let col_b = &col("b", &schema)?; - let required_columns = [Arc::clone(col_b), Arc::clone(col_a)]; - let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); - eq_properties.add_new_orderings([ - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("c", 2)), - options: sort_options, - }]), - LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::new(Column::new("b", 1)), - options: sort_options_not, - }, - PhysicalSortExpr { - expr: Arc::new(Column::new("a", 0)), - options: sort_options, - }, - ]), - ]); - let (result, idxs) = eq_properties.find_longest_permutation(&required_columns); - assert_eq!(idxs, vec![0, 1]); - assert_eq!( - result, - LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(col_b), - options: sort_options_not - }, - PhysicalSortExpr { - expr: Arc::clone(col_a), - options: sort_options - } - ]) - ); - - let required_columns = [ - Arc::new(Column::new("b", 1)) as _, - Arc::new(Column::new("a", 0)) as _, - ]; - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - ]); - let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); - - // not satisfied orders - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::new(Column::new("b", 1)), - options: sort_options_not, - }, - PhysicalSortExpr { - expr: Arc::new(Column::new("c", 2)), - options: sort_options, - }, - PhysicalSortExpr { - expr: Arc::new(Column::new("a", 0)), - options: sort_options, - }, - ])]); - let (_, idxs) = eq_properties.find_longest_permutation(&required_columns); - assert_eq!(idxs, vec![0]); - - Ok(()) - } - - #[test] - fn test_update_properties() -> Result<()> { - let schema = Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - Field::new("d", DataType::Int32, true), - ]); - - let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone())); - let col_a = &col("a", &schema)?; - let col_b = &col("b", &schema)?; - let col_c = &col("c", &schema)?; - let col_d = &col("d", &schema)?; - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - // b=a (e.g they are aliases) - eq_properties.add_equal_conditions(col_b, col_a)?; - // [b ASC], [d ASC] - eq_properties.add_new_orderings(vec![ - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(col_b), - options: option_asc, - }]), - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(col_d), - options: option_asc, - }]), - ]); - - let test_cases = vec![ - // d + b - ( - Arc::new(BinaryExpr::new( - Arc::clone(col_d), - Operator::Plus, - Arc::clone(col_b), - )) as Arc, - SortProperties::Ordered(option_asc), - ), - // b - (Arc::clone(col_b), SortProperties::Ordered(option_asc)), - // a - (Arc::clone(col_a), SortProperties::Ordered(option_asc)), - // a + c - ( - Arc::new(BinaryExpr::new( - Arc::clone(col_a), - Operator::Plus, - Arc::clone(col_c), - )), - SortProperties::Unordered, - ), - ]; - for (expr, expected) in test_cases { - let leading_orderings = eq_properties - .oeq_class() - .iter() - .flat_map(|ordering| ordering.first().cloned()) - .collect::>(); - let expr_props = eq_properties.get_expr_properties(Arc::clone(&expr)); - let err_msg = format!( - "expr:{:?}, expected: {:?}, actual: {:?}, leading_orderings: {leading_orderings:?}", - expr, expected, expr_props.sort_properties - ); - assert_eq!(expr_props.sort_properties, expected, "{}", err_msg); - } - - Ok(()) - } - - #[test] - fn test_find_longest_permutation() -> Result<()> { - // Schema satisfies following orderings: - // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] - // and - // Column [a=c] (e.g they are aliases). - // At below we add [d ASC, h DESC] also, for test purposes - let (test_schema, mut eq_properties) = create_test_params()?; - let col_a = &col("a", &test_schema)?; - let col_b = &col("b", &test_schema)?; - let col_c = &col("c", &test_schema)?; - let col_d = &col("d", &test_schema)?; - let col_e = &col("e", &test_schema)?; - let col_f = &col("f", &test_schema)?; - let col_h = &col("h", &test_schema)?; - // a + d - let a_plus_d = Arc::new(BinaryExpr::new( - Arc::clone(col_a), - Operator::Plus, - Arc::clone(col_d), - )) as Arc; - - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - let option_desc = SortOptions { - descending: true, - nulls_first: true, - }; - // [d ASC, h DESC] also satisfies schema. - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(col_d), - options: option_asc, - }, - PhysicalSortExpr { - expr: Arc::clone(col_h), - options: option_desc, - }, - ])]); - let test_cases = vec![ - // TEST CASE 1 - (vec![col_a], vec![(col_a, option_asc)]), - // TEST CASE 2 - (vec![col_c], vec![(col_c, option_asc)]), - // TEST CASE 3 - ( - vec![col_d, col_e, col_b], - vec![ - (col_d, option_asc), - (col_e, option_desc), - (col_b, option_asc), - ], - ), - // TEST CASE 4 - (vec![col_b], vec![]), - // TEST CASE 5 - (vec![col_d], vec![(col_d, option_asc)]), - // TEST CASE 5 - (vec![&a_plus_d], vec![(&a_plus_d, option_asc)]), - // TEST CASE 6 - ( - vec![col_b, col_d], - vec![(col_d, option_asc), (col_b, option_asc)], - ), - // TEST CASE 6 - ( - vec![col_c, col_e], - vec![(col_c, option_asc), (col_e, option_desc)], - ), - // TEST CASE 7 - ( - vec![col_d, col_h, col_e, col_f, col_b], - vec![ - (col_d, option_asc), - (col_e, option_desc), - (col_h, option_desc), - (col_f, option_asc), - (col_b, option_asc), - ], - ), - // TEST CASE 8 - ( - vec![col_e, col_d, col_h, col_f, col_b], - vec![ - (col_e, option_desc), - (col_d, option_asc), - (col_h, option_desc), - (col_f, option_asc), - (col_b, option_asc), - ], - ), - // TEST CASE 9 - ( - vec![col_e, col_d, col_b, col_h, col_f], - vec![ - (col_e, option_desc), - (col_d, option_asc), - (col_b, option_asc), - (col_h, option_desc), - (col_f, option_asc), - ], - ), - ]; - for (exprs, expected) in test_cases { - let exprs = exprs.into_iter().cloned().collect::>(); - let expected = convert_to_sort_exprs(&expected); - let (actual, _) = eq_properties.find_longest_permutation(&exprs); - assert_eq!(actual, expected); - } - - Ok(()) - } - - #[test] - fn test_find_longest_permutation2() -> Result<()> { - // Schema satisfies following orderings: - // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] - // and - // Column [a=c] (e.g they are aliases). - // At below we add [d ASC, h DESC] also, for test purposes - let (test_schema, mut eq_properties) = create_test_params()?; - let col_h = &col("h", &test_schema)?; - - // Add column h as constant - eq_properties = eq_properties.with_constants(vec![ConstExpr::from(col_h)]); - - let test_cases = vec![ - // TEST CASE 1 - // ordering of the constants are treated as default ordering. - // This is the convention currently used. - (vec![col_h], vec![(col_h, SortOptions::default())]), - ]; - for (exprs, expected) in test_cases { - let exprs = exprs.into_iter().cloned().collect::>(); - let expected = convert_to_sort_exprs(&expected); - let (actual, _) = eq_properties.find_longest_permutation(&exprs); - assert_eq!(actual, expected); - } - - Ok(()) - } - - #[test] - fn test_get_finer() -> Result<()> { - let schema = create_test_schema()?; - let col_a = &col("a", &schema)?; - let col_b = &col("b", &schema)?; - let col_c = &col("c", &schema)?; - let eq_properties = EquivalenceProperties::new(schema); - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - let option_desc = SortOptions { - descending: true, - nulls_first: true, - }; - // First entry, and second entry are the physical sort requirement that are argument for get_finer_requirement. - // Third entry is the expected result. - let tests_cases = vec![ - // Get finer requirement between [a Some(ASC)] and [a None, b Some(ASC)] - // result should be [a Some(ASC), b Some(ASC)] - ( - vec![(col_a, Some(option_asc))], - vec![(col_a, None), (col_b, Some(option_asc))], - Some(vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))]), - ), - // Get finer requirement between [a Some(ASC), b Some(ASC), c Some(ASC)] and [a Some(ASC), b Some(ASC)] - // result should be [a Some(ASC), b Some(ASC), c Some(ASC)] - ( - vec![ - (col_a, Some(option_asc)), - (col_b, Some(option_asc)), - (col_c, Some(option_asc)), - ], - vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))], - Some(vec![ - (col_a, Some(option_asc)), - (col_b, Some(option_asc)), - (col_c, Some(option_asc)), - ]), - ), - // Get finer requirement between [a Some(ASC), b Some(ASC)] and [a Some(ASC), b Some(DESC)] - // result should be None - ( - vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))], - vec![(col_a, Some(option_asc)), (col_b, Some(option_desc))], - None, - ), - ]; - for (lhs, rhs, expected) in tests_cases { - let lhs = convert_to_sort_reqs(&lhs); - let rhs = convert_to_sort_reqs(&rhs); - let expected = expected.map(|expected| convert_to_sort_reqs(&expected)); - let finer = eq_properties.get_finer_requirement(&lhs, &rhs); - assert_eq!(finer, expected) - } - - Ok(()) - } - - #[test] - fn test_normalize_sort_reqs() -> Result<()> { - // Schema satisfies following properties - // a=c - // and following orderings are valid - // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] - let (test_schema, eq_properties) = create_test_params()?; - let col_a = &col("a", &test_schema)?; - let col_b = &col("b", &test_schema)?; - let col_c = &col("c", &test_schema)?; - let col_d = &col("d", &test_schema)?; - let col_e = &col("e", &test_schema)?; - let col_f = &col("f", &test_schema)?; - let option_asc = SortOptions { - descending: false, - nulls_first: false, - }; - let option_desc = SortOptions { - descending: true, - nulls_first: true, - }; - // First element in the tuple stores vector of requirement, second element is the expected return value for ordering_satisfy function - let requirements = vec![ - ( - vec![(col_a, Some(option_asc))], - vec![(col_a, Some(option_asc))], - ), - ( - vec![(col_a, Some(option_desc))], - vec![(col_a, Some(option_desc))], - ), - (vec![(col_a, None)], vec![(col_a, None)]), - // Test whether equivalence works as expected - ( - vec![(col_c, Some(option_asc))], - vec![(col_a, Some(option_asc))], - ), - (vec![(col_c, None)], vec![(col_a, None)]), - // Test whether ordering equivalence works as expected - ( - vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))], - vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))], - ), - ( - vec![(col_d, None), (col_b, None)], - vec![(col_d, None), (col_b, None)], - ), - ( - vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))], - vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))], - ), - // We should be able to normalize in compatible requirements also (not exactly equal) - ( - vec![(col_e, Some(option_desc)), (col_f, None)], - vec![(col_e, Some(option_desc)), (col_f, None)], - ), - ( - vec![(col_e, None), (col_f, None)], - vec![(col_e, None), (col_f, None)], - ), - ]; - - for (reqs, expected_normalized) in requirements.into_iter() { - let req = convert_to_sort_reqs(&reqs); - let expected_normalized = convert_to_sort_reqs(&expected_normalized); - - assert_eq!( - eq_properties.normalize_sort_requirements(&req), - expected_normalized - ); - } - - Ok(()) - } - - #[test] - fn test_schema_normalize_sort_requirement_with_equivalence() -> Result<()> { - let option1 = SortOptions { - descending: false, - nulls_first: false, - }; - // Assume that column a and c are aliases. - let (test_schema, eq_properties) = create_test_params()?; - let col_a = &col("a", &test_schema)?; - let col_c = &col("c", &test_schema)?; - let col_d = &col("d", &test_schema)?; - - // Test cases for equivalence normalization - // First entry in the tuple is PhysicalSortRequirement, second entry in the tuple is - // expected PhysicalSortRequirement after normalization. - let test_cases = vec![ - (vec![(col_a, Some(option1))], vec![(col_a, Some(option1))]), - // In the normalized version column c should be replace with column a - (vec![(col_c, Some(option1))], vec![(col_a, Some(option1))]), - (vec![(col_c, None)], vec![(col_a, None)]), - (vec![(col_d, Some(option1))], vec![(col_d, Some(option1))]), - ]; - for (reqs, expected) in test_cases.into_iter() { - let reqs = convert_to_sort_reqs(&reqs); - let expected = convert_to_sort_reqs(&expected); - - let normalized = eq_properties.normalize_sort_requirements(&reqs); - assert!( - expected.eq(&normalized), - "error in test: reqs: {reqs:?}, expected: {expected:?}, normalized: {normalized:?}" - ); - } - - Ok(()) - } - - #[test] - fn test_eliminate_redundant_monotonic_sorts() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Date32, true), - Field::new("b", DataType::Utf8, true), - Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), - ])); - let base_properties = EquivalenceProperties::new(Arc::clone(&schema)) - .with_reorder(LexOrdering::new( - ["a", "b", "c"] - .into_iter() - .map(|c| { - col(c, schema.as_ref()).map(|expr| PhysicalSortExpr { - expr, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }) - }) - .collect::>>()?, - )); - - struct TestCase { - name: &'static str, - constants: Vec>, - equal_conditions: Vec<[Arc; 2]>, - sort_columns: &'static [&'static str], - should_satisfy_ordering: bool, - } - - let col_a = col("a", schema.as_ref())?; - let col_b = col("b", schema.as_ref())?; - let col_c = col("c", schema.as_ref())?; - let cast_c = Arc::new(CastExpr::new(col_c, DataType::Date32, None)); - - let cases = vec![ - TestCase { - name: "(a, b, c) -> (c)", - // b is constant, so it should be removed from the sort order - constants: vec![Arc::clone(&col_b)], - equal_conditions: vec![[ - Arc::clone(&cast_c) as Arc, - Arc::clone(&col_a), - ]], - sort_columns: &["c"], - should_satisfy_ordering: true, - }, - // Same test with above test, where equality order is swapped. - // Algorithm shouldn't depend on this order. - TestCase { - name: "(a, b, c) -> (c)", - // b is constant, so it should be removed from the sort order - constants: vec![col_b], - equal_conditions: vec![[ - Arc::clone(&col_a), - Arc::clone(&cast_c) as Arc, - ]], - sort_columns: &["c"], - should_satisfy_ordering: true, - }, - TestCase { - name: "not ordered because (b) is not constant", - // b is not constant anymore - constants: vec![], - // a and c are still compatible, but this is irrelevant since the original ordering is (a, b, c) - equal_conditions: vec![[ - Arc::clone(&cast_c) as Arc, - Arc::clone(&col_a), - ]], - sort_columns: &["c"], - should_satisfy_ordering: false, - }, - ]; - - for case in cases { - // Construct the equivalence properties in different orders - // to exercise different code paths - // (The resulting properties _should_ be the same) - for properties in [ - // Equal conditions before constants - { - let mut properties = base_properties.clone(); - for [left, right] in &case.equal_conditions { - properties.add_equal_conditions(left, right)? - } - properties.with_constants( - case.constants.iter().cloned().map(ConstExpr::from), - ) - }, - // Constants before equal conditions - { - let mut properties = base_properties.clone().with_constants( - case.constants.iter().cloned().map(ConstExpr::from), - ); - for [left, right] in &case.equal_conditions { - properties.add_equal_conditions(left, right)? - } - properties - }, - ] { - let sort = case - .sort_columns - .iter() - .map(|&name| { - col(name, &schema).map(|col| PhysicalSortExpr { - expr: col, - options: SortOptions::default(), - }) - }) - .collect::>()?; - - assert_eq!( - properties.ordering_satisfy(sort.as_ref()), - case.should_satisfy_ordering, - "failed test '{}'", - case.name - ); - } - } - - Ok(()) - } - - /// Return a new schema with the same types, but new field names - /// - /// The new field names are the old field names with `text` appended. - /// - /// For example, the schema "a", "b", "c" becomes "a1", "b1", "c1" - /// if `text` is "1". - fn append_fields(schema: &SchemaRef, text: &str) -> SchemaRef { - Arc::new(Schema::new( - schema - .fields() - .iter() - .map(|field| { - Field::new( - // Annotate name with `text`: - format!("{}{}", field.name(), text), - field.data_type().clone(), - field.is_nullable(), - ) - }) - .collect::>(), - )) - } - - #[test] - fn test_union_equivalence_properties_multi_children_1() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - let schema3 = append_fields(&schema, "2"); - UnionEquivalenceTest::new(&schema) - // Children 1 - .with_child_sort(vec![vec!["a", "b", "c"]], &schema) - // Children 2 - .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) - // Children 3 - .with_child_sort(vec![vec!["a2", "b2"]], &schema3) - .with_expected_sort(vec![vec!["a", "b"]]) - .run() - } - - #[test] - fn test_union_equivalence_properties_multi_children_2() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - let schema3 = append_fields(&schema, "2"); - UnionEquivalenceTest::new(&schema) - // Children 1 - .with_child_sort(vec![vec!["a", "b", "c"]], &schema) - // Children 2 - .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) - // Children 3 - .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3) - .with_expected_sort(vec![vec!["a", "b", "c"]]) - .run() - } - - #[test] - fn test_union_equivalence_properties_multi_children_3() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - let schema3 = append_fields(&schema, "2"); - UnionEquivalenceTest::new(&schema) - // Children 1 - .with_child_sort(vec![vec!["a", "b"]], &schema) - // Children 2 - .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) - // Children 3 - .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3) - .with_expected_sort(vec![vec!["a", "b"]]) - .run() - } - - #[test] - fn test_union_equivalence_properties_multi_children_4() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - let schema3 = append_fields(&schema, "2"); - UnionEquivalenceTest::new(&schema) - // Children 1 - .with_child_sort(vec![vec!["a", "b"]], &schema) - // Children 2 - .with_child_sort(vec![vec!["a1", "b1"]], &schema2) - // Children 3 - .with_child_sort(vec![vec!["b2", "c2"]], &schema3) - .with_expected_sort(vec![]) - .run() - } - - #[test] - fn test_union_equivalence_properties_multi_children_5() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - UnionEquivalenceTest::new(&schema) - // Children 1 - .with_child_sort(vec![vec!["a", "b"], vec!["c"]], &schema) - // Children 2 - .with_child_sort(vec![vec!["a1", "b1"], vec!["c1"]], &schema2) - .with_expected_sort(vec![vec!["a", "b"], vec!["c"]]) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_common_constants() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child: [a ASC], const [b, c] - vec![vec!["a"]], - vec!["b", "c"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child: [b ASC], const [a, c] - vec![vec!["b"]], - vec!["a", "c"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union expected orderings: [[a ASC], [b ASC]], const [c] - vec![vec!["a"], vec!["b"]], - vec!["c"], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_prefix() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child: [a ASC], const [] - vec![vec!["a"]], - vec![], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child: [a ASC, b ASC], const [] - vec![vec!["a", "b"]], - vec![], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [a ASC], const [] - vec![vec!["a"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_asc_desc_mismatch() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child: [a ASC], const [] - vec![vec!["a"]], - vec![], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [a DESC], const [] - vec![vec!["a DESC"]], - vec![], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union doesn't have any ordering or constant - vec![], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_different_schemas() { - let schema = create_test_schema().unwrap(); - let schema2 = append_fields(&schema, "1"); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child orderings: [a ASC], const [] - vec![vec!["a"]], - vec![], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [a1 ASC, b1 ASC], const [] - vec![vec!["a1", "b1"]], - vec![], - &schema2, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [a ASC] - // - // Note that a, and a1 are at the same index for their - // corresponding schemas. - vec![vec!["a"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_fill_gaps() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child orderings: [a ASC, c ASC], const [b] - vec![vec!["a", "c"]], - vec!["b"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [b ASC, c ASC], const [a] - vec![vec!["b", "c"]], - vec!["a"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [ - // [a ASC, b ASC, c ASC], - // [b ASC, a ASC, c ASC] - // ], const [] - vec![vec!["a", "b", "c"], vec!["b", "a", "c"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_no_fill_gaps() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child orderings: [a ASC, c ASC], const [d] // some other constant - vec![vec!["a", "c"]], - vec!["d"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [b ASC, c ASC], const [a] - vec![vec!["b", "c"]], - vec!["a"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [[a]] (only a is constant) - vec![vec!["a"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_fill_some_gaps() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child orderings: [c ASC], const [a, b] // some other constant - vec![vec!["c"]], - vec!["a", "b"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [a DESC, b], const [] - vec![vec!["a DESC", "b"]], - vec![], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [[a, b]] (can fill in the a/b with constants) - vec![vec!["a DESC", "b"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_fill_gaps_non_symmetric() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child orderings: [a ASC, c ASC], const [b] - vec![vec!["a", "c"]], - vec!["b"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child orderings: [b ASC, c ASC], const [a] - vec![vec!["b DESC", "c"]], - vec!["a"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: [ - // [a ASC, b ASC, c ASC], - // [b ASC, a ASC, c ASC] - // ], const [] - vec![vec!["a", "b DESC", "c"], vec!["b DESC", "a", "c"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_gap_fill_symmetric() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child: [a ASC, b ASC, d ASC], const [c] - vec![vec!["a", "b", "d"]], - vec!["c"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child: [a ASC, c ASC, d ASC], const [b] - vec![vec!["a", "c", "d"]], - vec!["b"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: - // [a, b, c, d] - // [a, c, b, d] - vec![vec!["a", "c", "b", "d"], vec!["a", "b", "c", "d"]], - vec![], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_gap_fill_and_common() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // First child: [a DESC, d ASC], const [b, c] - vec![vec!["a DESC", "d"]], - vec!["b", "c"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child: [a DESC, c ASC, d ASC], const [b] - vec![vec!["a DESC", "c", "d"]], - vec!["b"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: - // [a DESC, c, d] [b] - vec![vec!["a DESC", "c", "d"]], - vec!["b"], - ) - .run() - } - - #[test] - fn test_union_equivalence_properties_constants_middle_desc() { - let schema = create_test_schema().unwrap(); - UnionEquivalenceTest::new(&schema) - .with_child_sort_and_const_exprs( - // NB `b DESC` in the first child - // - // First child: [a ASC, b DESC, d ASC], const [c] - vec![vec!["a", "b DESC", "d"]], - vec!["c"], - &schema, - ) - .with_child_sort_and_const_exprs( - // Second child: [a ASC, c ASC, d ASC], const [b] - vec![vec!["a", "c", "d"]], - vec!["b"], - &schema, - ) - .with_expected_sort_and_const_exprs( - // Union orderings: - // [a, b, d] (c constant) - // [a, c, d] (b constant) - vec![vec!["a", "c", "b DESC", "d"], vec!["a", "b DESC", "c", "d"]], - vec![], - ) - .run() - } - - // TODO tests with multiple constants - - #[derive(Debug)] - struct UnionEquivalenceTest { - /// The schema of the output of the Union - output_schema: SchemaRef, - /// The equivalence properties of each child to the union - child_properties: Vec, - /// The expected output properties of the union. Must be set before - /// running `build` - expected_properties: Option, - } - - impl UnionEquivalenceTest { - fn new(output_schema: &SchemaRef) -> Self { - Self { - output_schema: Arc::clone(output_schema), - child_properties: vec![], - expected_properties: None, - } - } - - /// Add a union input with the specified orderings - /// - /// See [`Self::make_props`] for the format of the strings in `orderings` - fn with_child_sort( - mut self, - orderings: Vec>, - schema: &SchemaRef, - ) -> Self { - let properties = self.make_props(orderings, vec![], schema); - self.child_properties.push(properties); - self - } - - /// Add a union input with the specified orderings and constant - /// equivalences - /// - /// See [`Self::make_props`] for the format of the strings in - /// `orderings` and `constants` - fn with_child_sort_and_const_exprs( - mut self, - orderings: Vec>, - constants: Vec<&str>, - schema: &SchemaRef, - ) -> Self { - let properties = self.make_props(orderings, constants, schema); - self.child_properties.push(properties); - self - } - - /// Set the expected output sort order for the union of the children - /// - /// See [`Self::make_props`] for the format of the strings in `orderings` - fn with_expected_sort(mut self, orderings: Vec>) -> Self { - let properties = self.make_props(orderings, vec![], &self.output_schema); - self.expected_properties = Some(properties); - self - } - - /// Set the expected output sort order and constant expressions for the - /// union of the children - /// - /// See [`Self::make_props`] for the format of the strings in - /// `orderings` and `constants`. - fn with_expected_sort_and_const_exprs( - mut self, - orderings: Vec>, - constants: Vec<&str>, - ) -> Self { - let properties = self.make_props(orderings, constants, &self.output_schema); - self.expected_properties = Some(properties); - self - } - - /// compute the union's output equivalence properties from the child - /// properties, and compare them to the expected properties - fn run(self) { - let Self { - output_schema, - child_properties, - expected_properties, - } = self; - - let expected_properties = - expected_properties.expect("expected_properties not set"); - - // try all permutations of the children - // as the code treats lhs and rhs differently - for child_properties in child_properties - .iter() - .cloned() - .permutations(child_properties.len()) - { - println!("--- permutation ---"); - for c in &child_properties { - println!("{c}"); - } - let actual_properties = - calculate_union(child_properties, Arc::clone(&output_schema)) - .expect("failed to calculate union equivalence properties"); - assert_eq_properties_same( - &actual_properties, - &expected_properties, - format!( - "expected: {expected_properties:?}\nactual: {actual_properties:?}" - ), - ); - } - } - - /// Make equivalence properties for the specified columns named in orderings and constants - /// - /// orderings: strings formatted like `"a"` or `"a DESC"`. See [`parse_sort_expr`] - /// constants: strings formatted like `"a"`. - fn make_props( - &self, - orderings: Vec>, - constants: Vec<&str>, - schema: &SchemaRef, - ) -> EquivalenceProperties { - let orderings = orderings - .iter() - .map(|ordering| { - ordering - .iter() - .map(|name| parse_sort_expr(name, schema)) - .collect::() - }) - .collect::>(); - - let constants = constants - .iter() - .map(|col_name| ConstExpr::new(col(col_name, schema).unwrap())) - .collect::>(); - - EquivalenceProperties::new_with_orderings(Arc::clone(schema), &orderings) - .with_constants(constants) - } - } - - fn assert_eq_properties_same( - lhs: &EquivalenceProperties, - rhs: &EquivalenceProperties, - err_msg: String, - ) { - // Check whether constants are same - let lhs_constants = lhs.constants(); - let rhs_constants = rhs.constants(); - for rhs_constant in rhs_constants { - assert!( - const_exprs_contains(lhs_constants, rhs_constant.expr()), - "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" - ); - } - assert_eq!( - lhs_constants.len(), - rhs_constants.len(), - "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" - ); - - // Check whether orderings are same. - let lhs_orderings = lhs.oeq_class(); - let rhs_orderings = rhs.oeq_class(); - for rhs_ordering in rhs_orderings.iter() { - assert!( - lhs_orderings.contains(rhs_ordering), - "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" - ); - } - assert_eq!( - lhs_orderings.len(), - rhs_orderings.len(), - "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" - ); - } - - /// Converts a string to a physical sort expression - /// - /// # Example - /// * `"a"` -> (`"a"`, `SortOptions::default()`) - /// * `"a ASC"` -> (`"a"`, `SortOptions { descending: false, nulls_first: false }`) - fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr { - let mut parts = name.split_whitespace(); - let name = parts.next().expect("empty sort expression"); - let mut sort_expr = PhysicalSortExpr::new( - col(name, schema).expect("invalid column name"), - SortOptions::default(), - ); - - if let Some(options) = parts.next() { - sort_expr = match options { - "ASC" => sort_expr.asc(), - "DESC" => sort_expr.desc(), - _ => panic!( - "unknown sort options. Expected 'ASC' or 'DESC', got {}", - options - ), - } - } - - assert!( - parts.next().is_none(), - "unexpected tokens in column name. Expected 'name' / 'name ASC' / 'name DESC' but got '{name}'" - ); - - sort_expr - } - - #[test] - fn test_ordering_equivalence_with_lex_monotonic_concat() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Utf8, false), - ])); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - let a_concat_b: Arc = Arc::new(ScalarFunctionExpr::new( - "concat", - concat(), - vec![Arc::clone(&col_a), Arc::clone(&col_b)], - DataType::Utf8, - )); - - // Assume existing ordering is [c ASC, a ASC, b ASC] - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - eq_properties.add_new_ordering(LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ])); - - // Add equality condition c = concat(a, b) - eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; - - let orderings = eq_properties.oeq_class(); - - let expected_ordering1 = - LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc() - ]); - let expected_ordering2 = LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ]); - - // The ordering should be [c ASC] and [a ASC, b ASC] - assert_eq!(orderings.len(), 2); - assert!(orderings.contains(&expected_ordering1)); - assert!(orderings.contains(&expected_ordering2)); - - Ok(()) - } - - #[test] - fn test_ordering_equivalence_with_non_lex_monotonic_multiply() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - Field::new("c", DataType::Int32, false), - ])); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - let a_times_b: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&col_a), - Operator::Multiply, - Arc::clone(&col_b), - )); - - // Assume existing ordering is [c ASC, a ASC, b ASC] - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let initial_ordering = LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ]); - - eq_properties.add_new_ordering(initial_ordering.clone()); - - // Add equality condition c = a * b - eq_properties.add_equal_conditions(&col_c, &a_times_b)?; - - let orderings = eq_properties.oeq_class(); - - // The ordering should remain unchanged since multiplication is not lex-monotonic - assert_eq!(orderings.len(), 1); - assert!(orderings.contains(&initial_ordering)); - - Ok(()) - } - - #[test] - fn test_ordering_equivalence_with_concat_equality() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Utf8, false), - ])); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - let a_concat_b: Arc = Arc::new(ScalarFunctionExpr::new( - "concat", - concat(), - vec![Arc::clone(&col_a), Arc::clone(&col_b)], - DataType::Utf8, - )); - - // Assume existing ordering is [concat(a, b) ASC, a ASC, b ASC] - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - eq_properties.add_new_ordering(LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&a_concat_b)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ])); - - // Add equality condition c = concat(a, b) - eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; - - let orderings = eq_properties.oeq_class(); - - let expected_ordering1 = LexOrdering::from(vec![PhysicalSortExpr::new_default( - Arc::clone(&a_concat_b), - ) - .asc()]); - let expected_ordering2 = LexOrdering::from(vec![ - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ]); - - // The ordering should be [concat(a, b) ASC] and [a ASC, b ASC] - assert_eq!(orderings.len(), 2); - assert!(orderings.contains(&expected_ordering1)); - assert!(orderings.contains(&expected_ordering2)); - - Ok(()) - } - - #[test] - fn test_with_reorder_constant_filtering() -> Result<()> { - let schema = create_test_schema()?; - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - // Setup constant columns - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - eq_properties = eq_properties.with_constants([ConstExpr::from(&col_a)]); - - let sort_exprs = LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: SortOptions::default(), - }, - PhysicalSortExpr { - expr: Arc::clone(&col_b), - options: SortOptions::default(), - }, - ]); - - let result = eq_properties.with_reorder(sort_exprs); - - // Should only contain b since a is constant - assert_eq!(result.oeq_class().len(), 1); - let ordering = result.oeq_class().iter().next().unwrap(); - assert_eq!(ordering.len(), 1); - assert!(ordering[0].expr.eq(&col_b)); - - Ok(()) - } - - #[test] - fn test_with_reorder_preserve_suffix() -> Result<()> { - let schema = create_test_schema()?; - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - let asc = SortOptions::default(); - let desc = SortOptions { - descending: true, - nulls_first: true, - }; - - // Initial ordering: [a ASC, b DESC, c ASC] - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: asc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_b), - options: desc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_c), - options: asc, - }, - ])]); - - // New ordering: [a ASC] - let new_order = LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: asc, - }]); - - let result = eq_properties.with_reorder(new_order); - - // Should only contain [a ASC, b DESC, c ASC] - assert_eq!(result.oeq_class().len(), 1); - let ordering = result.oeq_class().iter().next().unwrap(); - assert_eq!(ordering.len(), 3); - assert!(ordering[0].expr.eq(&col_a)); - assert!(ordering[0].options.eq(&asc)); - assert!(ordering[1].expr.eq(&col_b)); - assert!(ordering[1].options.eq(&desc)); - assert!(ordering[2].expr.eq(&col_c)); - assert!(ordering[2].options.eq(&asc)); - - Ok(()) - } - - #[test] - fn test_with_reorder_equivalent_expressions() -> Result<()> { - let schema = create_test_schema()?; - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - // Make a and b equivalent - eq_properties.add_equal_conditions(&col_a, &col_b)?; - - let asc = SortOptions::default(); - - // Initial ordering: [a ASC, c ASC] - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: asc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_c), - options: asc, - }, - ])]); - - // New ordering: [b ASC] - let new_order = LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_b), - options: asc, - }]); - - let result = eq_properties.with_reorder(new_order); - - // Should only contain [b ASC, c ASC] - assert_eq!(result.oeq_class().len(), 1); - - // Verify orderings - let ordering = result.oeq_class().iter().next().unwrap(); - assert_eq!(ordering.len(), 2); - assert!(ordering[0].expr.eq(&col_b)); - assert!(ordering[0].options.eq(&asc)); - assert!(ordering[1].expr.eq(&col_c)); - assert!(ordering[1].options.eq(&asc)); - - Ok(()) - } - - #[test] - fn test_with_reorder_incompatible_prefix() -> Result<()> { - let schema = create_test_schema()?; - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - - let asc = SortOptions::default(); - let desc = SortOptions { - descending: true, - nulls_first: true, - }; - - // Initial ordering: [a ASC, b DESC] - eq_properties.add_new_orderings([LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: asc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_b), - options: desc, - }, - ])]); - - // New ordering: [a DESC] - let new_order = LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: desc, - }]); - - let result = eq_properties.with_reorder(new_order.clone()); - - // Should only contain the new ordering since options don't match - assert_eq!(result.oeq_class().len(), 1); - let ordering = result.oeq_class().iter().next().unwrap(); - assert_eq!(ordering, &new_order); - - Ok(()) - } - - #[test] - fn test_with_reorder_comprehensive() -> Result<()> { - let schema = create_test_schema()?; - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - let col_d = col("d", &schema)?; - let col_e = col("e", &schema)?; - - let asc = SortOptions::default(); - - // Constants: c is constant - eq_properties = eq_properties.with_constants([ConstExpr::from(&col_c)]); - - // Equality: b = d - eq_properties.add_equal_conditions(&col_b, &col_d)?; - - // Orderings: [d ASC, a ASC], [e ASC] - eq_properties.add_new_orderings([ - LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_d), - options: asc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: asc, - }, - ]), - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_e), - options: asc, - }]), - ]); - - // Initial ordering: [b ASC, c ASC] - let new_order = LexOrdering::new(vec![ - PhysicalSortExpr { - expr: Arc::clone(&col_b), - options: asc, - }, - PhysicalSortExpr { - expr: Arc::clone(&col_c), - options: asc, - }, - ]); - - let result = eq_properties.with_reorder(new_order); - - // Should preserve the original [d ASC, a ASC] ordering - assert_eq!(result.oeq_class().len(), 1); - let ordering = result.oeq_class().iter().next().unwrap(); - assert_eq!(ordering.len(), 2); - - // First expression should be either b or d (they're equivalent) - assert!( - ordering[0].expr.eq(&col_b) || ordering[0].expr.eq(&col_d), - "Expected b or d as first expression, got {:?}", - ordering[0].expr - ); - assert!(ordering[0].options.eq(&asc)); - - // Second expression should be a - assert!(ordering[1].expr.eq(&col_a)); - assert!(ordering[1].options.eq(&asc)); - - Ok(()) - } - - #[test] - fn test_union_constant_value_preservation() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - ])); - - let col_a = col("a", &schema)?; - let literal_10 = ScalarValue::Int32(Some(10)); - - // Create first input with a=10 - let const_expr1 = ConstExpr::new(Arc::clone(&col_a)) - .with_across_partitions(AcrossPartitions::Uniform(Some(literal_10.clone()))); - let input1 = EquivalenceProperties::new(Arc::clone(&schema)) - .with_constants(vec![const_expr1]); - - // Create second input with a=10 - let const_expr2 = ConstExpr::new(Arc::clone(&col_a)) - .with_across_partitions(AcrossPartitions::Uniform(Some(literal_10.clone()))); - let input2 = EquivalenceProperties::new(Arc::clone(&schema)) - .with_constants(vec![const_expr2]); - - // Calculate union properties - let union_props = calculate_union(vec![input1, input2], schema)?; - - // Verify column 'a' remains constant with value 10 - let const_a = &union_props.constants()[0]; - assert!(const_a.expr().eq(&col_a)); - assert_eq!( - const_a.across_partitions(), - AcrossPartitions::Uniform(Some(literal_10)) - ); - - Ok(()) - } - - #[test] - fn test_ordering_satisfaction_with_key_constraints() -> Result<()> { - let pk_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Int32, true), - Field::new("c", DataType::Int32, true), - Field::new("d", DataType::Int32, true), - ])); - - let unique_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Int32, false), - Field::new("c", DataType::Int32, true), - Field::new("d", DataType::Int32, true), - ])); - - // Test cases to run - let test_cases = vec![ - // (name, schema, constraint, base_ordering, satisfied_orderings, unsatisfied_orderings) - ( - "single column primary key", - &pk_schema, - vec![Constraint::PrimaryKey(vec![0])], - vec!["a"], // base ordering - vec![vec!["a", "b"], vec!["a", "c", "d"]], - vec![vec!["b", "a"], vec!["c", "a"]], - ), - ( - "single column unique", - &unique_schema, - vec![Constraint::Unique(vec![0])], - vec!["a"], // base ordering - vec![vec!["a", "b"], vec!["a", "c", "d"]], - vec![vec!["b", "a"], vec!["c", "a"]], - ), - ( - "multi-column primary key", - &pk_schema, - vec![Constraint::PrimaryKey(vec![0, 1])], - vec!["a", "b"], // base ordering - vec![vec!["a", "b", "c"], vec!["a", "b", "d"]], - vec![vec!["b", "a"], vec!["a", "c", "b"]], - ), - ( - "multi-column unique", - &unique_schema, - vec![Constraint::Unique(vec![0, 1])], - vec!["a", "b"], // base ordering - vec![vec!["a", "b", "c"], vec!["a", "b", "d"]], - vec![vec!["b", "a"], vec!["c", "a", "b"]], - ), - ( - "nullable unique", - &unique_schema, - vec![Constraint::Unique(vec![2, 3])], - vec!["c", "d"], // base ordering - vec![], - vec![vec!["c", "d", "a"]], - ), - ( - "ordering with arbitrary column unique", - &unique_schema, - vec![Constraint::Unique(vec![0, 1])], - vec!["a", "c", "b"], // base ordering - vec![vec!["a", "c", "b", "d"]], - vec![vec!["a", "b", "d"]], - ), - ( - "ordering with arbitrary column pk", - &pk_schema, - vec![Constraint::PrimaryKey(vec![0, 1])], - vec!["a", "c", "b"], // base ordering - vec![vec!["a", "c", "b", "d"]], - vec![vec!["a", "b", "d"]], - ), - ( - "ordering with arbitrary column pk complex", - &pk_schema, - vec![Constraint::PrimaryKey(vec![3, 1])], - vec!["b", "a", "d"], // base ordering - vec![vec!["b", "a", "d", "c"]], - vec![vec!["b", "c", "d", "a"], vec!["b", "a", "c", "d"]], - ), - ]; - - for ( - name, - schema, - constraints, - base_order, - satisfied_orders, - unsatisfied_orders, - ) in test_cases - { - let mut eq_properties = EquivalenceProperties::new(Arc::clone(schema)); - - // Convert base ordering - let base_ordering = LexOrdering::new( - base_order - .iter() - .map(|col_name| PhysicalSortExpr { - expr: col(col_name, schema).unwrap(), - options: SortOptions::default(), - }) - .collect(), - ); - - // Convert string column names to orderings - let satisfied_orderings: Vec = satisfied_orders - .iter() - .map(|cols| { - LexOrdering::new( - cols.iter() - .map(|col_name| PhysicalSortExpr { - expr: col(col_name, schema).unwrap(), - options: SortOptions::default(), - }) - .collect(), - ) - }) - .collect(); - - let unsatisfied_orderings: Vec = unsatisfied_orders - .iter() - .map(|cols| { - LexOrdering::new( - cols.iter() - .map(|col_name| PhysicalSortExpr { - expr: col(col_name, schema).unwrap(), - options: SortOptions::default(), - }) - .collect(), - ) - }) - .collect(); - - // Test that orderings are not satisfied before adding constraints - for ordering in &satisfied_orderings { - assert!( - !eq_properties.ordering_satisfy(ordering), - "{}: ordering {:?} should not be satisfied before adding constraints", - name, - ordering - ); - } - - // Add base ordering - eq_properties.add_new_ordering(base_ordering); - - // Add constraints - eq_properties = - eq_properties.with_constraints(Constraints::new_unverified(constraints)); - - // Test that expected orderings are now satisfied - for ordering in &satisfied_orderings { - assert!( - eq_properties.ordering_satisfy(ordering), - "{}: ordering {:?} should be satisfied after adding constraints", - name, - ordering - ); - } - - // Test that unsatisfied orderings remain unsatisfied - for ordering in &unsatisfied_orderings { - assert!( - !eq_properties.ordering_satisfy(ordering), - "{}: ordering {:?} should not be satisfied after adding constraints", - name, - ordering - ); - } - } - - Ok(()) - } -} diff --git a/datafusion/physical-expr/src/equivalence/properties/dependency.rs b/datafusion/physical-expr/src/equivalence/properties/dependency.rs new file mode 100644 index 000000000000..9eba295e562e --- /dev/null +++ b/datafusion/physical-expr/src/equivalence/properties/dependency.rs @@ -0,0 +1,1774 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display}; +use std::sync::Arc; + +use crate::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use indexmap::IndexSet; + +use indexmap::IndexMap; +use itertools::Itertools; + +use super::{expr_refers, ExprWrapper}; + +// A list of sort expressions that can be calculated from a known set of +/// dependencies. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct Dependencies { + inner: IndexSet, +} + +impl Display for Dependencies { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "[")?; + let mut iter = self.inner.iter(); + if let Some(dep) = iter.next() { + write!(f, "{}", dep)?; + } + for dep in iter { + write!(f, ", {}", dep)?; + } + write!(f, "]") + } +} + +impl Dependencies { + /// Create a new empty `Dependencies` instance. + fn new() -> Self { + Self { + inner: IndexSet::new(), + } + } + + /// Create a new `Dependencies` from an iterator of `PhysicalSortExpr`. + pub fn new_from_iter(iter: impl IntoIterator) -> Self { + Self { + inner: iter.into_iter().collect(), + } + } + + /// Insert a new dependency into the set. + pub fn insert(&mut self, sort_expr: PhysicalSortExpr) { + self.inner.insert(sort_expr); + } + + /// Iterator over dependencies in the set + pub fn iter(&self) -> impl Iterator + Clone { + self.inner.iter() + } + + /// Return the inner set of dependencies + pub fn into_inner(self) -> IndexSet { + self.inner + } + + /// Returns true if there are no dependencies + fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} + +/// Contains a mapping of all dependencies we have processed for each sort expr +pub struct DependencyEnumerator<'a> { + /// Maps `expr` --> `[exprs]` that have previously been processed + seen: IndexMap<&'a PhysicalSortExpr, IndexSet<&'a PhysicalSortExpr>>, +} + +impl<'a> DependencyEnumerator<'a> { + pub fn new() -> Self { + Self { + seen: IndexMap::new(), + } + } + + /// Insert a new dependency, + /// + /// returns false if the dependency was already in the map + /// returns true if the dependency was newly inserted + fn insert( + &mut self, + target: &'a PhysicalSortExpr, + dep: &'a PhysicalSortExpr, + ) -> bool { + self.seen.entry(target).or_default().insert(dep) + } + + /// This function recursively analyzes the dependencies of the given sort + /// expression within the given dependency map to construct lexicographical + /// orderings that include the sort expression and its dependencies. + /// + /// # Parameters + /// + /// - `referred_sort_expr`: A reference to the sort expression (`PhysicalSortExpr`) + /// for which lexicographical orderings satisfying its dependencies are to be + /// constructed. + /// - `dependency_map`: A reference to the `DependencyMap` that contains + /// dependencies for different `PhysicalSortExpr`s. + /// + /// # Returns + /// + /// A vector of lexicographical orderings (`Vec`) based on the given + /// sort expression and its dependencies. + pub fn construct_orderings( + &mut self, + referred_sort_expr: &'a PhysicalSortExpr, + dependency_map: &'a DependencyMap, + ) -> Vec { + let node = dependency_map + .get(referred_sort_expr) + .expect("`referred_sort_expr` should be inside `dependency_map`"); + // Since we work on intermediate nodes, we are sure `val.target_sort_expr` + // exists. + let target_sort_expr = node.target_sort_expr.as_ref().unwrap(); + // An empty dependency means the referred_sort_expr represents a global ordering. + // Return its projected version, which is the target_expression. + if node.dependencies.is_empty() { + return vec![LexOrdering::new(vec![target_sort_expr.clone()])]; + }; + + node.dependencies + .iter() + .flat_map(|dep| { + let mut orderings = if self.insert(target_sort_expr, dep) { + self.construct_orderings(dep, dependency_map) + } else { + vec![] + }; + + for ordering in orderings.iter_mut() { + ordering.push(target_sort_expr.clone()) + } + orderings + }) + .collect() + } +} + +/// Maps an expression --> DependencyNode +/// +/// # Debugging / deplaying `DependencyMap` +/// +/// This structure implements `Display` to assist debugging. For example: +/// +/// ```text +/// DependencyMap: { +/// a@0 ASC --> (target: a@0 ASC, dependencies: [[]]) +/// b@1 ASC --> (target: b@1 ASC, dependencies: [[a@0 ASC, c@2 ASC]]) +/// c@2 ASC --> (target: c@2 ASC, dependencies: [[b@1 ASC, a@0 ASC]]) +/// d@3 ASC --> (target: d@3 ASC, dependencies: [[c@2 ASC, b@1 ASC]]) +/// } +/// ``` +/// +/// # Note on IndexMap Rationale +/// +/// Using `IndexMap` (which preserves insert order) to ensure consistent results +/// across different executions for the same query. We could have used +/// `HashSet`, `HashMap` in place of them without any loss of functionality. +/// +/// As an example, if existing orderings are +/// 1. `[a ASC, b ASC]` +/// 2. `[c ASC]` for +/// +/// Then both the following output orderings are valid +/// 1. `[a ASC, b ASC, c ASC]` +/// 2. `[c ASC, a ASC, b ASC]` +/// +/// (this are both valid as they are concatenated versions of the alternative +/// orderings). When using `HashSet`, `HashMap` it is not guaranteed to generate +/// consistent result, among the possible 2 results in the example above. +#[derive(Debug)] +pub struct DependencyMap { + inner: IndexMap, +} + +impl DependencyMap { + pub fn new() -> Self { + Self { + inner: IndexMap::new(), + } + } + + /// Insert a new dependency `sort_expr` --> `dependency` into the map. + /// + /// If `target_sort_expr` is none, a new entry is created with empty dependencies. + pub fn insert( + &mut self, + sort_expr: &PhysicalSortExpr, + target_sort_expr: Option<&PhysicalSortExpr>, + dependency: Option<&PhysicalSortExpr>, + ) { + self.inner + .entry(sort_expr.clone()) + .or_insert_with(|| DependencyNode { + target_sort_expr: target_sort_expr.cloned(), + dependencies: Dependencies::new(), + }) + .insert_dependency(dependency) + } + + /// Iterator over (sort_expr, DependencyNode) pairs + pub fn iter(&self) -> impl Iterator { + self.inner.iter() + } + + /// iterator over all sort exprs + pub fn sort_exprs(&self) -> impl Iterator { + self.inner.keys() + } + + /// Return the dependency node for the given sort expression, if any + pub fn get(&self, sort_expr: &PhysicalSortExpr) -> Option<&DependencyNode> { + self.inner.get(sort_expr) + } +} + +impl Display for DependencyMap { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "DependencyMap: {{")?; + for (sort_expr, node) in self.inner.iter() { + writeln!(f, " {sort_expr} --> {node}")?; + } + writeln!(f, "}}") + } +} + +/// Represents a node in the dependency map used to construct projected orderings. +/// +/// A `DependencyNode` contains information about a particular sort expression, +/// including its target sort expression and a set of dependencies on other sort +/// expressions. +/// +/// # Fields +/// +/// - `target_sort_expr`: An optional `PhysicalSortExpr` representing the target +/// sort expression associated with the node. It is `None` if the sort expression +/// cannot be projected. +/// - `dependencies`: A [`Dependencies`] containing dependencies on other sort +/// expressions that are referred to by the target sort expression. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DependencyNode { + pub target_sort_expr: Option, + pub dependencies: Dependencies, +} + +impl DependencyNode { + /// Insert dependency to the state (if exists). + fn insert_dependency(&mut self, dependency: Option<&PhysicalSortExpr>) { + if let Some(dep) = dependency { + self.dependencies.insert(dep.clone()); + } + } +} + +impl Display for DependencyNode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if let Some(target) = &self.target_sort_expr { + write!(f, "(target: {}, ", target)?; + } else { + write!(f, "(")?; + } + write!(f, "dependencies: [{}])", self.dependencies) + } +} + +/// This function analyzes the dependency map to collect referred dependencies for +/// a given source expression. +/// +/// # Parameters +/// +/// - `dependency_map`: A reference to the `DependencyMap` where each +/// `PhysicalSortExpr` is associated with a `DependencyNode`. +/// - `source`: A reference to the source expression (`Arc`) +/// for which relevant dependencies need to be identified. +/// +/// # Returns +/// +/// A `Vec` containing the dependencies for the given source +/// expression. These dependencies are expressions that are referred to by +/// the source expression based on the provided dependency map. +pub fn referred_dependencies( + dependency_map: &DependencyMap, + source: &Arc, +) -> Vec { + // Associate `PhysicalExpr`s with `PhysicalSortExpr`s that contain them: + let mut expr_to_sort_exprs = IndexMap::::new(); + for sort_expr in dependency_map + .sort_exprs() + .filter(|sort_expr| expr_refers(source, &sort_expr.expr)) + { + let key = ExprWrapper(Arc::clone(&sort_expr.expr)); + expr_to_sort_exprs + .entry(key) + .or_default() + .insert(sort_expr.clone()); + } + + // Generate all valid dependencies for the source. For example, if the source + // is `a + b` and the map is `[a -> (a ASC, a DESC), b -> (b ASC)]`, we get + // `vec![HashSet(a ASC, b ASC), HashSet(a DESC, b ASC)]`. + let dependencies = expr_to_sort_exprs + .into_values() + .map(Dependencies::into_inner) + .collect::>(); + dependencies + .iter() + .multi_cartesian_product() + .map(|referred_deps| { + Dependencies::new_from_iter(referred_deps.into_iter().cloned()) + }) + .collect() +} + +/// This function retrieves the dependencies of the given relevant sort expression +/// from the given dependency map. It then constructs prefix orderings by recursively +/// analyzing the dependencies and include them in the orderings. +/// +/// # Parameters +/// +/// - `relevant_sort_expr`: A reference to the relevant sort expression +/// (`PhysicalSortExpr`) for which prefix orderings are to be constructed. +/// - `dependency_map`: A reference to the `DependencyMap` containing dependencies. +/// +/// # Returns +/// +/// A vector of prefix orderings (`Vec`) based on the given relevant +/// sort expression and its dependencies. +pub fn construct_prefix_orderings( + relevant_sort_expr: &PhysicalSortExpr, + dependency_map: &DependencyMap, +) -> Vec { + let mut dep_enumerator = DependencyEnumerator::new(); + dependency_map + .get(relevant_sort_expr) + .expect("no relevant sort expr found") + .dependencies + .iter() + .flat_map(|dep| dep_enumerator.construct_orderings(dep, dependency_map)) + .collect() +} + +/// Generates all possible orderings where dependencies are satisfied for the +/// current projection expression. +/// +/// # Example +/// If `dependencies` is `a + b ASC` and the dependency map holds dependencies +/// * `a ASC` --> `[c ASC]` +/// * `b ASC` --> `[d DESC]`, +/// +/// This function generates these two sort orders +/// * `[c ASC, d DESC, a + b ASC]` +/// * `[d DESC, c ASC, a + b ASC]` +/// +/// # Parameters +/// +/// * `dependencies` - Set of relevant expressions. +/// * `dependency_map` - Map of dependencies for expressions that may appear in `dependencies` +/// +/// # Returns +/// +/// A vector of lexical orderings (`Vec`) representing all valid orderings +/// based on the given dependencies. +pub fn generate_dependency_orderings( + dependencies: &Dependencies, + dependency_map: &DependencyMap, +) -> Vec { + // Construct all the valid prefix orderings for each expression appearing + // in the projection: + let relevant_prefixes = dependencies + .iter() + .flat_map(|dep| { + let prefixes = construct_prefix_orderings(dep, dependency_map); + (!prefixes.is_empty()).then_some(prefixes) + }) + .collect::>(); + + // No dependency, dependent is a leading ordering. + if relevant_prefixes.is_empty() { + // Return an empty ordering: + return vec![LexOrdering::default()]; + } + + relevant_prefixes + .into_iter() + .multi_cartesian_product() + .flat_map(|prefix_orderings| { + prefix_orderings + .iter() + .permutations(prefix_orderings.len()) + .map(|prefixes| { + prefixes + .into_iter() + .flat_map(|ordering| ordering.clone()) + .collect() + }) + .collect::>() + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::ops::Not; + use std::sync::Arc; + + use super::*; + use crate::equivalence::tests::{ + convert_to_sort_exprs, convert_to_sort_reqs, create_test_params, + create_test_schema, output_schema, parse_sort_expr, + }; + use crate::equivalence::ProjectionMapping; + use crate::expressions::{col, BinaryExpr, CastExpr, Column}; + use crate::{ConstExpr, EquivalenceProperties, ScalarFunctionExpr}; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion_common::{Constraint, Constraints, Result}; + use datafusion_expr::sort_properties::SortProperties; + use datafusion_expr::Operator; + + use datafusion_functions::string::concat; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + + #[test] + fn project_equivalence_properties_test() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + Field::new("c", DataType::Int64, true), + ])); + + let input_properties = EquivalenceProperties::new(Arc::clone(&input_schema)); + let col_a = col("a", &input_schema)?; + + // a as a1, a as a2, a as a3, a as a3 + let proj_exprs = vec![ + (Arc::clone(&col_a), "a1".to_string()), + (Arc::clone(&col_a), "a2".to_string()), + (Arc::clone(&col_a), "a3".to_string()), + (Arc::clone(&col_a), "a4".to_string()), + ]; + let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; + + let out_schema = output_schema(&projection_mapping, &input_schema)?; + // a as a1, a as a2, a as a3, a as a3 + let proj_exprs = vec![ + (Arc::clone(&col_a), "a1".to_string()), + (Arc::clone(&col_a), "a2".to_string()), + (Arc::clone(&col_a), "a3".to_string()), + (Arc::clone(&col_a), "a4".to_string()), + ]; + let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; + + // a as a1, a as a2, a as a3, a as a3 + let col_a1 = &col("a1", &out_schema)?; + let col_a2 = &col("a2", &out_schema)?; + let col_a3 = &col("a3", &out_schema)?; + let col_a4 = &col("a4", &out_schema)?; + let out_properties = input_properties.project(&projection_mapping, out_schema); + + // At the output a1=a2=a3=a4 + assert_eq!(out_properties.eq_group().len(), 1); + let eq_class = out_properties.eq_group().iter().next().unwrap(); + assert_eq!(eq_class.len(), 4); + assert!(eq_class.contains(col_a1)); + assert!(eq_class.contains(col_a2)); + assert!(eq_class.contains(col_a3)); + assert!(eq_class.contains(col_a4)); + + Ok(()) + } + + #[test] + fn project_equivalence_properties_test_multi() -> Result<()> { + // test multiple input orderings with equivalence properties + let input_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + Field::new("c", DataType::Int64, true), + Field::new("d", DataType::Int64, true), + ])); + + let mut input_properties = EquivalenceProperties::new(Arc::clone(&input_schema)); + // add equivalent ordering [a, b, c, d] + input_properties.add_new_ordering(LexOrdering::new(vec![ + parse_sort_expr("a", &input_schema), + parse_sort_expr("b", &input_schema), + parse_sort_expr("c", &input_schema), + parse_sort_expr("d", &input_schema), + ])); + + // add equivalent ordering [a, c, b, d] + input_properties.add_new_ordering(LexOrdering::new(vec![ + parse_sort_expr("a", &input_schema), + parse_sort_expr("c", &input_schema), + parse_sort_expr("b", &input_schema), // NB b and c are swapped + parse_sort_expr("d", &input_schema), + ])); + + // simply project all the columns in order + let proj_exprs = vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + (col("c", &input_schema)?, "c".to_string()), + (col("d", &input_schema)?, "d".to_string()), + ]; + let projection_mapping = ProjectionMapping::try_new(&proj_exprs, &input_schema)?; + let out_properties = input_properties.project(&projection_mapping, input_schema); + + assert_eq!( + out_properties.to_string(), + "order: [[a@0 ASC, c@2 ASC, b@1 ASC, d@3 ASC], [a@0 ASC, b@1 ASC, c@2 ASC, d@3 ASC]]" + ); + + Ok(()) + } + + #[test] + fn test_normalize_ordering_equivalence_classes() -> Result<()> { + let sort_options = SortOptions::default(); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let col_a_expr = col("a", &schema)?; + let col_b_expr = col("b", &schema)?; + let col_c_expr = col("c", &schema)?; + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone())); + + eq_properties.add_equal_conditions(&col_a_expr, &col_c_expr)?; + let others = vec![ + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_b_expr), + options: sort_options, + }]), + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_c_expr), + options: sort_options, + }]), + ]; + eq_properties.add_new_orderings(others); + + let mut expected_eqs = EquivalenceProperties::new(Arc::new(schema)); + expected_eqs.add_new_orderings([ + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_b_expr), + options: sort_options, + }]), + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_c_expr), + options: sort_options, + }]), + ]); + + let oeq_class = eq_properties.oeq_class().clone(); + let expected = expected_eqs.oeq_class(); + assert!(oeq_class.eq(expected)); + + Ok(()) + } + + #[test] + fn test_get_indices_of_matching_sort_exprs_with_order_eq() -> Result<()> { + let sort_options = SortOptions::default(); + let sort_options_not = SortOptions::default().not(); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]); + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let required_columns = [Arc::clone(col_b), Arc::clone(col_a)]; + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::new(Column::new("b", 1)), + options: sort_options_not, + }, + PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: sort_options, + }, + ])]); + let (result, idxs) = eq_properties.find_longest_permutation(&required_columns); + assert_eq!(idxs, vec![0, 1]); + assert_eq!( + result, + LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(col_b), + options: sort_options_not + }, + PhysicalSortExpr { + expr: Arc::clone(col_a), + options: sort_options + } + ]) + ); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let required_columns = [Arc::clone(col_b), Arc::clone(col_a)]; + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); + eq_properties.add_new_orderings([ + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::new(Column::new("c", 2)), + options: sort_options, + }]), + LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::new(Column::new("b", 1)), + options: sort_options_not, + }, + PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: sort_options, + }, + ]), + ]); + let (result, idxs) = eq_properties.find_longest_permutation(&required_columns); + assert_eq!(idxs, vec![0, 1]); + assert_eq!( + result, + LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(col_b), + options: sort_options_not + }, + PhysicalSortExpr { + expr: Arc::clone(col_a), + options: sort_options + } + ]) + ); + + let required_columns = [ + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("a", 0)) as _, + ]; + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema)); + + // not satisfied orders + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::new(Column::new("b", 1)), + options: sort_options_not, + }, + PhysicalSortExpr { + expr: Arc::new(Column::new("c", 2)), + options: sort_options, + }, + PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: sort_options, + }, + ])]); + let (_, idxs) = eq_properties.find_longest_permutation(&required_columns); + assert_eq!(idxs, vec![0]); + + Ok(()) + } + + #[test] + fn test_update_properties() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + ]); + + let mut eq_properties = EquivalenceProperties::new(Arc::new(schema.clone())); + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let col_c = &col("c", &schema)?; + let col_d = &col("d", &schema)?; + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + // b=a (e.g they are aliases) + eq_properties.add_equal_conditions(col_b, col_a)?; + // [b ASC], [d ASC] + eq_properties.add_new_orderings(vec![ + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(col_b), + options: option_asc, + }]), + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(col_d), + options: option_asc, + }]), + ]); + + let test_cases = vec![ + // d + b + ( + Arc::new(BinaryExpr::new( + Arc::clone(col_d), + Operator::Plus, + Arc::clone(col_b), + )) as Arc, + SortProperties::Ordered(option_asc), + ), + // b + (Arc::clone(col_b), SortProperties::Ordered(option_asc)), + // a + (Arc::clone(col_a), SortProperties::Ordered(option_asc)), + // a + c + ( + Arc::new(BinaryExpr::new( + Arc::clone(col_a), + Operator::Plus, + Arc::clone(col_c), + )), + SortProperties::Unordered, + ), + ]; + for (expr, expected) in test_cases { + let leading_orderings = eq_properties + .oeq_class() + .iter() + .flat_map(|ordering| ordering.first().cloned()) + .collect::>(); + let expr_props = eq_properties.get_expr_properties(Arc::clone(&expr)); + let err_msg = format!( + "expr:{:?}, expected: {:?}, actual: {:?}, leading_orderings: {leading_orderings:?}", + expr, expected, expr_props.sort_properties + ); + assert_eq!(expr_props.sort_properties, expected, "{}", err_msg); + } + + Ok(()) + } + + #[test] + fn test_find_longest_permutation() -> Result<()> { + // Schema satisfies following orderings: + // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] + // and + // Column [a=c] (e.g they are aliases). + // At below we add [d ASC, h DESC] also, for test purposes + let (test_schema, mut eq_properties) = create_test_params()?; + let col_a = &col("a", &test_schema)?; + let col_b = &col("b", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let col_d = &col("d", &test_schema)?; + let col_e = &col("e", &test_schema)?; + let col_f = &col("f", &test_schema)?; + let col_h = &col("h", &test_schema)?; + // a + d + let a_plus_d = Arc::new(BinaryExpr::new( + Arc::clone(col_a), + Operator::Plus, + Arc::clone(col_d), + )) as Arc; + + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let option_desc = SortOptions { + descending: true, + nulls_first: true, + }; + // [d ASC, h DESC] also satisfies schema. + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(col_d), + options: option_asc, + }, + PhysicalSortExpr { + expr: Arc::clone(col_h), + options: option_desc, + }, + ])]); + let test_cases = vec![ + // TEST CASE 1 + (vec![col_a], vec![(col_a, option_asc)]), + // TEST CASE 2 + (vec![col_c], vec![(col_c, option_asc)]), + // TEST CASE 3 + ( + vec![col_d, col_e, col_b], + vec![ + (col_d, option_asc), + (col_e, option_desc), + (col_b, option_asc), + ], + ), + // TEST CASE 4 + (vec![col_b], vec![]), + // TEST CASE 5 + (vec![col_d], vec![(col_d, option_asc)]), + // TEST CASE 5 + (vec![&a_plus_d], vec![(&a_plus_d, option_asc)]), + // TEST CASE 6 + ( + vec![col_b, col_d], + vec![(col_d, option_asc), (col_b, option_asc)], + ), + // TEST CASE 6 + ( + vec![col_c, col_e], + vec![(col_c, option_asc), (col_e, option_desc)], + ), + // TEST CASE 7 + ( + vec![col_d, col_h, col_e, col_f, col_b], + vec![ + (col_d, option_asc), + (col_e, option_desc), + (col_h, option_desc), + (col_f, option_asc), + (col_b, option_asc), + ], + ), + // TEST CASE 8 + ( + vec![col_e, col_d, col_h, col_f, col_b], + vec![ + (col_e, option_desc), + (col_d, option_asc), + (col_h, option_desc), + (col_f, option_asc), + (col_b, option_asc), + ], + ), + // TEST CASE 9 + ( + vec![col_e, col_d, col_b, col_h, col_f], + vec![ + (col_e, option_desc), + (col_d, option_asc), + (col_b, option_asc), + (col_h, option_desc), + (col_f, option_asc), + ], + ), + ]; + for (exprs, expected) in test_cases { + let exprs = exprs.into_iter().cloned().collect::>(); + let expected = convert_to_sort_exprs(&expected); + let (actual, _) = eq_properties.find_longest_permutation(&exprs); + assert_eq!(actual, expected); + } + + Ok(()) + } + + #[test] + fn test_find_longest_permutation2() -> Result<()> { + // Schema satisfies following orderings: + // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] + // and + // Column [a=c] (e.g they are aliases). + // At below we add [d ASC, h DESC] also, for test purposes + let (test_schema, mut eq_properties) = create_test_params()?; + let col_h = &col("h", &test_schema)?; + + // Add column h as constant + eq_properties = eq_properties.with_constants(vec![ConstExpr::from(col_h)]); + + let test_cases = vec![ + // TEST CASE 1 + // ordering of the constants are treated as default ordering. + // This is the convention currently used. + (vec![col_h], vec![(col_h, SortOptions::default())]), + ]; + for (exprs, expected) in test_cases { + let exprs = exprs.into_iter().cloned().collect::>(); + let expected = convert_to_sort_exprs(&expected); + let (actual, _) = eq_properties.find_longest_permutation(&exprs); + assert_eq!(actual, expected); + } + + Ok(()) + } + + #[test] + fn test_get_finer() -> Result<()> { + let schema = create_test_schema()?; + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let col_c = &col("c", &schema)?; + let eq_properties = EquivalenceProperties::new(schema); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let option_desc = SortOptions { + descending: true, + nulls_first: true, + }; + // First entry, and second entry are the physical sort requirement that are argument for get_finer_requirement. + // Third entry is the expected result. + let tests_cases = vec![ + // Get finer requirement between [a Some(ASC)] and [a None, b Some(ASC)] + // result should be [a Some(ASC), b Some(ASC)] + ( + vec![(col_a, Some(option_asc))], + vec![(col_a, None), (col_b, Some(option_asc))], + Some(vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))]), + ), + // Get finer requirement between [a Some(ASC), b Some(ASC), c Some(ASC)] and [a Some(ASC), b Some(ASC)] + // result should be [a Some(ASC), b Some(ASC), c Some(ASC)] + ( + vec![ + (col_a, Some(option_asc)), + (col_b, Some(option_asc)), + (col_c, Some(option_asc)), + ], + vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))], + Some(vec![ + (col_a, Some(option_asc)), + (col_b, Some(option_asc)), + (col_c, Some(option_asc)), + ]), + ), + // Get finer requirement between [a Some(ASC), b Some(ASC)] and [a Some(ASC), b Some(DESC)] + // result should be None + ( + vec![(col_a, Some(option_asc)), (col_b, Some(option_asc))], + vec![(col_a, Some(option_asc)), (col_b, Some(option_desc))], + None, + ), + ]; + for (lhs, rhs, expected) in tests_cases { + let lhs = convert_to_sort_reqs(&lhs); + let rhs = convert_to_sort_reqs(&rhs); + let expected = expected.map(|expected| convert_to_sort_reqs(&expected)); + let finer = eq_properties.get_finer_requirement(&lhs, &rhs); + assert_eq!(finer, expected) + } + + Ok(()) + } + + #[test] + fn test_normalize_sort_reqs() -> Result<()> { + // Schema satisfies following properties + // a=c + // and following orderings are valid + // [a ASC], [d ASC, b ASC], [e DESC, f ASC, g ASC] + let (test_schema, eq_properties) = create_test_params()?; + let col_a = &col("a", &test_schema)?; + let col_b = &col("b", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let col_d = &col("d", &test_schema)?; + let col_e = &col("e", &test_schema)?; + let col_f = &col("f", &test_schema)?; + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let option_desc = SortOptions { + descending: true, + nulls_first: true, + }; + // First element in the tuple stores vector of requirement, second element is the expected return value for ordering_satisfy function + let requirements = vec![ + ( + vec![(col_a, Some(option_asc))], + vec![(col_a, Some(option_asc))], + ), + ( + vec![(col_a, Some(option_desc))], + vec![(col_a, Some(option_desc))], + ), + (vec![(col_a, None)], vec![(col_a, None)]), + // Test whether equivalence works as expected + ( + vec![(col_c, Some(option_asc))], + vec![(col_a, Some(option_asc))], + ), + (vec![(col_c, None)], vec![(col_a, None)]), + // Test whether ordering equivalence works as expected + ( + vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))], + vec![(col_d, Some(option_asc)), (col_b, Some(option_asc))], + ), + ( + vec![(col_d, None), (col_b, None)], + vec![(col_d, None), (col_b, None)], + ), + ( + vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))], + vec![(col_e, Some(option_desc)), (col_f, Some(option_asc))], + ), + // We should be able to normalize in compatible requirements also (not exactly equal) + ( + vec![(col_e, Some(option_desc)), (col_f, None)], + vec![(col_e, Some(option_desc)), (col_f, None)], + ), + ( + vec![(col_e, None), (col_f, None)], + vec![(col_e, None), (col_f, None)], + ), + ]; + + for (reqs, expected_normalized) in requirements.into_iter() { + let req = convert_to_sort_reqs(&reqs); + let expected_normalized = convert_to_sort_reqs(&expected_normalized); + + assert_eq!( + eq_properties.normalize_sort_requirements(&req), + expected_normalized + ); + } + + Ok(()) + } + + #[test] + fn test_schema_normalize_sort_requirement_with_equivalence() -> Result<()> { + let option1 = SortOptions { + descending: false, + nulls_first: false, + }; + // Assume that column a and c are aliases. + let (test_schema, eq_properties) = create_test_params()?; + let col_a = &col("a", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let col_d = &col("d", &test_schema)?; + + // Test cases for equivalence normalization + // First entry in the tuple is PhysicalSortRequirement, second entry in the tuple is + // expected PhysicalSortRequirement after normalization. + let test_cases = vec![ + (vec![(col_a, Some(option1))], vec![(col_a, Some(option1))]), + // In the normalized version column c should be replace with column a + (vec![(col_c, Some(option1))], vec![(col_a, Some(option1))]), + (vec![(col_c, None)], vec![(col_a, None)]), + (vec![(col_d, Some(option1))], vec![(col_d, Some(option1))]), + ]; + for (reqs, expected) in test_cases.into_iter() { + let reqs = convert_to_sort_reqs(&reqs); + let expected = convert_to_sort_reqs(&expected); + + let normalized = eq_properties.normalize_sort_requirements(&reqs); + assert!( + expected.eq(&normalized), + "error in test: reqs: {reqs:?}, expected: {expected:?}, normalized: {normalized:?}" + ); + } + + Ok(()) + } + + #[test] + fn test_eliminate_redundant_monotonic_sorts() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let base_properties = EquivalenceProperties::new(Arc::clone(&schema)) + .with_reorder(LexOrdering::new( + ["a", "b", "c"] + .into_iter() + .map(|c| { + col(c, schema.as_ref()).map(|expr| PhysicalSortExpr { + expr, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }) + }) + .collect::>>()?, + )); + + struct TestCase { + name: &'static str, + constants: Vec>, + equal_conditions: Vec<[Arc; 2]>, + sort_columns: &'static [&'static str], + should_satisfy_ordering: bool, + } + + let col_a = col("a", schema.as_ref())?; + let col_b = col("b", schema.as_ref())?; + let col_c = col("c", schema.as_ref())?; + let cast_c = Arc::new(CastExpr::new(col_c, DataType::Date32, None)); + + let cases = vec![ + TestCase { + name: "(a, b, c) -> (c)", + // b is constant, so it should be removed from the sort order + constants: vec![Arc::clone(&col_b)], + equal_conditions: vec![[ + Arc::clone(&cast_c) as Arc, + Arc::clone(&col_a), + ]], + sort_columns: &["c"], + should_satisfy_ordering: true, + }, + // Same test with above test, where equality order is swapped. + // Algorithm shouldn't depend on this order. + TestCase { + name: "(a, b, c) -> (c)", + // b is constant, so it should be removed from the sort order + constants: vec![col_b], + equal_conditions: vec![[ + Arc::clone(&col_a), + Arc::clone(&cast_c) as Arc, + ]], + sort_columns: &["c"], + should_satisfy_ordering: true, + }, + TestCase { + name: "not ordered because (b) is not constant", + // b is not constant anymore + constants: vec![], + // a and c are still compatible, but this is irrelevant since the original ordering is (a, b, c) + equal_conditions: vec![[ + Arc::clone(&cast_c) as Arc, + Arc::clone(&col_a), + ]], + sort_columns: &["c"], + should_satisfy_ordering: false, + }, + ]; + + for case in cases { + // Construct the equivalence properties in different orders + // to exercise different code paths + // (The resulting properties _should_ be the same) + for properties in [ + // Equal conditions before constants + { + let mut properties = base_properties.clone(); + for [left, right] in &case.equal_conditions { + properties.add_equal_conditions(left, right)? + } + properties.with_constants( + case.constants.iter().cloned().map(ConstExpr::from), + ) + }, + // Constants before equal conditions + { + let mut properties = base_properties.clone().with_constants( + case.constants.iter().cloned().map(ConstExpr::from), + ); + for [left, right] in &case.equal_conditions { + properties.add_equal_conditions(left, right)? + } + properties + }, + ] { + let sort = case + .sort_columns + .iter() + .map(|&name| { + col(name, &schema).map(|col| PhysicalSortExpr { + expr: col, + options: SortOptions::default(), + }) + }) + .collect::>()?; + + assert_eq!( + properties.ordering_satisfy(sort.as_ref()), + case.should_satisfy_ordering, + "failed test '{}'", + case.name + ); + } + } + + Ok(()) + } + + #[test] + fn test_ordering_equivalence_with_lex_monotonic_concat() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::Utf8, false), + ])); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + + let a_concat_b: Arc = Arc::new(ScalarFunctionExpr::new( + "concat", + concat(), + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + DataType::Utf8, + )); + + // Assume existing ordering is [c ASC, a ASC, b ASC] + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + eq_properties.add_new_ordering(LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), + ])); + + // Add equality condition c = concat(a, b) + eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; + + let orderings = eq_properties.oeq_class(); + + let expected_ordering1 = + LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc() + ]); + let expected_ordering2 = LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), + ]); + + // The ordering should be [c ASC] and [a ASC, b ASC] + assert_eq!(orderings.len(), 2); + assert!(orderings.contains(&expected_ordering1)); + assert!(orderings.contains(&expected_ordering2)); + + Ok(()) + } + + #[test] + fn test_ordering_equivalence_with_non_lex_monotonic_multiply() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + + let a_times_b: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + Operator::Multiply, + Arc::clone(&col_b), + )); + + // Assume existing ordering is [c ASC, a ASC, b ASC] + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let initial_ordering = LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), + ]); + + eq_properties.add_new_ordering(initial_ordering.clone()); + + // Add equality condition c = a * b + eq_properties.add_equal_conditions(&col_c, &a_times_b)?; + + let orderings = eq_properties.oeq_class(); + + // The ordering should remain unchanged since multiplication is not lex-monotonic + assert_eq!(orderings.len(), 1); + assert!(orderings.contains(&initial_ordering)); + + Ok(()) + } + + #[test] + fn test_ordering_equivalence_with_concat_equality() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::Utf8, false), + ])); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + + let a_concat_b: Arc = Arc::new(ScalarFunctionExpr::new( + "concat", + concat(), + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + DataType::Utf8, + )); + + // Assume existing ordering is [concat(a, b) ASC, a ASC, b ASC] + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + eq_properties.add_new_ordering(LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&a_concat_b)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), + ])); + + // Add equality condition c = concat(a, b) + eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; + + let orderings = eq_properties.oeq_class(); + + let expected_ordering1 = LexOrdering::from(vec![PhysicalSortExpr::new_default( + Arc::clone(&a_concat_b), + ) + .asc()]); + let expected_ordering2 = LexOrdering::from(vec![ + PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), + PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), + ]); + + // The ordering should be [concat(a, b) ASC] and [a ASC, b ASC] + assert_eq!(orderings.len(), 2); + assert!(orderings.contains(&expected_ordering1)); + assert!(orderings.contains(&expected_ordering2)); + + Ok(()) + } + + #[test] + fn test_with_reorder_constant_filtering() -> Result<()> { + let schema = create_test_schema()?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + // Setup constant columns + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + eq_properties = eq_properties.with_constants([ConstExpr::from(&col_a)]); + + let sort_exprs = LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: Arc::clone(&col_b), + options: SortOptions::default(), + }, + ]); + + let result = eq_properties.with_reorder(sort_exprs); + + // Should only contain b since a is constant + assert_eq!(result.oeq_class().len(), 1); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 1); + assert!(ordering[0].expr.eq(&col_b)); + + Ok(()) + } + + #[test] + fn test_with_reorder_preserve_suffix() -> Result<()> { + let schema = create_test_schema()?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + + let asc = SortOptions::default(); + let desc = SortOptions { + descending: true, + nulls_first: true, + }; + + // Initial ordering: [a ASC, b DESC, c ASC] + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: asc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_b), + options: desc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_c), + options: asc, + }, + ])]); + + // New ordering: [a ASC] + let new_order = LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: asc, + }]); + + let result = eq_properties.with_reorder(new_order); + + // Should only contain [a ASC, b DESC, c ASC] + assert_eq!(result.oeq_class().len(), 1); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 3); + assert!(ordering[0].expr.eq(&col_a)); + assert!(ordering[0].options.eq(&asc)); + assert!(ordering[1].expr.eq(&col_b)); + assert!(ordering[1].options.eq(&desc)); + assert!(ordering[2].expr.eq(&col_c)); + assert!(ordering[2].options.eq(&asc)); + + Ok(()) + } + + #[test] + fn test_with_reorder_equivalent_expressions() -> Result<()> { + let schema = create_test_schema()?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + + // Make a and b equivalent + eq_properties.add_equal_conditions(&col_a, &col_b)?; + + let asc = SortOptions::default(); + + // Initial ordering: [a ASC, c ASC] + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: asc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_c), + options: asc, + }, + ])]); + + // New ordering: [b ASC] + let new_order = LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_b), + options: asc, + }]); + + let result = eq_properties.with_reorder(new_order); + + // Should only contain [b ASC, c ASC] + assert_eq!(result.oeq_class().len(), 1); + + // Verify orderings + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 2); + assert!(ordering[0].expr.eq(&col_b)); + assert!(ordering[0].options.eq(&asc)); + assert!(ordering[1].expr.eq(&col_c)); + assert!(ordering[1].options.eq(&asc)); + + Ok(()) + } + + #[test] + fn test_with_reorder_incompatible_prefix() -> Result<()> { + let schema = create_test_schema()?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + + let asc = SortOptions::default(); + let desc = SortOptions { + descending: true, + nulls_first: true, + }; + + // Initial ordering: [a ASC, b DESC] + eq_properties.add_new_orderings([LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: asc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_b), + options: desc, + }, + ])]); + + // New ordering: [a DESC] + let new_order = LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: desc, + }]); + + let result = eq_properties.with_reorder(new_order.clone()); + + // Should only contain the new ordering since options don't match + assert_eq!(result.oeq_class().len(), 1); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering, &new_order); + + Ok(()) + } + + #[test] + fn test_with_reorder_comprehensive() -> Result<()> { + let schema = create_test_schema()?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_c = col("c", &schema)?; + let col_d = col("d", &schema)?; + let col_e = col("e", &schema)?; + + let asc = SortOptions::default(); + + // Constants: c is constant + eq_properties = eq_properties.with_constants([ConstExpr::from(&col_c)]); + + // Equality: b = d + eq_properties.add_equal_conditions(&col_b, &col_d)?; + + // Orderings: [d ASC, a ASC], [e ASC] + eq_properties.add_new_orderings([ + LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_d), + options: asc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: asc, + }, + ]), + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_e), + options: asc, + }]), + ]); + + // Initial ordering: [b ASC, c ASC] + let new_order = LexOrdering::new(vec![ + PhysicalSortExpr { + expr: Arc::clone(&col_b), + options: asc, + }, + PhysicalSortExpr { + expr: Arc::clone(&col_c), + options: asc, + }, + ]); + + let result = eq_properties.with_reorder(new_order); + + // Should preserve the original [d ASC, a ASC] ordering + assert_eq!(result.oeq_class().len(), 1); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 2); + + // First expression should be either b or d (they're equivalent) + assert!( + ordering[0].expr.eq(&col_b) || ordering[0].expr.eq(&col_d), + "Expected b or d as first expression, got {:?}", + ordering[0].expr + ); + assert!(ordering[0].options.eq(&asc)); + + // Second expression should be a + assert!(ordering[1].expr.eq(&col_a)); + assert!(ordering[1].options.eq(&asc)); + + Ok(()) + } + + #[test] + fn test_ordering_satisfaction_with_key_constraints() -> Result<()> { + let pk_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + ])); + + let unique_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + ])); + + // Test cases to run + let test_cases = vec![ + // (name, schema, constraint, base_ordering, satisfied_orderings, unsatisfied_orderings) + ( + "single column primary key", + &pk_schema, + vec![Constraint::PrimaryKey(vec![0])], + vec!["a"], // base ordering + vec![vec!["a", "b"], vec!["a", "c", "d"]], + vec![vec!["b", "a"], vec!["c", "a"]], + ), + ( + "single column unique", + &unique_schema, + vec![Constraint::Unique(vec![0])], + vec!["a"], // base ordering + vec![vec!["a", "b"], vec!["a", "c", "d"]], + vec![vec!["b", "a"], vec!["c", "a"]], + ), + ( + "multi-column primary key", + &pk_schema, + vec![Constraint::PrimaryKey(vec![0, 1])], + vec!["a", "b"], // base ordering + vec![vec!["a", "b", "c"], vec!["a", "b", "d"]], + vec![vec!["b", "a"], vec!["a", "c", "b"]], + ), + ( + "multi-column unique", + &unique_schema, + vec![Constraint::Unique(vec![0, 1])], + vec!["a", "b"], // base ordering + vec![vec!["a", "b", "c"], vec!["a", "b", "d"]], + vec![vec!["b", "a"], vec!["c", "a", "b"]], + ), + ( + "nullable unique", + &unique_schema, + vec![Constraint::Unique(vec![2, 3])], + vec!["c", "d"], // base ordering + vec![], + vec![vec!["c", "d", "a"]], + ), + ( + "ordering with arbitrary column unique", + &unique_schema, + vec![Constraint::Unique(vec![0, 1])], + vec!["a", "c", "b"], // base ordering + vec![vec!["a", "c", "b", "d"]], + vec![vec!["a", "b", "d"]], + ), + ( + "ordering with arbitrary column pk", + &pk_schema, + vec![Constraint::PrimaryKey(vec![0, 1])], + vec!["a", "c", "b"], // base ordering + vec![vec!["a", "c", "b", "d"]], + vec![vec!["a", "b", "d"]], + ), + ( + "ordering with arbitrary column pk complex", + &pk_schema, + vec![Constraint::PrimaryKey(vec![3, 1])], + vec!["b", "a", "d"], // base ordering + vec![vec!["b", "a", "d", "c"]], + vec![vec!["b", "c", "d", "a"], vec!["b", "a", "c", "d"]], + ), + ]; + + for ( + name, + schema, + constraints, + base_order, + satisfied_orders, + unsatisfied_orders, + ) in test_cases + { + let mut eq_properties = EquivalenceProperties::new(Arc::clone(schema)); + + // Convert base ordering + let base_ordering = LexOrdering::new( + base_order + .iter() + .map(|col_name| PhysicalSortExpr { + expr: col(col_name, schema).unwrap(), + options: SortOptions::default(), + }) + .collect(), + ); + + // Convert string column names to orderings + let satisfied_orderings: Vec = satisfied_orders + .iter() + .map(|cols| { + LexOrdering::new( + cols.iter() + .map(|col_name| PhysicalSortExpr { + expr: col(col_name, schema).unwrap(), + options: SortOptions::default(), + }) + .collect(), + ) + }) + .collect(); + + let unsatisfied_orderings: Vec = unsatisfied_orders + .iter() + .map(|cols| { + LexOrdering::new( + cols.iter() + .map(|col_name| PhysicalSortExpr { + expr: col(col_name, schema).unwrap(), + options: SortOptions::default(), + }) + .collect(), + ) + }) + .collect(); + + // Test that orderings are not satisfied before adding constraints + for ordering in &satisfied_orderings { + assert!( + !eq_properties.ordering_satisfy(ordering), + "{}: ordering {:?} should not be satisfied before adding constraints", + name, + ordering + ); + } + + // Add base ordering + eq_properties.add_new_ordering(base_ordering); + + // Add constraints + eq_properties = + eq_properties.with_constraints(Constraints::new_unverified(constraints)); + + // Test that expected orderings are now satisfied + for ordering in &satisfied_orderings { + assert!( + eq_properties.ordering_satisfy(ordering), + "{}: ordering {:?} should be satisfied after adding constraints", + name, + ordering + ); + } + + // Test that unsatisfied orderings remain unsatisfied + for ordering in &unsatisfied_orderings { + assert!( + !eq_properties.ordering_satisfy(ordering), + "{}: ordering {:?} should not be satisfied after adding constraints", + name, + ordering + ); + } + } + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/equivalence/properties/joins.rs b/datafusion/physical-expr/src/equivalence/properties/joins.rs new file mode 100644 index 000000000000..7944e89d0305 --- /dev/null +++ b/datafusion/physical-expr/src/equivalence/properties/joins.rs @@ -0,0 +1,301 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::{equivalence::OrderingEquivalenceClass, PhysicalExprRef}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{JoinSide, JoinType}; + +use super::EquivalenceProperties; + +/// Calculate ordering equivalence properties for the given join operation. +pub fn join_equivalence_properties( + left: EquivalenceProperties, + right: EquivalenceProperties, + join_type: &JoinType, + join_schema: SchemaRef, + maintains_input_order: &[bool], + probe_side: Option, + on: &[(PhysicalExprRef, PhysicalExprRef)], +) -> EquivalenceProperties { + let left_size = left.schema.fields.len(); + let mut result = EquivalenceProperties::new(join_schema); + result.add_equivalence_group(left.eq_group().join( + right.eq_group(), + join_type, + left_size, + on, + )); + + let EquivalenceProperties { + constants: left_constants, + oeq_class: left_oeq_class, + .. + } = left; + let EquivalenceProperties { + constants: right_constants, + oeq_class: mut right_oeq_class, + .. + } = right; + match maintains_input_order { + [true, false] => { + // In this special case, right side ordering can be prefixed with + // the left side ordering. + if let (Some(JoinSide::Left), JoinType::Inner) = (probe_side, join_type) { + updated_right_ordering_equivalence_class( + &mut right_oeq_class, + join_type, + left_size, + ); + + // Right side ordering equivalence properties should be prepended + // with those of the left side while constructing output ordering + // equivalence properties since stream side is the left side. + // + // For example, if the right side ordering equivalences contain + // `b ASC`, and the left side ordering equivalences contain `a ASC`, + // then we should add `a ASC, b ASC` to the ordering equivalences + // of the join output. + let out_oeq_class = left_oeq_class.join_suffix(&right_oeq_class); + result.add_ordering_equivalence_class(out_oeq_class); + } else { + result.add_ordering_equivalence_class(left_oeq_class); + } + } + [false, true] => { + updated_right_ordering_equivalence_class( + &mut right_oeq_class, + join_type, + left_size, + ); + // In this special case, left side ordering can be prefixed with + // the right side ordering. + if let (Some(JoinSide::Right), JoinType::Inner) = (probe_side, join_type) { + // Left side ordering equivalence properties should be prepended + // with those of the right side while constructing output ordering + // equivalence properties since stream side is the right side. + // + // For example, if the left side ordering equivalences contain + // `a ASC`, and the right side ordering equivalences contain `b ASC`, + // then we should add `b ASC, a ASC` to the ordering equivalences + // of the join output. + let out_oeq_class = right_oeq_class.join_suffix(&left_oeq_class); + result.add_ordering_equivalence_class(out_oeq_class); + } else { + result.add_ordering_equivalence_class(right_oeq_class); + } + } + [false, false] => {} + [true, true] => unreachable!("Cannot maintain ordering of both sides"), + _ => unreachable!("Join operators can not have more than two children"), + } + match join_type { + JoinType::LeftAnti | JoinType::LeftSemi => { + result = result.with_constants(left_constants); + } + JoinType::RightAnti | JoinType::RightSemi => { + result = result.with_constants(right_constants); + } + _ => {} + } + result +} + +/// In the context of a join, update the right side `OrderingEquivalenceClass` +/// so that they point to valid indices in the join output schema. +/// +/// To do so, we increment column indices by the size of the left table when +/// join schema consists of a combination of the left and right schemas. This +/// is the case for `Inner`, `Left`, `Full` and `Right` joins. For other cases, +/// indices do not change. +pub fn updated_right_ordering_equivalence_class( + right_oeq_class: &mut OrderingEquivalenceClass, + join_type: &JoinType, + left_size: usize, +) { + if matches!( + join_type, + JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right + ) { + right_oeq_class.add_offset(left_size); + } +} + +#[cfg(test)] +mod tests { + + use std::sync::Arc; + + use super::*; + use crate::equivalence::add_offset_to_expr; + use crate::equivalence::tests::{convert_to_orderings, create_test_schema}; + use crate::expressions::col; + use datafusion_common::Result; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + + #[test] + fn test_join_equivalence_properties() -> Result<()> { + let schema = create_test_schema()?; + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let col_c = &col("c", &schema)?; + let offset = schema.fields.len(); + let col_a2 = &add_offset_to_expr(Arc::clone(col_a), offset); + let col_b2 = &add_offset_to_expr(Arc::clone(col_b), offset); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let test_cases = vec![ + // ------- TEST CASE 1 -------- + // [a ASC], [b ASC] + ( + // [a ASC], [b ASC] + vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], + // [a ASC], [b ASC] + vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], + // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC] + vec![ + vec![(col_a, option_asc), (col_a2, option_asc)], + vec![(col_a, option_asc), (col_b2, option_asc)], + vec![(col_b, option_asc), (col_a2, option_asc)], + vec![(col_b, option_asc), (col_b2, option_asc)], + ], + ), + // ------- TEST CASE 2 -------- + // [a ASC], [b ASC] + ( + // [a ASC], [b ASC], [c ASC] + vec![ + vec![(col_a, option_asc)], + vec![(col_b, option_asc)], + vec![(col_c, option_asc)], + ], + // [a ASC], [b ASC] + vec![vec![(col_a, option_asc)], vec![(col_b, option_asc)]], + // expected [a ASC, a2 ASC], [a ASC, b2 ASC], [b ASC, a2 ASC], [b ASC, b2 ASC], [c ASC, a2 ASC], [c ASC, b2 ASC] + vec![ + vec![(col_a, option_asc), (col_a2, option_asc)], + vec![(col_a, option_asc), (col_b2, option_asc)], + vec![(col_b, option_asc), (col_a2, option_asc)], + vec![(col_b, option_asc), (col_b2, option_asc)], + vec![(col_c, option_asc), (col_a2, option_asc)], + vec![(col_c, option_asc), (col_b2, option_asc)], + ], + ), + ]; + for (left_orderings, right_orderings, expected) in test_cases { + let mut left_eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let mut right_eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let left_orderings = convert_to_orderings(&left_orderings); + let right_orderings = convert_to_orderings(&right_orderings); + let expected = convert_to_orderings(&expected); + left_eq_properties.add_new_orderings(left_orderings); + right_eq_properties.add_new_orderings(right_orderings); + let join_eq = join_equivalence_properties( + left_eq_properties, + right_eq_properties, + &JoinType::Inner, + Arc::new(Schema::empty()), + &[true, false], + Some(JoinSide::Left), + &[], + ); + let err_msg = + format!("expected: {:?}, actual:{:?}", expected, &join_eq.oeq_class); + assert_eq!(join_eq.oeq_class.len(), expected.len(), "{}", err_msg); + for ordering in join_eq.oeq_class { + assert!( + expected.contains(&ordering), + "{}, ordering: {:?}", + err_msg, + ordering + ); + } + } + Ok(()) + } + + #[test] + fn test_get_updated_right_ordering_equivalence_properties() -> Result<()> { + let join_type = JoinType::Inner; + // Join right child schema + let child_fields: Fields = ["x", "y", "z", "w"] + .into_iter() + .map(|name| Field::new(name, DataType::Int32, true)) + .collect(); + let child_schema = Schema::new(child_fields); + let col_x = &col("x", &child_schema)?; + let col_y = &col("y", &child_schema)?; + let col_z = &col("z", &child_schema)?; + let col_w = &col("w", &child_schema)?; + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + // [x ASC, y ASC], [z ASC, w ASC] + let orderings = vec![ + vec![(col_x, option_asc), (col_y, option_asc)], + vec![(col_z, option_asc), (col_w, option_asc)], + ]; + let orderings = convert_to_orderings(&orderings); + // Right child ordering equivalences + let mut right_oeq_class = OrderingEquivalenceClass::new(orderings); + + let left_columns_len = 4; + + let fields: Fields = ["a", "b", "c", "d", "x", "y", "z", "w"] + .into_iter() + .map(|name| Field::new(name, DataType::Int32, true)) + .collect(); + + // Join Schema + let schema = Schema::new(fields); + let col_a = &col("a", &schema)?; + let col_d = &col("d", &schema)?; + let col_x = &col("x", &schema)?; + let col_y = &col("y", &schema)?; + let col_z = &col("z", &schema)?; + let col_w = &col("w", &schema)?; + + let mut join_eq_properties = EquivalenceProperties::new(Arc::new(schema)); + // a=x and d=w + join_eq_properties.add_equal_conditions(col_a, col_x)?; + join_eq_properties.add_equal_conditions(col_d, col_w)?; + + updated_right_ordering_equivalence_class( + &mut right_oeq_class, + &join_type, + left_columns_len, + ); + join_eq_properties.add_ordering_equivalence_class(right_oeq_class); + let result = join_eq_properties.oeq_class().clone(); + + // [x ASC, y ASC], [z ASC, w ASC] + let orderings = vec![ + vec![(col_x, option_asc), (col_y, option_asc)], + vec![(col_z, option_asc), (col_w, option_asc)], + ]; + let orderings = convert_to_orderings(&orderings); + let expected = OrderingEquivalenceClass::new(orderings); + + assert_eq!(result, expected); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs new file mode 100644 index 000000000000..080587c0e231 --- /dev/null +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -0,0 +1,1639 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod dependency; // Submodule containing DependencyMap and Dependencies +mod joins; // Submodule containing join_equivalence_properties +mod union; // Submodule containing calculate_union + +use dependency::{ + construct_prefix_orderings, generate_dependency_orderings, referred_dependencies, + Dependencies, DependencyMap, +}; +pub use joins::*; +pub use union::*; + +use std::fmt::Display; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; +use std::{fmt, mem}; + +use crate::equivalence::class::{const_exprs_contains, AcrossPartitions}; +use crate::equivalence::{ + EquivalenceClass, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping, +}; +use crate::expressions::{with_new_schema, CastExpr, Column, Literal}; +use crate::{ + physical_exprs_contains, ConstExpr, LexOrdering, LexRequirement, PhysicalExpr, + PhysicalSortExpr, PhysicalSortRequirement, +}; + +use arrow::compute::SortOptions; +use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{plan_err, Constraint, Constraints, HashMap, Result}; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; +use datafusion_physical_expr_common::utils::ExprPropertiesNode; + +use indexmap::IndexSet; +use itertools::Itertools; + +/// A `EquivalenceProperties` object stores information known about the output +/// of a plan node, that can be used to optimize the plan. +/// +/// Currently, it keeps track of: +/// - Sort expressions (orderings) +/// - Equivalent expressions: expressions that are known to have same value. +/// - Constants expressions: expressions that are known to contain a single +/// constant value. +/// +/// # Example equivalent sort expressions +/// +/// Consider table below: +/// +/// ```text +/// ┌-------┐ +/// | a | b | +/// |---|---| +/// | 1 | 9 | +/// | 2 | 8 | +/// | 3 | 7 | +/// | 5 | 5 | +/// └---┴---┘ +/// ``` +/// +/// In this case, both `a ASC` and `b DESC` can describe the table ordering. +/// `EquivalenceProperties`, tracks these different valid sort expressions and +/// treat `a ASC` and `b DESC` on an equal footing. For example if the query +/// specifies the output sorted by EITHER `a ASC` or `b DESC`, the sort can be +/// avoided. +/// +/// # Example equivalent expressions +/// +/// Similarly, consider the table below: +/// +/// ```text +/// ┌-------┐ +/// | a | b | +/// |---|---| +/// | 1 | 1 | +/// | 2 | 2 | +/// | 3 | 3 | +/// | 5 | 5 | +/// └---┴---┘ +/// ``` +/// +/// In this case, columns `a` and `b` always have the same value, which can of +/// such equivalences inside this object. With this information, Datafusion can +/// optimize operations such as. For example, if the partition requirement is +/// `Hash(a)` and output partitioning is `Hash(b)`, then DataFusion avoids +/// repartitioning the data as the existing partitioning satisfies the +/// requirement. +/// +/// # Code Example +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::{Schema, Field, DataType, SchemaRef}; +/// # use datafusion_physical_expr::{ConstExpr, EquivalenceProperties}; +/// # use datafusion_physical_expr::expressions::col; +/// use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +/// # let schema: SchemaRef = Arc::new(Schema::new(vec![ +/// # Field::new("a", DataType::Int32, false), +/// # Field::new("b", DataType::Int32, false), +/// # Field::new("c", DataType::Int32, false), +/// # ])); +/// # let col_a = col("a", &schema).unwrap(); +/// # let col_b = col("b", &schema).unwrap(); +/// # let col_c = col("c", &schema).unwrap(); +/// // This object represents data that is sorted by a ASC, c DESC +/// // with a single constant value of b +/// let mut eq_properties = EquivalenceProperties::new(schema) +/// .with_constants(vec![ConstExpr::from(col_b)]); +/// eq_properties.add_new_ordering(LexOrdering::new(vec![ +/// PhysicalSortExpr::new_default(col_a).asc(), +/// PhysicalSortExpr::new_default(col_c).desc(), +/// ])); +/// +/// assert_eq!(eq_properties.to_string(), "order: [[a@0 ASC, c@2 DESC]], const: [b@1(heterogeneous)]") +/// ``` +#[derive(Debug, Clone)] +pub struct EquivalenceProperties { + /// Distinct equivalence classes (exprs known to have the same expressions) + eq_group: EquivalenceGroup, + /// Equivalent sort expressions + oeq_class: OrderingEquivalenceClass, + /// Expressions whose values are constant + /// + /// TODO: We do not need to track constants separately, they can be tracked + /// inside `eq_group` as `Literal` expressions. + constants: Vec, + /// Table constraints + constraints: Constraints, + /// Schema associated with this object. + schema: SchemaRef, +} + +impl EquivalenceProperties { + /// Creates an empty `EquivalenceProperties` object. + pub fn new(schema: SchemaRef) -> Self { + Self { + eq_group: EquivalenceGroup::empty(), + oeq_class: OrderingEquivalenceClass::empty(), + constants: vec![], + constraints: Constraints::empty(), + schema, + } + } + + /// Adds constraints to the properties. + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.constraints = constraints; + self + } + + /// Creates a new `EquivalenceProperties` object with the given orderings. + pub fn new_with_orderings(schema: SchemaRef, orderings: &[LexOrdering]) -> Self { + Self { + eq_group: EquivalenceGroup::empty(), + oeq_class: OrderingEquivalenceClass::new(orderings.to_vec()), + constants: vec![], + constraints: Constraints::empty(), + schema, + } + } + + /// Returns the associated schema. + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Returns a reference to the ordering equivalence class within. + pub fn oeq_class(&self) -> &OrderingEquivalenceClass { + &self.oeq_class + } + + /// Return the inner OrderingEquivalenceClass, consuming self + pub fn into_oeq_class(self) -> OrderingEquivalenceClass { + self.oeq_class + } + + /// Returns a reference to the equivalence group within. + pub fn eq_group(&self) -> &EquivalenceGroup { + &self.eq_group + } + + /// Returns a reference to the constant expressions + pub fn constants(&self) -> &[ConstExpr] { + &self.constants + } + + pub fn constraints(&self) -> &Constraints { + &self.constraints + } + + /// Returns the output ordering of the properties. + pub fn output_ordering(&self) -> Option { + let constants = self.constants(); + let mut output_ordering = self.oeq_class().output_ordering().unwrap_or_default(); + // Prune out constant expressions + output_ordering + .retain(|sort_expr| !const_exprs_contains(constants, &sort_expr.expr)); + (!output_ordering.is_empty()).then_some(output_ordering) + } + + /// Returns the normalized version of the ordering equivalence class within. + /// Normalization removes constants and duplicates as well as standardizing + /// expressions according to the equivalence group within. + pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass { + OrderingEquivalenceClass::new( + self.oeq_class + .iter() + .map(|ordering| self.normalize_sort_exprs(ordering)) + .collect(), + ) + } + + /// Extends this `EquivalenceProperties` with the `other` object. + pub fn extend(mut self, other: Self) -> Self { + self.eq_group.extend(other.eq_group); + self.oeq_class.extend(other.oeq_class); + self.with_constants(other.constants) + } + + /// Clears (empties) the ordering equivalence class within this object. + /// Call this method when existing orderings are invalidated. + pub fn clear_orderings(&mut self) { + self.oeq_class.clear(); + } + + /// Removes constant expressions that may change across partitions. + /// This method should be used when data from different partitions are merged. + pub fn clear_per_partition_constants(&mut self) { + self.constants.retain(|item| { + matches!(item.across_partitions(), AcrossPartitions::Uniform(_)) + }) + } + + /// Extends this `EquivalenceProperties` by adding the orderings inside the + /// ordering equivalence class `other`. + pub fn add_ordering_equivalence_class(&mut self, other: OrderingEquivalenceClass) { + self.oeq_class.extend(other); + } + + /// Adds new orderings into the existing ordering equivalence class. + pub fn add_new_orderings( + &mut self, + orderings: impl IntoIterator, + ) { + self.oeq_class.add_new_orderings(orderings); + } + + /// Adds a single ordering to the existing ordering equivalence class. + pub fn add_new_ordering(&mut self, ordering: LexOrdering) { + self.add_new_orderings([ordering]); + } + + /// Incorporates the given equivalence group to into the existing + /// equivalence group within. + pub fn add_equivalence_group(&mut self, other_eq_group: EquivalenceGroup) { + self.eq_group.extend(other_eq_group); + } + + /// Adds a new equality condition into the existing equivalence group. + /// If the given equality defines a new equivalence class, adds this new + /// equivalence class to the equivalence group. + pub fn add_equal_conditions( + &mut self, + left: &Arc, + right: &Arc, + ) -> Result<()> { + // Discover new constants in light of new the equality: + if self.is_expr_constant(left) { + // Left expression is constant, add right as constant + if !const_exprs_contains(&self.constants, right) { + let const_expr = ConstExpr::from(right) + .with_across_partitions(self.get_expr_constant_value(left)); + self.constants.push(const_expr); + } + } else if self.is_expr_constant(right) { + // Right expression is constant, add left as constant + if !const_exprs_contains(&self.constants, left) { + let const_expr = ConstExpr::from(left) + .with_across_partitions(self.get_expr_constant_value(right)); + self.constants.push(const_expr); + } + } + + // Add equal expressions to the state + self.eq_group.add_equal_conditions(left, right); + + // Discover any new orderings + self.discover_new_orderings(left)?; + Ok(()) + } + + /// Track/register physical expressions with constant values. + #[deprecated(since = "43.0.0", note = "Use [`with_constants`] instead")] + pub fn add_constants(self, constants: impl IntoIterator) -> Self { + self.with_constants(constants) + } + + /// Remove the specified constant + pub fn remove_constant(mut self, c: &ConstExpr) -> Self { + self.constants.retain(|existing| existing != c); + self + } + + /// Track/register physical expressions with constant values. + pub fn with_constants( + mut self, + constants: impl IntoIterator, + ) -> Self { + let normalized_constants = constants + .into_iter() + .filter_map(|c| { + let across_partitions = c.across_partitions(); + let expr = c.owned_expr(); + let normalized_expr = self.eq_group.normalize_expr(expr); + + if const_exprs_contains(&self.constants, &normalized_expr) { + return None; + } + + let const_expr = ConstExpr::from(normalized_expr) + .with_across_partitions(across_partitions); + + Some(const_expr) + }) + .collect::>(); + + // Add all new normalized constants + self.constants.extend(normalized_constants); + + // Discover any new orderings based on the constants + for ordering in self.normalized_oeq_class().iter() { + if let Err(e) = self.discover_new_orderings(&ordering[0].expr) { + log::debug!("error discovering new orderings: {e}"); + } + } + + self + } + + // Discover new valid orderings in light of a new equality. + // Accepts a single argument (`expr`) which is used to determine + // which orderings should be updated. + // When constants or equivalence classes are changed, there may be new orderings + // that can be discovered with the new equivalence properties. + // For a discussion, see: https://github.com/apache/datafusion/issues/9812 + fn discover_new_orderings(&mut self, expr: &Arc) -> Result<()> { + let normalized_expr = self.eq_group().normalize_expr(Arc::clone(expr)); + let eq_class = self + .eq_group + .iter() + .find_map(|class| { + class + .contains(&normalized_expr) + .then(|| class.clone().into_vec()) + }) + .unwrap_or_else(|| vec![Arc::clone(&normalized_expr)]); + + let mut new_orderings: Vec = vec![]; + for ordering in self.normalized_oeq_class().iter() { + if !ordering[0].expr.eq(&normalized_expr) { + continue; + } + + let leading_ordering_options = ordering[0].options; + + for equivalent_expr in &eq_class { + let children = equivalent_expr.children(); + if children.is_empty() { + continue; + } + + // Check if all children match the next expressions in the ordering + let mut all_children_match = true; + let mut child_properties = vec![]; + + // Build properties for each child based on the next expressions + for (i, child) in children.iter().enumerate() { + if let Some(next) = ordering.get(i + 1) { + if !child.as_ref().eq(next.expr.as_ref()) { + all_children_match = false; + break; + } + child_properties.push(ExprProperties { + sort_properties: SortProperties::Ordered(next.options), + range: Interval::make_unbounded( + &child.data_type(&self.schema)?, + )?, + preserves_lex_ordering: true, + }); + } else { + all_children_match = false; + break; + } + } + + if all_children_match { + // Check if the expression is monotonic in all arguments + if let Ok(expr_properties) = + equivalent_expr.get_properties(&child_properties) + { + if expr_properties.preserves_lex_ordering + && SortProperties::Ordered(leading_ordering_options) + == expr_properties.sort_properties + { + // Assume existing ordering is [c ASC, a ASC, b ASC] + // When equality c = f(a,b) is given, if we know that given ordering `[a ASC, b ASC]`, + // ordering `[f(a,b) ASC]` is valid, then we can deduce that ordering `[a ASC, b ASC]` is also valid. + // Hence, ordering `[a ASC, b ASC]` can be added to the state as a valid ordering. + // (e.g. existing ordering where leading ordering is removed) + new_orderings.push(LexOrdering::new(ordering[1..].to_vec())); + break; + } + } + } + } + } + + self.oeq_class.add_new_orderings(new_orderings); + Ok(()) + } + + /// Updates the ordering equivalence group within assuming that the table + /// is re-sorted according to the argument `sort_exprs`. Note that constants + /// and equivalence classes are unchanged as they are unaffected by a re-sort. + /// If the given ordering is already satisfied, the function does nothing. + pub fn with_reorder(mut self, sort_exprs: LexOrdering) -> Self { + // Filter out constant expressions as they don't affect ordering + let filtered_exprs = LexOrdering::new( + sort_exprs + .into_iter() + .filter(|expr| !self.is_expr_constant(&expr.expr)) + .collect(), + ); + + if filtered_exprs.is_empty() { + return self; + } + + let mut new_orderings = vec![filtered_exprs.clone()]; + + // Preserve valid suffixes from existing orderings + let oeq_class = mem::take(&mut self.oeq_class); + for existing in oeq_class { + if self.is_prefix_of(&filtered_exprs, &existing) { + let mut extended = filtered_exprs.clone(); + extended.extend(existing.into_iter().skip(filtered_exprs.len())); + new_orderings.push(extended); + } + } + + self.oeq_class = OrderingEquivalenceClass::new(new_orderings); + self + } + + /// Checks if the new ordering matches a prefix of the existing ordering + /// (considering expression equivalences) + fn is_prefix_of(&self, new_order: &LexOrdering, existing: &LexOrdering) -> bool { + // Check if new order is longer than existing - can't be a prefix + if new_order.len() > existing.len() { + return false; + } + + // Check if new order matches existing prefix (considering equivalences) + new_order.iter().zip(existing).all(|(new, existing)| { + self.eq_group.exprs_equal(&new.expr, &existing.expr) + && new.options == existing.options + }) + } + + /// Normalizes the given sort expressions (i.e. `sort_exprs`) using the + /// equivalence group and the ordering equivalence class within. + /// + /// Assume that `self.eq_group` states column `a` and `b` are aliases. + /// Also assume that `self.oeq_class` states orderings `d ASC` and `a ASC, c ASC` + /// are equivalent (in the sense that both describe the ordering of the table). + /// If the `sort_exprs` argument were `vec![b ASC, c ASC, a ASC]`, then this + /// function would return `vec![a ASC, c ASC]`. Internally, it would first + /// normalize to `vec![a ASC, c ASC, a ASC]` and end up with the final result + /// after deduplication. + fn normalize_sort_exprs(&self, sort_exprs: &LexOrdering) -> LexOrdering { + // Convert sort expressions to sort requirements: + let sort_reqs = LexRequirement::from(sort_exprs.clone()); + // Normalize the requirements: + let normalized_sort_reqs = self.normalize_sort_requirements(&sort_reqs); + // Convert sort requirements back to sort expressions: + LexOrdering::from(normalized_sort_reqs) + } + + /// Normalizes the given sort requirements (i.e. `sort_reqs`) using the + /// equivalence group and the ordering equivalence class within. It works by: + /// - Removing expressions that have a constant value from the given requirement. + /// - Replacing sections that belong to some equivalence class in the equivalence + /// group with the first entry in the matching equivalence class. + /// + /// Assume that `self.eq_group` states column `a` and `b` are aliases. + /// Also assume that `self.oeq_class` states orderings `d ASC` and `a ASC, c ASC` + /// are equivalent (in the sense that both describe the ordering of the table). + /// If the `sort_reqs` argument were `vec![b ASC, c ASC, a ASC]`, then this + /// function would return `vec![a ASC, c ASC]`. Internally, it would first + /// normalize to `vec![a ASC, c ASC, a ASC]` and end up with the final result + /// after deduplication. + fn normalize_sort_requirements(&self, sort_reqs: &LexRequirement) -> LexRequirement { + let normalized_sort_reqs = self.eq_group.normalize_sort_requirements(sort_reqs); + let mut constant_exprs = vec![]; + constant_exprs.extend( + self.constants + .iter() + .map(|const_expr| Arc::clone(const_expr.expr())), + ); + let constants_normalized = self.eq_group.normalize_exprs(constant_exprs); + // Prune redundant sections in the requirement: + normalized_sort_reqs + .iter() + .filter(|&order| !physical_exprs_contains(&constants_normalized, &order.expr)) + .cloned() + .collect::() + .collapse() + } + + /// Checks whether the given ordering is satisfied by any of the existing + /// orderings. + pub fn ordering_satisfy(&self, given: &LexOrdering) -> bool { + // Convert the given sort expressions to sort requirements: + let sort_requirements = LexRequirement::from(given.clone()); + self.ordering_satisfy_requirement(&sort_requirements) + } + + /// Checks whether the given sort requirements are satisfied by any of the + /// existing orderings. + pub fn ordering_satisfy_requirement(&self, reqs: &LexRequirement) -> bool { + let mut eq_properties = self.clone(); + // First, standardize the given requirement: + let normalized_reqs = eq_properties.normalize_sort_requirements(reqs); + + // Check whether given ordering is satisfied by constraints first + if self.satisfied_by_constraints(&normalized_reqs) { + return true; + } + + for normalized_req in normalized_reqs { + // Check whether given ordering is satisfied + if !eq_properties.ordering_satisfy_single(&normalized_req) { + return false; + } + // Treat satisfied keys as constants in subsequent iterations. We + // can do this because the "next" key only matters in a lexicographical + // ordering when the keys to its left have the same values. + // + // Note that these expressions are not properly "constants". This is just + // an implementation strategy confined to this function. + // + // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, + // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. + // From the analysis above, we know that `[a ASC]` is satisfied. Then, + // we add column `a` as constant to the algorithm state. This enables us + // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. + eq_properties = eq_properties + .with_constants(std::iter::once(ConstExpr::from(normalized_req.expr))); + } + true + } + + /// Checks if the sort requirements are satisfied by any of the table constraints (primary key or unique). + /// Returns true if any constraint fully satisfies the requirements. + fn satisfied_by_constraints( + &self, + normalized_reqs: &[PhysicalSortRequirement], + ) -> bool { + self.constraints.iter().any(|constraint| match constraint { + Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => self + .satisfied_by_constraint( + normalized_reqs, + indices, + matches!(constraint, Constraint::Unique(_)), + ), + }) + } + + /// Checks if sort requirements are satisfied by a constraint (primary key or unique). + /// Returns true if the constraint indices form a valid prefix of an existing ordering + /// that matches the requirements. For unique constraints, also verifies nullable columns. + fn satisfied_by_constraint( + &self, + normalized_reqs: &[PhysicalSortRequirement], + indices: &[usize], + check_null: bool, + ) -> bool { + // Requirements must contain indices + if indices.len() > normalized_reqs.len() { + return false; + } + + // Iterate over all orderings + self.oeq_class.iter().any(|ordering| { + if indices.len() > ordering.len() { + return false; + } + + // Build a map of column positions in the ordering + let mut col_positions = HashMap::with_capacity(ordering.len()); + for (pos, req) in ordering.iter().enumerate() { + if let Some(col) = req.expr.as_any().downcast_ref::() { + col_positions.insert( + col.index(), + (pos, col.nullable(&self.schema).unwrap_or(true)), + ); + } + } + + // Check if all constraint indices appear in valid positions + if !indices.iter().all(|&idx| { + col_positions + .get(&idx) + .map(|&(pos, nullable)| { + // For unique constraints, verify column is not nullable if it's first/last + !check_null + || (pos != 0 && pos != ordering.len() - 1) + || !nullable + }) + .unwrap_or(false) + }) { + return false; + } + + // Check if this ordering matches requirements prefix + let ordering_len = ordering.len(); + normalized_reqs.len() >= ordering_len + && normalized_reqs[..ordering_len].iter().zip(ordering).all( + |(req, existing)| { + req.expr.eq(&existing.expr) + && req + .options + .is_none_or(|req_opts| req_opts == existing.options) + }, + ) + }) + } + + /// Determines whether the ordering specified by the given sort requirement + /// is satisfied based on the orderings within, equivalence classes, and + /// constant expressions. + /// + /// # Parameters + /// + /// - `req`: A reference to a `PhysicalSortRequirement` for which the ordering + /// satisfaction check will be done. + /// + /// # Returns + /// + /// Returns `true` if the specified ordering is satisfied, `false` otherwise. + fn ordering_satisfy_single(&self, req: &PhysicalSortRequirement) -> bool { + let ExprProperties { + sort_properties, .. + } = self.get_expr_properties(Arc::clone(&req.expr)); + match sort_properties { + SortProperties::Ordered(options) => { + let sort_expr = PhysicalSortExpr { + expr: Arc::clone(&req.expr), + options, + }; + sort_expr.satisfy(req, self.schema()) + } + // Singleton expressions satisfies any ordering. + SortProperties::Singleton => true, + SortProperties::Unordered => false, + } + } + + /// Checks whether the `given` sort requirements are equal or more specific + /// than the `reference` sort requirements. + pub fn requirements_compatible( + &self, + given: &LexRequirement, + reference: &LexRequirement, + ) -> bool { + let normalized_given = self.normalize_sort_requirements(given); + let normalized_reference = self.normalize_sort_requirements(reference); + + (normalized_reference.len() <= normalized_given.len()) + && normalized_reference + .into_iter() + .zip(normalized_given) + .all(|(reference, given)| given.compatible(&reference)) + } + + /// Returns the finer ordering among the orderings `lhs` and `rhs`, breaking + /// any ties by choosing `lhs`. + /// + /// The finer ordering is the ordering that satisfies both of the orderings. + /// If the orderings are incomparable, returns `None`. + /// + /// For example, the finer ordering among `[a ASC]` and `[a ASC, b ASC]` is + /// the latter. + pub fn get_finer_ordering( + &self, + lhs: &LexOrdering, + rhs: &LexOrdering, + ) -> Option { + // Convert the given sort expressions to sort requirements: + let lhs = LexRequirement::from(lhs.clone()); + let rhs = LexRequirement::from(rhs.clone()); + let finer = self.get_finer_requirement(&lhs, &rhs); + // Convert the chosen sort requirements back to sort expressions: + finer.map(LexOrdering::from) + } + + /// Returns the finer ordering among the requirements `lhs` and `rhs`, + /// breaking any ties by choosing `lhs`. + /// + /// The finer requirements are the ones that satisfy both of the given + /// requirements. If the requirements are incomparable, returns `None`. + /// + /// For example, the finer requirements among `[a ASC]` and `[a ASC, b ASC]` + /// is the latter. + pub fn get_finer_requirement( + &self, + req1: &LexRequirement, + req2: &LexRequirement, + ) -> Option { + let mut lhs = self.normalize_sort_requirements(req1); + let mut rhs = self.normalize_sort_requirements(req2); + lhs.inner + .iter_mut() + .zip(rhs.inner.iter_mut()) + .all(|(lhs, rhs)| { + lhs.expr.eq(&rhs.expr) + && match (lhs.options, rhs.options) { + (Some(lhs_opt), Some(rhs_opt)) => lhs_opt == rhs_opt, + (Some(options), None) => { + rhs.options = Some(options); + true + } + (None, Some(options)) => { + lhs.options = Some(options); + true + } + (None, None) => true, + } + }) + .then_some(if lhs.len() >= rhs.len() { lhs } else { rhs }) + } + + /// we substitute the ordering according to input expression type, this is a simplified version + /// In this case, we just substitute when the expression satisfy the following condition: + /// I. just have one column and is a CAST expression + /// TODO: Add one-to-ones analysis for monotonic ScalarFunctions. + /// TODO: we could precompute all the scenario that is computable, for example: atan(x + 1000) should also be substituted if + /// x is DESC or ASC + /// After substitution, we may generate more than 1 `LexOrdering`. As an example, + /// `[a ASC, b ASC]` will turn into `[a ASC, b ASC], [CAST(a) ASC, b ASC]` when projection expressions `a, b, CAST(a)` is applied. + pub fn substitute_ordering_component( + &self, + mapping: &ProjectionMapping, + sort_expr: &LexOrdering, + ) -> Result> { + let new_orderings = sort_expr + .iter() + .map(|sort_expr| { + let referring_exprs: Vec<_> = mapping + .iter() + .map(|(source, _target)| source) + .filter(|source| expr_refers(source, &sort_expr.expr)) + .cloned() + .collect(); + let mut res = LexOrdering::new(vec![sort_expr.clone()]); + // TODO: Add one-to-ones analysis for ScalarFunctions. + for r_expr in referring_exprs { + // we check whether this expression is substitutable or not + if let Some(cast_expr) = r_expr.as_any().downcast_ref::() { + // we need to know whether the Cast Expr matches or not + let expr_type = sort_expr.expr.data_type(&self.schema)?; + if cast_expr.expr.eq(&sort_expr.expr) + && cast_expr.is_bigger_cast(expr_type) + { + res.push(PhysicalSortExpr { + expr: Arc::clone(&r_expr), + options: sort_expr.options, + }); + } + } + } + Ok(res) + }) + .collect::>>()?; + // Generate all valid orderings, given substituted expressions. + let res = new_orderings + .into_iter() + .multi_cartesian_product() + .map(LexOrdering::new) + .collect::>(); + Ok(res) + } + + /// In projection, supposed we have a input function 'A DESC B DESC' and the output shares the same expression + /// with A and B, we could surely use the ordering of the original ordering, However, if the A has been changed, + /// for example, A-> Cast(A, Int64) or any other form, it is invalid if we continue using the original ordering + /// Since it would cause bug in dependency constructions, we should substitute the input order in order to get correct + /// dependency map, happen in issue 8838: + pub fn substitute_oeq_class(&mut self, mapping: &ProjectionMapping) -> Result<()> { + let new_order = self + .oeq_class + .iter() + .map(|order| self.substitute_ordering_component(mapping, order)) + .collect::>>()?; + let new_order = new_order.into_iter().flatten().collect(); + self.oeq_class = OrderingEquivalenceClass::new(new_order); + Ok(()) + } + /// Projects argument `expr` according to `projection_mapping`, taking + /// equivalences into account. + /// + /// For example, assume that columns `a` and `c` are always equal, and that + /// `projection_mapping` encodes following mapping: + /// + /// ```text + /// a -> a1 + /// b -> b1 + /// ``` + /// + /// Then, this function projects `a + b` to `Some(a1 + b1)`, `c + b` to + /// `Some(a1 + b1)` and `d` to `None`, meaning that it cannot be projected. + pub fn project_expr( + &self, + expr: &Arc, + projection_mapping: &ProjectionMapping, + ) -> Option> { + self.eq_group.project_expr(projection_mapping, expr) + } + + /// Constructs a dependency map based on existing orderings referred to in + /// the projection. + /// + /// This function analyzes the orderings in the normalized order-equivalence + /// class and builds a dependency map. The dependency map captures relationships + /// between expressions within the orderings, helping to identify dependencies + /// and construct valid projected orderings during projection operations. + /// + /// # Parameters + /// + /// - `mapping`: A reference to the `ProjectionMapping` that defines the + /// relationship between source and target expressions. + /// + /// # Returns + /// + /// A [`DependencyMap`] representing the dependency map, where each + /// \[`DependencyNode`\] contains dependencies for the key [`PhysicalSortExpr`]. + /// + /// # Example + /// + /// Assume we have two equivalent orderings: `[a ASC, b ASC]` and `[a ASC, c ASC]`, + /// and the projection mapping is `[a -> a_new, b -> b_new, b + c -> b + c]`. + /// Then, the dependency map will be: + /// + /// ```text + /// a ASC: Node {Some(a_new ASC), HashSet{}} + /// b ASC: Node {Some(b_new ASC), HashSet{a ASC}} + /// c ASC: Node {None, HashSet{a ASC}} + /// ``` + fn construct_dependency_map(&self, mapping: &ProjectionMapping) -> DependencyMap { + let mut dependency_map = DependencyMap::new(); + for ordering in self.normalized_oeq_class().iter() { + for (idx, sort_expr) in ordering.iter().enumerate() { + let target_sort_expr = + self.project_expr(&sort_expr.expr, mapping).map(|expr| { + PhysicalSortExpr { + expr, + options: sort_expr.options, + } + }); + let is_projected = target_sort_expr.is_some(); + if is_projected + || mapping + .iter() + .any(|(source, _)| expr_refers(source, &sort_expr.expr)) + { + // Previous ordering is a dependency. Note that there is no, + // dependency for a leading ordering (i.e. the first sort + // expression). + let dependency = idx.checked_sub(1).map(|a| &ordering[a]); + // Add sort expressions that can be projected or referred to + // by any of the projection expressions to the dependency map: + dependency_map.insert( + sort_expr, + target_sort_expr.as_ref(), + dependency, + ); + } + if !is_projected { + // If we can not project, stop constructing the dependency + // map as remaining dependencies will be invalid after projection. + break; + } + } + } + dependency_map + } + + /// Returns a new `ProjectionMapping` where source expressions are normalized. + /// + /// This normalization ensures that source expressions are transformed into a + /// consistent representation. This is beneficial for algorithms that rely on + /// exact equalities, as it allows for more precise and reliable comparisons. + /// + /// # Parameters + /// + /// - `mapping`: A reference to the original `ProjectionMapping` to be normalized. + /// + /// # Returns + /// + /// A new `ProjectionMapping` with normalized source expressions. + fn normalized_mapping(&self, mapping: &ProjectionMapping) -> ProjectionMapping { + // Construct the mapping where source expressions are normalized. In this way + // In the algorithms below we can work on exact equalities + ProjectionMapping { + map: mapping + .iter() + .map(|(source, target)| { + let normalized_source = + self.eq_group.normalize_expr(Arc::clone(source)); + (normalized_source, Arc::clone(target)) + }) + .collect(), + } + } + + /// Computes projected orderings based on a given projection mapping. + /// + /// This function takes a `ProjectionMapping` and computes the possible + /// orderings for the projected expressions. It considers dependencies + /// between expressions and generates valid orderings according to the + /// specified sort properties. + /// + /// # Parameters + /// + /// - `mapping`: A reference to the `ProjectionMapping` that defines the + /// relationship between source and target expressions. + /// + /// # Returns + /// + /// A vector of `LexOrdering` containing all valid orderings after projection. + fn projected_orderings(&self, mapping: &ProjectionMapping) -> Vec { + let mapping = self.normalized_mapping(mapping); + + // Get dependency map for existing orderings: + let dependency_map = self.construct_dependency_map(&mapping); + let orderings = mapping.iter().flat_map(|(source, target)| { + referred_dependencies(&dependency_map, source) + .into_iter() + .filter_map(|relevant_deps| { + if let Ok(SortProperties::Ordered(options)) = + get_expr_properties(source, &relevant_deps, &self.schema) + .map(|prop| prop.sort_properties) + { + Some((options, relevant_deps)) + } else { + // Do not consider unordered cases + None + } + }) + .flat_map(|(options, relevant_deps)| { + let sort_expr = PhysicalSortExpr { + expr: Arc::clone(target), + options, + }; + // Generate dependent orderings (i.e. prefixes for `sort_expr`): + let mut dependency_orderings = + generate_dependency_orderings(&relevant_deps, &dependency_map); + // Append `sort_expr` to the dependent orderings: + for ordering in dependency_orderings.iter_mut() { + ordering.push(sort_expr.clone()); + } + dependency_orderings + }) + }); + + // Add valid projected orderings. For example, if existing ordering is + // `a + b` and projection is `[a -> a_new, b -> b_new]`, we need to + // preserve `a_new + b_new` as ordered. Please note that `a_new` and + // `b_new` themselves need not be ordered. Such dependencies cannot be + // deduced via the pass above. + let projected_orderings = dependency_map.iter().flat_map(|(sort_expr, node)| { + let mut prefixes = construct_prefix_orderings(sort_expr, &dependency_map); + if prefixes.is_empty() { + // If prefix is empty, there is no dependency. Insert + // empty ordering: + prefixes = vec![LexOrdering::default()]; + } + // Append current ordering on top its dependencies: + for ordering in prefixes.iter_mut() { + if let Some(target) = &node.target_sort_expr { + ordering.push(target.clone()) + } + } + prefixes + }); + + // Simplify each ordering by removing redundant sections: + orderings + .chain(projected_orderings) + .map(|lex_ordering| lex_ordering.collapse()) + .collect() + } + + /// Projects constants based on the provided `ProjectionMapping`. + /// + /// This function takes a `ProjectionMapping` and identifies/projects + /// constants based on the existing constants and the mapping. It ensures + /// that constants are appropriately propagated through the projection. + /// + /// # Parameters + /// + /// - `mapping`: A reference to a `ProjectionMapping` representing the + /// mapping of source expressions to target expressions in the projection. + /// + /// # Returns + /// + /// Returns a `Vec>` containing the projected constants. + fn projected_constants(&self, mapping: &ProjectionMapping) -> Vec { + // First, project existing constants. For example, assume that `a + b` + // is known to be constant. If the projection were `a as a_new`, `b as b_new`, + // then we would project constant `a + b` as `a_new + b_new`. + let mut projected_constants = self + .constants + .iter() + .flat_map(|const_expr| { + const_expr + .map(|expr| self.eq_group.project_expr(mapping, expr)) + .map(|projected_expr| { + projected_expr + .with_across_partitions(const_expr.across_partitions()) + }) + }) + .collect::>(); + + // Add projection expressions that are known to be constant: + for (source, target) in mapping.iter() { + if self.is_expr_constant(source) + && !const_exprs_contains(&projected_constants, target) + { + if self.is_expr_constant_across_partitions(source) { + projected_constants.push( + ConstExpr::from(target) + .with_across_partitions(self.get_expr_constant_value(source)), + ) + } else { + projected_constants.push( + ConstExpr::from(target) + .with_across_partitions(AcrossPartitions::Heterogeneous), + ) + } + } + } + projected_constants + } + + /// Projects constraints according to the given projection mapping. + /// + /// This function takes a projection mapping and extracts the column indices of the target columns. + /// It then projects the constraints to only include relationships between + /// columns that exist in the projected output. + /// + /// # Arguments + /// + /// * `mapping` - A reference to `ProjectionMapping` that defines how expressions are mapped + /// in the projection operation + /// + /// # Returns + /// + /// Returns a new `Constraints` object containing only the constraints + /// that are valid for the projected columns. + fn projected_constraints(&self, mapping: &ProjectionMapping) -> Option { + let indices = mapping + .iter() + .filter_map(|(_, target)| target.as_any().downcast_ref::()) + .map(|col| col.index()) + .collect::>(); + debug_assert_eq!(mapping.map.len(), indices.len()); + self.constraints.project(&indices) + } + + /// Projects the equivalences within according to `mapping` + /// and `output_schema`. + pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self { + let eq_group = self.eq_group.project(mapping); + let oeq_class = OrderingEquivalenceClass::new(self.projected_orderings(mapping)); + let constants = self.projected_constants(mapping); + let constraints = self + .projected_constraints(mapping) + .unwrap_or_else(Constraints::empty); + Self { + schema: output_schema, + eq_group, + oeq_class, + constants, + constraints, + } + } + + /// Returns the longest (potentially partial) permutation satisfying the + /// existing ordering. For example, if we have the equivalent orderings + /// `[a ASC, b ASC]` and `[c DESC]`, with `exprs` containing `[c, b, a, d]`, + /// then this function returns `([a ASC, b ASC, c DESC], [2, 1, 0])`. + /// This means that the specification `[a ASC, b ASC, c DESC]` is satisfied + /// by the existing ordering, and `[a, b, c]` resides at indices: `2, 1, 0` + /// inside the argument `exprs` (respectively). For the mathematical + /// definition of "partial permutation", see: + /// + /// + pub fn find_longest_permutation( + &self, + exprs: &[Arc], + ) -> (LexOrdering, Vec) { + let mut eq_properties = self.clone(); + let mut result = vec![]; + // The algorithm is as follows: + // - Iterate over all the expressions and insert ordered expressions + // into the result. + // - Treat inserted expressions as constants (i.e. add them as constants + // to the state). + // - Continue the above procedure until no expression is inserted; i.e. + // the algorithm reaches a fixed point. + // This algorithm should reach a fixed point in at most `exprs.len()` + // iterations. + let mut search_indices = (0..exprs.len()).collect::>(); + for _idx in 0..exprs.len() { + // Get ordered expressions with their indices. + let ordered_exprs = search_indices + .iter() + .flat_map(|&idx| { + let ExprProperties { + sort_properties, .. + } = eq_properties.get_expr_properties(Arc::clone(&exprs[idx])); + match sort_properties { + SortProperties::Ordered(options) => Some(( + PhysicalSortExpr { + expr: Arc::clone(&exprs[idx]), + options, + }, + idx, + )), + SortProperties::Singleton => { + // Assign default ordering to constant expressions + let options = SortOptions::default(); + Some(( + PhysicalSortExpr { + expr: Arc::clone(&exprs[idx]), + options, + }, + idx, + )) + } + SortProperties::Unordered => None, + } + }) + .collect::>(); + // We reached a fixed point, exit. + if ordered_exprs.is_empty() { + break; + } + // Remove indices that have an ordering from `search_indices`, and + // treat ordered expressions as constants in subsequent iterations. + // We can do this because the "next" key only matters in a lexicographical + // ordering when the keys to its left have the same values. + // + // Note that these expressions are not properly "constants". This is just + // an implementation strategy confined to this function. + for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs { + eq_properties = + eq_properties.with_constants(std::iter::once(ConstExpr::from(expr))); + search_indices.shift_remove(idx); + } + // Add new ordered section to the state. + result.extend(ordered_exprs); + } + let (left, right) = result.into_iter().unzip(); + (LexOrdering::new(left), right) + } + + /// This function determines whether the provided expression is constant + /// based on the known constants. + /// + /// # Parameters + /// + /// - `expr`: A reference to a `Arc` representing the + /// expression to be checked. + /// + /// # Returns + /// + /// Returns `true` if the expression is constant according to equivalence + /// group, `false` otherwise. + pub fn is_expr_constant(&self, expr: &Arc) -> bool { + // As an example, assume that we know columns `a` and `b` are constant. + // Then, `a`, `b` and `a + b` will all return `true` whereas `c` will + // return `false`. + let const_exprs = self + .constants + .iter() + .map(|const_expr| Arc::clone(const_expr.expr())); + let normalized_constants = self.eq_group.normalize_exprs(const_exprs); + let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); + is_constant_recurse(&normalized_constants, &normalized_expr) + } + + /// This function determines whether the provided expression is constant + /// across partitions based on the known constants. + /// + /// # Parameters + /// + /// - `expr`: A reference to a `Arc` representing the + /// expression to be checked. + /// + /// # Returns + /// + /// Returns `true` if the expression is constant across all partitions according + /// to equivalence group, `false` otherwise + #[deprecated( + since = "45.0.0", + note = "Use [`is_expr_constant_across_partitions`] instead" + )] + pub fn is_expr_constant_accross_partitions( + &self, + expr: &Arc, + ) -> bool { + self.is_expr_constant_across_partitions(expr) + } + + /// This function determines whether the provided expression is constant + /// across partitions based on the known constants. + /// + /// # Parameters + /// + /// - `expr`: A reference to a `Arc` representing the + /// expression to be checked. + /// + /// # Returns + /// + /// Returns `true` if the expression is constant across all partitions according + /// to equivalence group, `false` otherwise. + pub fn is_expr_constant_across_partitions( + &self, + expr: &Arc, + ) -> bool { + // As an example, assume that we know columns `a` and `b` are constant. + // Then, `a`, `b` and `a + b` will all return `true` whereas `c` will + // return `false`. + let const_exprs = self + .constants + .iter() + .filter_map(|const_expr| { + if matches!( + const_expr.across_partitions(), + AcrossPartitions::Uniform { .. } + ) { + Some(Arc::clone(const_expr.expr())) + } else { + None + } + }) + .collect::>(); + let normalized_constants = self.eq_group.normalize_exprs(const_exprs); + let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); + is_constant_recurse(&normalized_constants, &normalized_expr) + } + + /// Retrieves the constant value of a given physical expression, if it exists. + /// + /// Normalizes the input expression and checks if it matches any known constants + /// in the current context. Returns whether the expression has a uniform value, + /// varies across partitions, or is not constant. + /// + /// # Parameters + /// - `expr`: A reference to the physical expression to evaluate. + /// + /// # Returns + /// - `AcrossPartitions::Uniform(value)`: If the expression has the same value across partitions. + /// - `AcrossPartitions::Heterogeneous`: If the expression varies across partitions. + /// - `None`: If the expression is not recognized as constant. + pub fn get_expr_constant_value( + &self, + expr: &Arc, + ) -> AcrossPartitions { + let normalized_expr = self.eq_group.normalize_expr(Arc::clone(expr)); + + if let Some(lit) = normalized_expr.as_any().downcast_ref::() { + return AcrossPartitions::Uniform(Some(lit.value().clone())); + } + + for const_expr in self.constants.iter() { + if normalized_expr.eq(const_expr.expr()) { + return const_expr.across_partitions(); + } + } + + AcrossPartitions::Heterogeneous + } + + /// Retrieves the properties for a given physical expression. + /// + /// This function constructs an [`ExprProperties`] object for the given + /// expression, which encapsulates information about the expression's + /// properties, including its [`SortProperties`] and [`Interval`]. + /// + /// # Parameters + /// + /// - `expr`: An `Arc` representing the physical expression + /// for which ordering information is sought. + /// + /// # Returns + /// + /// Returns an [`ExprProperties`] object containing the ordering and range + /// information for the given expression. + pub fn get_expr_properties(&self, expr: Arc) -> ExprProperties { + ExprPropertiesNode::new_unknown(expr) + .transform_up(|expr| update_properties(expr, self)) + .data() + .map(|node| node.data) + .unwrap_or(ExprProperties::new_unknown()) + } + + /// Transforms this `EquivalenceProperties` into a new `EquivalenceProperties` + /// by mapping columns in the original schema to columns in the new schema + /// by index. + pub fn with_new_schema(self, schema: SchemaRef) -> Result { + // The new schema and the original schema is aligned when they have the + // same number of columns, and fields at the same index have the same + // type in both schemas. + let schemas_aligned = (self.schema.fields.len() == schema.fields.len()) + && self + .schema + .fields + .iter() + .zip(schema.fields.iter()) + .all(|(lhs, rhs)| lhs.data_type().eq(rhs.data_type())); + if !schemas_aligned { + // Rewriting equivalence properties in terms of new schema is not + // safe when schemas are not aligned: + return plan_err!( + "Cannot rewrite old_schema:{:?} with new schema: {:?}", + self.schema, + schema + ); + } + // Rewrite constants according to new schema: + let new_constants = self + .constants + .into_iter() + .map(|const_expr| { + let across_partitions = const_expr.across_partitions(); + let new_const_expr = with_new_schema(const_expr.owned_expr(), &schema)?; + Ok(ConstExpr::new(new_const_expr) + .with_across_partitions(across_partitions)) + }) + .collect::>>()?; + + // Rewrite orderings according to new schema: + let mut new_orderings = vec![]; + for ordering in self.oeq_class { + let new_ordering = ordering + .into_iter() + .map(|mut sort_expr| { + sort_expr.expr = with_new_schema(sort_expr.expr, &schema)?; + Ok(sort_expr) + }) + .collect::>()?; + new_orderings.push(new_ordering); + } + + // Rewrite equivalence classes according to the new schema: + let mut eq_classes = vec![]; + for eq_class in self.eq_group { + let new_eq_exprs = eq_class + .into_vec() + .into_iter() + .map(|expr| with_new_schema(expr, &schema)) + .collect::>()?; + eq_classes.push(EquivalenceClass::new(new_eq_exprs)); + } + + // Construct the resulting equivalence properties: + let mut result = EquivalenceProperties::new(schema); + result.constants = new_constants; + result.add_new_orderings(new_orderings); + result.add_equivalence_group(EquivalenceGroup::new(eq_classes)); + + Ok(result) + } +} + +/// More readable display version of the `EquivalenceProperties`. +/// +/// Format: +/// ```text +/// order: [[a ASC, b ASC], [a ASC, c ASC]], eq: [[a = b], [a = c]], const: [a = 1] +/// ``` +impl Display for EquivalenceProperties { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.eq_group.is_empty() + && self.oeq_class.is_empty() + && self.constants.is_empty() + { + return write!(f, "No properties"); + } + if !self.oeq_class.is_empty() { + write!(f, "order: {}", self.oeq_class)?; + } + if !self.eq_group.is_empty() { + write!(f, ", eq: {}", self.eq_group)?; + } + if !self.constants.is_empty() { + write!(f, ", const: [{}]", ConstExpr::format_list(&self.constants))?; + } + Ok(()) + } +} + +/// Calculates the properties of a given [`ExprPropertiesNode`]. +/// +/// Order information can be retrieved as: +/// - If it is a leaf node, we directly find the order of the node by looking +/// at the given sort expression and equivalence properties if it is a `Column` +/// leaf, or we mark it as unordered. In the case of a `Literal` leaf, we mark +/// it as singleton so that it can cooperate with all ordered columns. +/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` +/// and operator has its own rules on how to propagate the children orderings. +/// However, before we engage in recursion, we check whether this intermediate +/// node directly matches with the sort expression. If there is a match, the +/// sort expression emerges at that node immediately, discarding the recursive +/// result coming from its children. +/// +/// Range information is calculated as: +/// - If it is a `Literal` node, we set the range as a point value. If it is a +/// `Column` node, we set the datatype of the range, but cannot give an interval +/// for the range, yet. +/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr` +/// and operator has its own rules on how to propagate the children range. +fn update_properties( + mut node: ExprPropertiesNode, + eq_properties: &EquivalenceProperties, +) -> Result> { + // First, try to gather the information from the children: + if !node.expr.children().is_empty() { + // We have an intermediate (non-leaf) node, account for its children: + let children_props = node.children.iter().map(|c| c.data.clone()).collect_vec(); + node.data = node.expr.get_properties(&children_props)?; + } else if node.expr.as_any().is::() { + // We have a Literal, which is one of the two possible leaf node types: + node.data = node.expr.get_properties(&[])?; + } else if node.expr.as_any().is::() { + // We have a Column, which is the other possible leaf node type: + node.data.range = + Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)? + } + // Now, check what we know about orderings: + let normalized_expr = eq_properties + .eq_group + .normalize_expr(Arc::clone(&node.expr)); + let oeq_class = eq_properties.normalized_oeq_class(); + if eq_properties.is_expr_constant(&normalized_expr) + || oeq_class.is_expr_partial_const(&normalized_expr) + { + node.data.sort_properties = SortProperties::Singleton; + } else if let Some(options) = oeq_class.get_options(&normalized_expr) { + node.data.sort_properties = SortProperties::Ordered(options); + } + Ok(Transformed::yes(node)) +} + +/// This function determines whether the provided expression is constant +/// based on the known constants. +/// +/// # Parameters +/// +/// - `constants`: A `&[Arc]` containing expressions known to +/// be a constant. +/// - `expr`: A reference to a `Arc` representing the expression +/// to check. +/// +/// # Returns +/// +/// Returns `true` if the expression is constant according to equivalence +/// group, `false` otherwise. +fn is_constant_recurse( + constants: &[Arc], + expr: &Arc, +) -> bool { + if physical_exprs_contains(constants, expr) || expr.as_any().is::() { + return true; + } + let children = expr.children(); + !children.is_empty() && children.iter().all(|c| is_constant_recurse(constants, c)) +} + +/// This function examines whether a referring expression directly refers to a +/// given referred expression or if any of its children in the expression tree +/// refer to the specified expression. +/// +/// # Parameters +/// +/// - `referring_expr`: A reference to the referring expression (`Arc`). +/// - `referred_expr`: A reference to the referred expression (`Arc`) +/// +/// # Returns +/// +/// A boolean value indicating whether `referring_expr` refers (needs it to evaluate its result) +/// `referred_expr` or not. +fn expr_refers( + referring_expr: &Arc, + referred_expr: &Arc, +) -> bool { + referring_expr.eq(referred_expr) + || referring_expr + .children() + .iter() + .any(|child| expr_refers(child, referred_expr)) +} + +/// This function examines the given expression and its properties to determine +/// the ordering properties of the expression. The range knowledge is not utilized +/// yet in the scope of this function. +/// +/// # Parameters +/// +/// - `expr`: A reference to the source expression (`Arc`) for +/// which ordering properties need to be determined. +/// - `dependencies`: A reference to `Dependencies`, containing sort expressions +/// referred to by `expr`. +/// - `schema``: A reference to the schema which the `expr` columns refer. +/// +/// # Returns +/// +/// A `SortProperties` indicating the ordering information of the given expression. +fn get_expr_properties( + expr: &Arc, + dependencies: &Dependencies, + schema: &SchemaRef, +) -> Result { + if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) { + // If exact match is found, return its ordering. + Ok(ExprProperties { + sort_properties: SortProperties::Ordered(column_order.options), + range: Interval::make_unbounded(&expr.data_type(schema)?)?, + preserves_lex_ordering: false, + }) + } else if expr.as_any().downcast_ref::().is_some() { + Ok(ExprProperties { + sort_properties: SortProperties::Unordered, + range: Interval::make_unbounded(&expr.data_type(schema)?)?, + preserves_lex_ordering: false, + }) + } else if let Some(literal) = expr.as_any().downcast_ref::() { + Ok(ExprProperties { + sort_properties: SortProperties::Singleton, + range: Interval::try_new(literal.value().clone(), literal.value().clone())?, + preserves_lex_ordering: true, + }) + } else { + // Find orderings of its children + let child_states = expr + .children() + .iter() + .map(|child| get_expr_properties(child, dependencies, schema)) + .collect::>>()?; + // Calculate expression ordering using ordering of its children. + expr.get_properties(&child_states) + } +} + +/// Wrapper struct for `Arc` to use them as keys in a hash map. +#[derive(Debug, Clone)] +struct ExprWrapper(Arc); + +impl PartialEq for ExprWrapper { + fn eq(&self, other: &Self) -> bool { + self.0.eq(&other.0) + } +} + +impl Eq for ExprWrapper {} + +impl Hash for ExprWrapper { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::expressions::{col, BinaryExpr}; + + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use datafusion_expr::Operator; + + #[test] + fn test_expr_consists_of_constants() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + ])); + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let col_d = col("d", &schema)?; + let b_plus_d = Arc::new(BinaryExpr::new( + Arc::clone(&col_b), + Operator::Plus, + Arc::clone(&col_d), + )) as Arc; + + let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b)]; + let expr = Arc::clone(&b_plus_d); + assert!(!is_constant_recurse(&constants, &expr)); + + let constants = vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_d)]; + let expr = Arc::clone(&b_plus_d); + assert!(is_constant_recurse(&constants, &expr)); + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/equivalence/properties/union.rs b/datafusion/physical-expr/src/equivalence/properties/union.rs new file mode 100644 index 000000000000..64ef9278e248 --- /dev/null +++ b/datafusion/physical-expr/src/equivalence/properties/union.rs @@ -0,0 +1,927 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion_common::{internal_err, Result}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use std::iter::Peekable; +use std::sync::Arc; + +use crate::equivalence::class::AcrossPartitions; +use crate::ConstExpr; + +use super::EquivalenceProperties; +use crate::PhysicalSortExpr; +use arrow::datatypes::SchemaRef; +use std::slice::Iter; + +/// Calculates the union (in the sense of `UnionExec`) `EquivalenceProperties` +/// of `lhs` and `rhs` according to the schema of `lhs`. +/// +/// Rules: The UnionExec does not interleave its inputs: instead it passes each +/// input partition from the children as its own output. +/// +/// Since the output equivalence properties are properties that are true for +/// *all* output partitions, that is the same as being true for all *input* +/// partitions +fn calculate_union_binary( + lhs: EquivalenceProperties, + mut rhs: EquivalenceProperties, +) -> Result { + // Harmonize the schema of the rhs with the schema of the lhs (which is the accumulator schema): + if !rhs.schema.eq(&lhs.schema) { + rhs = rhs.with_new_schema(Arc::clone(&lhs.schema))?; + } + + // First, calculate valid constants for the union. An expression is constant + // at the output of the union if it is constant in both sides with matching values. + let constants = lhs + .constants() + .iter() + .filter_map(|lhs_const| { + // Find matching constant expression in RHS + rhs.constants() + .iter() + .find(|rhs_const| rhs_const.expr().eq(lhs_const.expr())) + .map(|rhs_const| { + let mut const_expr = ConstExpr::new(Arc::clone(lhs_const.expr())); + + // If both sides have matching constant values, preserve the value and set across_partitions=true + if let ( + AcrossPartitions::Uniform(Some(lhs_val)), + AcrossPartitions::Uniform(Some(rhs_val)), + ) = (lhs_const.across_partitions(), rhs_const.across_partitions()) + { + if lhs_val == rhs_val { + const_expr = const_expr.with_across_partitions( + AcrossPartitions::Uniform(Some(lhs_val)), + ) + } + } + const_expr + }) + }) + .collect::>(); + + // Next, calculate valid orderings for the union by searching for prefixes + // in both sides. + let mut orderings = UnionEquivalentOrderingBuilder::new(); + orderings.add_satisfied_orderings(lhs.normalized_oeq_class(), lhs.constants(), &rhs); + orderings.add_satisfied_orderings(rhs.normalized_oeq_class(), rhs.constants(), &lhs); + let orderings = orderings.build(); + + let mut eq_properties = + EquivalenceProperties::new(lhs.schema).with_constants(constants); + + eq_properties.add_new_orderings(orderings); + Ok(eq_properties) +} + +/// Calculates the union (in the sense of `UnionExec`) `EquivalenceProperties` +/// of the given `EquivalenceProperties` in `eqps` according to the given +/// output `schema` (which need not be the same with those of `lhs` and `rhs` +/// as details such as nullability may be different). +pub fn calculate_union( + eqps: Vec, + schema: SchemaRef, +) -> Result { + // TODO: In some cases, we should be able to preserve some equivalence + // classes. Add support for such cases. + let mut iter = eqps.into_iter(); + let Some(mut acc) = iter.next() else { + return internal_err!( + "Cannot calculate EquivalenceProperties for a union with no inputs" + ); + }; + + // Harmonize the schema of the init with the schema of the union: + if !acc.schema.eq(&schema) { + acc = acc.with_new_schema(schema)?; + } + // Fold in the rest of the EquivalenceProperties: + for props in iter { + acc = calculate_union_binary(acc, props)?; + } + Ok(acc) +} + +#[derive(Debug)] +enum AddedOrdering { + /// The ordering was added to the in progress result + Yes, + /// The ordering was not added + No(LexOrdering), +} + +/// Builds valid output orderings of a `UnionExec` +#[derive(Debug)] +struct UnionEquivalentOrderingBuilder { + orderings: Vec, +} + +impl UnionEquivalentOrderingBuilder { + fn new() -> Self { + Self { orderings: vec![] } + } + + /// Add all orderings from `orderings` that satisfy `properties`, + /// potentially augmented with`constants`. + /// + /// Note: any column that is known to be constant can be inserted into the + /// ordering without changing its meaning + /// + /// For example: + /// * `orderings` contains `[a ASC, c ASC]` and `constants` contains `b` + /// * `properties` has required ordering `[a ASC, b ASC]` + /// + /// Then this will add `[a ASC, b ASC]` to the `orderings` list (as `a` was + /// in the sort order and `b` was a constant). + fn add_satisfied_orderings( + &mut self, + orderings: impl IntoIterator, + constants: &[ConstExpr], + properties: &EquivalenceProperties, + ) { + for mut ordering in orderings.into_iter() { + // Progressively shorten the ordering to search for a satisfied prefix: + loop { + match self.try_add_ordering(ordering, constants, properties) { + AddedOrdering::Yes => break, + AddedOrdering::No(o) => { + ordering = o; + ordering.pop(); + } + } + } + } + } + + /// Adds `ordering`, potentially augmented with constants, if it satisfies + /// the target `properties` properties. + /// + /// Returns + /// + /// * [`AddedOrdering::Yes`] if the ordering was added (either directly or + /// augmented), or was empty. + /// + /// * [`AddedOrdering::No`] if the ordering was not added + fn try_add_ordering( + &mut self, + ordering: LexOrdering, + constants: &[ConstExpr], + properties: &EquivalenceProperties, + ) -> AddedOrdering { + if ordering.is_empty() { + AddedOrdering::Yes + } else if properties.ordering_satisfy(ordering.as_ref()) { + // If the ordering satisfies the target properties, no need to + // augment it with constants. + self.orderings.push(ordering); + AddedOrdering::Yes + } else { + // Did not satisfy target properties, try and augment with constants + // to match the properties + if self.try_find_augmented_ordering(&ordering, constants, properties) { + AddedOrdering::Yes + } else { + AddedOrdering::No(ordering) + } + } + } + + /// Attempts to add `constants` to `ordering` to satisfy the properties. + /// + /// returns true if any orderings were added, false otherwise + fn try_find_augmented_ordering( + &mut self, + ordering: &LexOrdering, + constants: &[ConstExpr], + properties: &EquivalenceProperties, + ) -> bool { + // can't augment if there is nothing to augment with + if constants.is_empty() { + return false; + } + let start_num_orderings = self.orderings.len(); + + // for each equivalent ordering in properties, try and augment + // `ordering` it with the constants to match + for existing_ordering in properties.oeq_class.iter() { + if let Some(augmented_ordering) = self.augment_ordering( + ordering, + constants, + existing_ordering, + &properties.constants, + ) { + if !augmented_ordering.is_empty() { + assert!(properties.ordering_satisfy(augmented_ordering.as_ref())); + self.orderings.push(augmented_ordering); + } + } + } + + self.orderings.len() > start_num_orderings + } + + /// Attempts to augment the ordering with constants to match the + /// `existing_ordering` + /// + /// Returns Some(ordering) if an augmented ordering was found, None otherwise + fn augment_ordering( + &mut self, + ordering: &LexOrdering, + constants: &[ConstExpr], + existing_ordering: &LexOrdering, + existing_constants: &[ConstExpr], + ) -> Option { + let mut augmented_ordering = LexOrdering::default(); + let mut sort_expr_iter = ordering.iter().peekable(); + let mut existing_sort_expr_iter = existing_ordering.iter().peekable(); + + // walk in parallel down the two orderings, trying to match them up + while sort_expr_iter.peek().is_some() || existing_sort_expr_iter.peek().is_some() + { + // If the next expressions are equal, add the next match + // otherwise try and match with a constant + if let Some(expr) = + advance_if_match(&mut sort_expr_iter, &mut existing_sort_expr_iter) + { + augmented_ordering.push(expr); + } else if let Some(expr) = + advance_if_matches_constant(&mut sort_expr_iter, existing_constants) + { + augmented_ordering.push(expr); + } else if let Some(expr) = + advance_if_matches_constant(&mut existing_sort_expr_iter, constants) + { + augmented_ordering.push(expr); + } else { + // no match, can't continue the ordering, return what we have + break; + } + } + + Some(augmented_ordering) + } + + fn build(self) -> Vec { + self.orderings + } +} + +/// Advances two iterators in parallel +/// +/// If the next expressions are equal, the iterators are advanced and returns +/// the matched expression . +/// +/// Otherwise, the iterators are left unchanged and return `None` +fn advance_if_match( + iter1: &mut Peekable>, + iter2: &mut Peekable>, +) -> Option { + if matches!((iter1.peek(), iter2.peek()), (Some(expr1), Some(expr2)) if expr1.eq(expr2)) + { + iter1.next().unwrap(); + iter2.next().cloned() + } else { + None + } +} + +/// Advances the iterator with a constant +/// +/// If the next expression matches one of the constants, advances the iterator +/// returning the matched expression +/// +/// Otherwise, the iterator is left unchanged and returns `None` +fn advance_if_matches_constant( + iter: &mut Peekable>, + constants: &[ConstExpr], +) -> Option { + let expr = iter.peek()?; + let const_expr = constants.iter().find(|c| c.eq_expr(expr))?; + let found_expr = PhysicalSortExpr::new(Arc::clone(const_expr.expr()), expr.options); + iter.next(); + Some(found_expr) +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::equivalence::class::const_exprs_contains; + use crate::equivalence::tests::{create_test_schema, parse_sort_expr}; + use crate::expressions::col; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; + + use itertools::Itertools; + + #[test] + fn test_union_equivalence_properties_multi_children_1() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + let schema3 = append_fields(&schema, "2"); + UnionEquivalenceTest::new(&schema) + // Children 1 + .with_child_sort(vec![vec!["a", "b", "c"]], &schema) + // Children 2 + .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) + // Children 3 + .with_child_sort(vec![vec!["a2", "b2"]], &schema3) + .with_expected_sort(vec![vec!["a", "b"]]) + .run() + } + + #[test] + fn test_union_equivalence_properties_multi_children_2() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + let schema3 = append_fields(&schema, "2"); + UnionEquivalenceTest::new(&schema) + // Children 1 + .with_child_sort(vec![vec!["a", "b", "c"]], &schema) + // Children 2 + .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) + // Children 3 + .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3) + .with_expected_sort(vec![vec!["a", "b", "c"]]) + .run() + } + + #[test] + fn test_union_equivalence_properties_multi_children_3() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + let schema3 = append_fields(&schema, "2"); + UnionEquivalenceTest::new(&schema) + // Children 1 + .with_child_sort(vec![vec!["a", "b"]], &schema) + // Children 2 + .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2) + // Children 3 + .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3) + .with_expected_sort(vec![vec!["a", "b"]]) + .run() + } + + #[test] + fn test_union_equivalence_properties_multi_children_4() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + let schema3 = append_fields(&schema, "2"); + UnionEquivalenceTest::new(&schema) + // Children 1 + .with_child_sort(vec![vec!["a", "b"]], &schema) + // Children 2 + .with_child_sort(vec![vec!["a1", "b1"]], &schema2) + // Children 3 + .with_child_sort(vec![vec!["b2", "c2"]], &schema3) + .with_expected_sort(vec![]) + .run() + } + + #[test] + fn test_union_equivalence_properties_multi_children_5() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + UnionEquivalenceTest::new(&schema) + // Children 1 + .with_child_sort(vec![vec!["a", "b"], vec!["c"]], &schema) + // Children 2 + .with_child_sort(vec![vec!["a1", "b1"], vec!["c1"]], &schema2) + .with_expected_sort(vec![vec!["a", "b"], vec!["c"]]) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_common_constants() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child: [a ASC], const [b, c] + vec![vec!["a"]], + vec!["b", "c"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child: [b ASC], const [a, c] + vec![vec!["b"]], + vec!["a", "c"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union expected orderings: [[a ASC], [b ASC]], const [c] + vec![vec!["a"], vec!["b"]], + vec!["c"], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_prefix() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child: [a ASC], const [] + vec![vec!["a"]], + vec![], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child: [a ASC, b ASC], const [] + vec![vec!["a", "b"]], + vec![], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [a ASC], const [] + vec![vec!["a"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_asc_desc_mismatch() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child: [a ASC], const [] + vec![vec!["a"]], + vec![], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [a DESC], const [] + vec![vec!["a DESC"]], + vec![], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union doesn't have any ordering or constant + vec![], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_different_schemas() { + let schema = create_test_schema().unwrap(); + let schema2 = append_fields(&schema, "1"); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child orderings: [a ASC], const [] + vec![vec!["a"]], + vec![], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [a1 ASC, b1 ASC], const [] + vec![vec!["a1", "b1"]], + vec![], + &schema2, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [a ASC] + // + // Note that a, and a1 are at the same index for their + // corresponding schemas. + vec![vec!["a"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_fill_gaps() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child orderings: [a ASC, c ASC], const [b] + vec![vec!["a", "c"]], + vec!["b"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [b ASC, c ASC], const [a] + vec![vec!["b", "c"]], + vec!["a"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [ + // [a ASC, b ASC, c ASC], + // [b ASC, a ASC, c ASC] + // ], const [] + vec![vec!["a", "b", "c"], vec!["b", "a", "c"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_no_fill_gaps() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child orderings: [a ASC, c ASC], const [d] // some other constant + vec![vec!["a", "c"]], + vec!["d"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [b ASC, c ASC], const [a] + vec![vec!["b", "c"]], + vec!["a"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [[a]] (only a is constant) + vec![vec!["a"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_fill_some_gaps() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child orderings: [c ASC], const [a, b] // some other constant + vec![vec!["c"]], + vec!["a", "b"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [a DESC, b], const [] + vec![vec!["a DESC", "b"]], + vec![], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [[a, b]] (can fill in the a/b with constants) + vec![vec!["a DESC", "b"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_fill_gaps_non_symmetric() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child orderings: [a ASC, c ASC], const [b] + vec![vec!["a", "c"]], + vec!["b"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child orderings: [b ASC, c ASC], const [a] + vec![vec!["b DESC", "c"]], + vec!["a"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: [ + // [a ASC, b ASC, c ASC], + // [b ASC, a ASC, c ASC] + // ], const [] + vec![vec!["a", "b DESC", "c"], vec!["b DESC", "a", "c"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_gap_fill_symmetric() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child: [a ASC, b ASC, d ASC], const [c] + vec![vec!["a", "b", "d"]], + vec!["c"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child: [a ASC, c ASC, d ASC], const [b] + vec![vec!["a", "c", "d"]], + vec!["b"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: + // [a, b, c, d] + // [a, c, b, d] + vec![vec!["a", "c", "b", "d"], vec!["a", "b", "c", "d"]], + vec![], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_gap_fill_and_common() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // First child: [a DESC, d ASC], const [b, c] + vec![vec!["a DESC", "d"]], + vec!["b", "c"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child: [a DESC, c ASC, d ASC], const [b] + vec![vec!["a DESC", "c", "d"]], + vec!["b"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: + // [a DESC, c, d] [b] + vec![vec!["a DESC", "c", "d"]], + vec!["b"], + ) + .run() + } + + #[test] + fn test_union_equivalence_properties_constants_middle_desc() { + let schema = create_test_schema().unwrap(); + UnionEquivalenceTest::new(&schema) + .with_child_sort_and_const_exprs( + // NB `b DESC` in the first child + // + // First child: [a ASC, b DESC, d ASC], const [c] + vec![vec!["a", "b DESC", "d"]], + vec!["c"], + &schema, + ) + .with_child_sort_and_const_exprs( + // Second child: [a ASC, c ASC, d ASC], const [b] + vec![vec!["a", "c", "d"]], + vec!["b"], + &schema, + ) + .with_expected_sort_and_const_exprs( + // Union orderings: + // [a, b, d] (c constant) + // [a, c, d] (b constant) + vec![vec!["a", "c", "b DESC", "d"], vec!["a", "b DESC", "c", "d"]], + vec![], + ) + .run() + } + + // TODO tests with multiple constants + + #[derive(Debug)] + struct UnionEquivalenceTest { + /// The schema of the output of the Union + output_schema: SchemaRef, + /// The equivalence properties of each child to the union + child_properties: Vec, + /// The expected output properties of the union. Must be set before + /// running `build` + expected_properties: Option, + } + + impl UnionEquivalenceTest { + fn new(output_schema: &SchemaRef) -> Self { + Self { + output_schema: Arc::clone(output_schema), + child_properties: vec![], + expected_properties: None, + } + } + + /// Add a union input with the specified orderings + /// + /// See [`Self::make_props`] for the format of the strings in `orderings` + fn with_child_sort( + mut self, + orderings: Vec>, + schema: &SchemaRef, + ) -> Self { + let properties = self.make_props(orderings, vec![], schema); + self.child_properties.push(properties); + self + } + + /// Add a union input with the specified orderings and constant + /// equivalences + /// + /// See [`Self::make_props`] for the format of the strings in + /// `orderings` and `constants` + fn with_child_sort_and_const_exprs( + mut self, + orderings: Vec>, + constants: Vec<&str>, + schema: &SchemaRef, + ) -> Self { + let properties = self.make_props(orderings, constants, schema); + self.child_properties.push(properties); + self + } + + /// Set the expected output sort order for the union of the children + /// + /// See [`Self::make_props`] for the format of the strings in `orderings` + fn with_expected_sort(mut self, orderings: Vec>) -> Self { + let properties = self.make_props(orderings, vec![], &self.output_schema); + self.expected_properties = Some(properties); + self + } + + /// Set the expected output sort order and constant expressions for the + /// union of the children + /// + /// See [`Self::make_props`] for the format of the strings in + /// `orderings` and `constants`. + fn with_expected_sort_and_const_exprs( + mut self, + orderings: Vec>, + constants: Vec<&str>, + ) -> Self { + let properties = self.make_props(orderings, constants, &self.output_schema); + self.expected_properties = Some(properties); + self + } + + /// compute the union's output equivalence properties from the child + /// properties, and compare them to the expected properties + fn run(self) { + let Self { + output_schema, + child_properties, + expected_properties, + } = self; + + let expected_properties = + expected_properties.expect("expected_properties not set"); + + // try all permutations of the children + // as the code treats lhs and rhs differently + for child_properties in child_properties + .iter() + .cloned() + .permutations(child_properties.len()) + { + println!("--- permutation ---"); + for c in &child_properties { + println!("{c}"); + } + let actual_properties = + calculate_union(child_properties, Arc::clone(&output_schema)) + .expect("failed to calculate union equivalence properties"); + Self::assert_eq_properties_same( + &actual_properties, + &expected_properties, + format!( + "expected: {expected_properties:?}\nactual: {actual_properties:?}" + ), + ); + } + } + + fn assert_eq_properties_same( + lhs: &EquivalenceProperties, + rhs: &EquivalenceProperties, + err_msg: String, + ) { + // Check whether constants are same + let lhs_constants = lhs.constants(); + let rhs_constants = rhs.constants(); + for rhs_constant in rhs_constants { + assert!( + const_exprs_contains(lhs_constants, rhs_constant.expr()), + "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" + ); + } + assert_eq!( + lhs_constants.len(), + rhs_constants.len(), + "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" + ); + + // Check whether orderings are same. + let lhs_orderings = lhs.oeq_class(); + let rhs_orderings = rhs.oeq_class(); + for rhs_ordering in rhs_orderings.iter() { + assert!( + lhs_orderings.contains(rhs_ordering), + "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" + ); + } + assert_eq!( + lhs_orderings.len(), + rhs_orderings.len(), + "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" + ); + } + + /// Make equivalence properties for the specified columns named in orderings and constants + /// + /// orderings: strings formatted like `"a"` or `"a DESC"`. See [`parse_sort_expr`] + /// constants: strings formatted like `"a"`. + fn make_props( + &self, + orderings: Vec>, + constants: Vec<&str>, + schema: &SchemaRef, + ) -> EquivalenceProperties { + let orderings = orderings + .iter() + .map(|ordering| { + ordering + .iter() + .map(|name| parse_sort_expr(name, schema)) + .collect::() + }) + .collect::>(); + + let constants = constants + .iter() + .map(|col_name| ConstExpr::new(col(col_name, schema).unwrap())) + .collect::>(); + + EquivalenceProperties::new_with_orderings(Arc::clone(schema), &orderings) + .with_constants(constants) + } + } + + #[test] + fn test_union_constant_value_preservation() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + let col_a = col("a", &schema)?; + let literal_10 = ScalarValue::Int32(Some(10)); + + // Create first input with a=10 + let const_expr1 = ConstExpr::new(Arc::clone(&col_a)) + .with_across_partitions(AcrossPartitions::Uniform(Some(literal_10.clone()))); + let input1 = EquivalenceProperties::new(Arc::clone(&schema)) + .with_constants(vec![const_expr1]); + + // Create second input with a=10 + let const_expr2 = ConstExpr::new(Arc::clone(&col_a)) + .with_across_partitions(AcrossPartitions::Uniform(Some(literal_10.clone()))); + let input2 = EquivalenceProperties::new(Arc::clone(&schema)) + .with_constants(vec![const_expr2]); + + // Calculate union properties + let union_props = calculate_union(vec![input1, input2], schema)?; + + // Verify column 'a' remains constant with value 10 + let const_a = &union_props.constants()[0]; + assert!(const_a.expr().eq(&col_a)); + assert_eq!( + const_a.across_partitions(), + AcrossPartitions::Uniform(Some(literal_10)) + ); + + Ok(()) + } + + /// Return a new schema with the same types, but new field names + /// + /// The new field names are the old field names with `text` appended. + /// + /// For example, the schema "a", "b", "c" becomes "a1", "b1", "c1" + /// if `text` is "1". + fn append_fields(schema: &SchemaRef, text: &str) -> SchemaRef { + Arc::new(Schema::new( + schema + .fields() + .iter() + .map(|field| { + Field::new( + // Annotate name with `text`: + format!("{}{}", field.name(), text), + field.data_type().clone(), + field.is_nullable(), + ) + }) + .collect::>(), + )) + } +} diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 0a448fa6a2e9..3671eaef1332 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -36,10 +36,6 @@ mod partitioning; mod physical_expr; pub mod planner; mod scalar_function; -pub mod udf { - #[allow(deprecated)] - pub use crate::scalar_function::create_physical_expr; -} pub mod statistics; pub mod utils; pub mod window; diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index bd38fb22ccbc..cf8cc6e00c80 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -39,12 +39,12 @@ use crate::PhysicalExpr; use arrow::array::{Array, RecordBatch}; use arrow::datatypes::{DataType, Schema}; -use datafusion_common::{internal_err, DFSchema, Result, ScalarValue}; +use datafusion_common::{internal_err, Result, ScalarValue}; use datafusion_expr::interval_arithmetic::Interval; use datafusion_expr::sort_properties::ExprProperties; use datafusion_expr::type_coercion::functions::data_types_with_scalar_udf; use datafusion_expr::{ - expr_vec_fmt, ColumnarValue, Expr, ReturnTypeArgs, ScalarFunctionArgs, ScalarUDF, + expr_vec_fmt, ColumnarValue, ReturnTypeArgs, ScalarFunctionArgs, ScalarUDF, }; /// Physical expression of a scalar function @@ -261,35 +261,3 @@ impl PhysicalExpr for ScalarFunctionExpr { }) } } - -/// Create a physical expression for the UDF. -#[deprecated(since = "45.0.0", note = "use ScalarFunctionExpr::new() instead")] -pub fn create_physical_expr( - fun: &ScalarUDF, - input_phy_exprs: &[Arc], - input_schema: &Schema, - args: &[Expr], - input_dfschema: &DFSchema, -) -> Result> { - let input_expr_types = input_phy_exprs - .iter() - .map(|e| e.data_type(input_schema)) - .collect::>>()?; - - // verify that input data types is consistent with function's `TypeSignature` - data_types_with_scalar_udf(&input_expr_types, fun)?; - - // Since we have arg_types, we don't need args and schema. - let return_type = - fun.return_type_from_exprs(args, input_dfschema, &input_expr_types)?; - - Ok(Arc::new( - ScalarFunctionExpr::new( - fun.name(), - Arc::new(fun.clone()), - input_phy_exprs.to_vec(), - return_type, - ) - .with_nullable(fun.is_nullable(args, input_dfschema)), - )) -} diff --git a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs index 17acb6272938..2e20608d0e9e 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs @@ -23,9 +23,7 @@ use crate::utils::{ }; use arrow::datatypes::SchemaRef; -use datafusion_common::tree_node::{ - ConcreteTreeNode, Transformed, TreeNode, TreeNodeRecursion, -}; +use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{plan_err, HashSet, JoinSide, Result}; use datafusion_expr::JoinType; use datafusion_physical_expr::expressions::Column; @@ -59,9 +57,9 @@ pub struct ParentRequirements { pub type SortPushDown = PlanContext; /// Assigns the ordering requirement of the root node to the its children. -pub fn assign_initial_requirements(node: &mut SortPushDown) { - let reqs = node.plan.required_input_ordering(); - for (child, requirement) in node.children.iter_mut().zip(reqs) { +pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { + let reqs = sort_push_down.plan.required_input_ordering(); + for (child, requirement) in sort_push_down.children.iter_mut().zip(reqs) { child.data = ParentRequirements { ordering_requirement: requirement, // If the parent has a fetch value, assign it to the children @@ -71,24 +69,26 @@ pub fn assign_initial_requirements(node: &mut SortPushDown) { } } -pub fn pushdown_sorts(sort_pushdown: SortPushDown) -> Result { - let mut new_node = pushdown_sorts_helper(sort_pushdown)?; - while new_node.tnr == TreeNodeRecursion::Stop { - new_node = pushdown_sorts_helper(new_node.data)?; +pub fn pushdown_sorts(sort_push_down: SortPushDown) -> Result { + sort_push_down + .transform_down(pushdown_sorts_helper) + .map(|transformed| transformed.data) +} + +fn min_fetch(f1: Option, f2: Option) -> Option { + match (f1, f2) { + (Some(f1), Some(f2)) => Some(f1.min(f2)), + (Some(_), _) => f1, + (_, Some(_)) => f2, + _ => None, } - let (new_node, children) = new_node.data.take_children(); - let new_children = children - .into_iter() - .map(pushdown_sorts) - .collect::>()?; - new_node.with_new_children(new_children) } fn pushdown_sorts_helper( - mut requirements: SortPushDown, + mut sort_push_down: SortPushDown, ) -> Result> { - let plan = &requirements.plan; - let parent_reqs = requirements + let plan = &sort_push_down.plan; + let parent_reqs = sort_push_down .data .ordering_requirement .clone() @@ -98,82 +98,102 @@ fn pushdown_sorts_helper( .ordering_satisfy_requirement(&parent_reqs); if is_sort(plan) { - let sort_fetch = plan.fetch(); - let required_ordering = plan + let current_sort_fetch = plan.fetch(); + let parent_req_fetch = sort_push_down.data.fetch; + + let current_plan_reqs = plan .output_ordering() .cloned() .map(LexRequirement::from) .unwrap_or_default(); - if !satisfy_parent { - // Make sure this `SortExec` satisfies parent requirements: - let sort_reqs = requirements.data.ordering_requirement.unwrap_or_default(); - // It's possible current plan (`SortExec`) has a fetch value. - // And if both of them have fetch values, we should use the minimum one. - if let Some(fetch) = sort_fetch { - if let Some(requirement_fetch) = requirements.data.fetch { - requirements.data.fetch = Some(fetch.min(requirement_fetch)); - } - } - let fetch = requirements.data.fetch.or(sort_fetch); - requirements = requirements.children.swap_remove(0); - requirements = add_sort_above(requirements, sort_reqs, fetch); - }; + let parent_is_stricter = plan + .equivalence_properties() + .requirements_compatible(&parent_reqs, ¤t_plan_reqs); + let current_is_stricter = plan + .equivalence_properties() + .requirements_compatible(¤t_plan_reqs, &parent_reqs); + + if !satisfy_parent && !parent_is_stricter { + // This new sort has different requirements than the ordering being pushed down. + // 1. add a `SortExec` here for the pushed down ordering (parent reqs). + // 2. continue sort pushdown, but with the new ordering of the new sort. + + // remove current sort (which will be the new ordering to pushdown) + let new_reqs = current_plan_reqs; + sort_push_down = sort_push_down.children.swap_remove(0); + sort_push_down = sort_push_down.update_plan_from_children()?; // changed plan + + // add back sort exec matching parent + sort_push_down = + add_sort_above(sort_push_down, parent_reqs, parent_req_fetch); + + // make pushdown requirements be the new ones. + sort_push_down.children[0].data = ParentRequirements { + ordering_requirement: Some(new_reqs), + fetch: current_sort_fetch, + }; + } else { + // Don't add a SortExec + // Do update what sort requirements to keep pushing down - // We can safely get the 0th index as we are dealing with a `SortExec`. - let mut child = requirements.children.swap_remove(0); - if let Some(adjusted) = - pushdown_requirement_to_children(&child.plan, &required_ordering)? - { - let fetch = sort_fetch.or_else(|| child.plan.fetch()); - for (grand_child, order) in child.children.iter_mut().zip(adjusted) { - grand_child.data = ParentRequirements { - ordering_requirement: order, - fetch, - }; + // remove current sort, and get the sort's child + sort_push_down = sort_push_down.children.swap_remove(0); + sort_push_down = sort_push_down.update_plan_from_children()?; // changed plan + + // set the stricter fetch + sort_push_down.data.fetch = min_fetch(current_sort_fetch, parent_req_fetch); + + // set the stricter ordering + if current_is_stricter { + sort_push_down.data.ordering_requirement = Some(current_plan_reqs); + } else { + sort_push_down.data.ordering_requirement = Some(parent_reqs); } - // Can push down requirements - child.data = ParentRequirements { - ordering_requirement: Some(required_ordering), - fetch, - }; - return Ok(Transformed { - data: child, - transformed: true, - tnr: TreeNodeRecursion::Stop, - }); - } else { - // Can not push down requirements - requirements.children = vec![child]; - assign_initial_requirements(&mut requirements); + // recursive call to helper, so it doesn't transform_down and miss the new node (previous child of sort) + return pushdown_sorts_helper(sort_push_down); } + } else if parent_reqs.is_empty() { + // note: this `satisfy_parent`, but we don't want to push down anything. + // Nothing to do. + return Ok(Transformed::no(sort_push_down)); } else if satisfy_parent { - // For non-sort operators, immediately return if parent requirements are met: + // For non-sort operators which satisfy ordering: let reqs = plan.required_input_ordering(); - for (child, order) in requirements.children.iter_mut().zip(reqs) { + let parent_req_fetch = sort_push_down.data.fetch; + + for (child, order) in sort_push_down.children.iter_mut().zip(reqs) { child.data.ordering_requirement = order; + child.data.fetch = min_fetch(parent_req_fetch, child.data.fetch); } } else if let Some(adjusted) = pushdown_requirement_to_children(plan, &parent_reqs)? { - // Can not satisfy the parent requirements, check whether we can push - // requirements down: - for (child, order) in requirements.children.iter_mut().zip(adjusted) { + // For operators that can take a sort pushdown. + + // Continue pushdown, with updated requirements: + let parent_fetch = sort_push_down.data.fetch; + let current_fetch = plan.fetch(); + for (child, order) in sort_push_down.children.iter_mut().zip(adjusted) { child.data.ordering_requirement = order; + child.data.fetch = min_fetch(current_fetch, parent_fetch); } - requirements.data.ordering_requirement = None; + sort_push_down.data.ordering_requirement = None; } else { // Can not push down requirements, add new `SortExec`: - let sort_reqs = requirements + let sort_reqs = sort_push_down .data .ordering_requirement .clone() .unwrap_or_default(); - let fetch = requirements.data.fetch; - requirements = add_sort_above(requirements, sort_reqs, fetch); - assign_initial_requirements(&mut requirements); + let fetch = sort_push_down.data.fetch; + sort_push_down = add_sort_above(sort_push_down, sort_reqs, fetch); + assign_initial_requirements(&mut sort_push_down); } - Ok(Transformed::yes(requirements)) + + Ok(Transformed::yes(sort_push_down)) } +/// Calculate the pushdown ordering requirements for children. +/// If sort cannot be pushed down, return None. fn pushdown_requirement_to_children( plan: &Arc, parent_required: &LexRequirement, diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index 88f11f53491e..bab31150e250 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -121,6 +121,10 @@ impl PhysicalOptimizer { // into an `order by max(x) limit y`. In this case it will copy the limit value down // to the aggregation, allowing it to use only y number of accumulators. Arc::new(TopKAggregation::new()), + // The LimitPushdown rule tries to push limits down as far as possible, + // replacing operators with fetching variants, or adding limits + // past operators that support limit pushdown. + Arc::new(LimitPushdown::new()), // The ProjectionPushdown rule tries to push projections towards // the sources in the execution plan. As a result of this process, // a projection can disappear if it reaches the source providers, and @@ -128,10 +132,6 @@ impl PhysicalOptimizer { // are not present, the load of executors such as join or union will be // reduced by narrowing their input tables. Arc::new(ProjectionPushdown::new()), - // The LimitPushdown rule tries to push limits down as far as possible, - // replacing operators with fetching variants, or adding limits - // past operators that support limit pushdown. - Arc::new(LimitPushdown::new()), // The SanityCheckPlan rule checks whether the order and // distribution requirements of each node in the plan // is satisfied. It will also reject non-runnable query diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 90a570894a44..3ca0547aa11d 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -132,10 +132,18 @@ impl OutputRequirementExec { impl DisplayAs for OutputRequirementExec { fn fmt_as( &self, - _t: DisplayFormatType, + t: DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { - write!(f, "OutputRequirementExec") + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "OutputRequirementExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } } } diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 636e78a06ce7..57a193315a5c 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -31,6 +31,10 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; /// This utility function adds a `SortExec` above an operator according to the /// given ordering requirements while preserving the original partitioning. +/// +/// Note that this updates the plan in both the [`PlanContext.children`] and +/// the [`PlanContext.plan`]'s children. Therefore its not required to sync +/// the child plans with [`PlanContext::update_plan_from_children`]. pub fn add_sort_above( node: PlanContext, sort_requirements: LexRequirement, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 0947a2ff5539..5dccc09fc722 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -57,41 +57,60 @@ mod row_hash; mod topk; mod topk_stream; -/// Hash aggregate modes +/// Aggregation modes /// /// See [`Accumulator::state`] for background information on multi-phase /// aggregation and how these modes are used. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub enum AggregateMode { + /// One of multiple layers of aggregation, any input partitioning + /// /// Partial aggregate that can be applied in parallel across input /// partitions. /// /// This is the first phase of a multi-phase aggregation. Partial, + /// *Final* of multiple layers of aggregation, in exactly one partition + /// /// Final aggregate that produces a single partition of output by combining /// the output of multiple partial aggregates. /// /// This is the second phase of a multi-phase aggregation. + /// + /// This mode requires that the input is a single partition + /// + /// Note: Adjacent `Partial` and `Final` mode aggregation is equivalent to a `Single` + /// mode aggregation node. The `Final` mode is required since this is used in an + /// intermediate step. The [`CombinePartialFinalAggregate`] physical optimizer rule + /// will replace this combination with `Single` mode for more efficient execution. + /// + /// [`CombinePartialFinalAggregate`]: https://docs.rs/datafusion/latest/datafusion/physical_optimizer/combine_partial_final_agg/struct.CombinePartialFinalAggregate.html Final, + /// *Final* of multiple layers of aggregation, input is *Partitioned* + /// /// Final aggregate that works on pre-partitioned data. /// - /// This requires the invariant that all rows with a particular - /// grouping key are in the same partitions, such as is the case - /// with Hash repartitioning on the group keys. If a group key is - /// duplicated, duplicate groups would be produced + /// This mode requires that all rows with a particular grouping key are in + /// the same partitions, such as is the case with Hash repartitioning on the + /// group keys. If a group key is duplicated, duplicate groups would be + /// produced FinalPartitioned, + /// *Single* layer of Aggregation, input is exactly one partition + /// /// Applies the entire logical aggregation operation in a single operator, /// as opposed to Partial / Final modes which apply the logical aggregation using /// two operators. /// /// This mode requires that the input is a single partition (like Final) Single, + /// *Single* layer of Aggregation, input is *Partitioned* + /// /// Applies the entire logical aggregation operation in a single operator, - /// as opposed to Partial / Final modes which apply the logical aggregation using - /// two operators. + /// as opposed to Partial / Final modes which apply the logical aggregation + /// using two operators. /// - /// This mode requires that the input is partitioned by group key (like - /// FinalPartitioned) + /// This mode requires that the input has more than one partition, and is + /// partitioned by group key (like FinalPartitioned). SinglePartitioned, } @@ -723,6 +742,15 @@ impl DisplayAs for AggregateExec { t: DisplayFormatType, f: &mut std::fmt::Formatter, ) -> std::fmt::Result { + let format_expr_with_alias = + |(e, alias): &(Arc, String)| -> String { + let e = e.to_string(); + if &e != alias { + format!("{e} as {alias}") + } else { + e + } + }; match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "AggregateExec: mode={:?}", self.mode)?; @@ -730,14 +758,7 @@ impl DisplayAs for AggregateExec { self.group_by .expr .iter() - .map(|(e, alias)| { - let e = e.to_string(); - if &e != alias { - format!("{e} as {alias}") - } else { - e - } - }) + .map(format_expr_with_alias) .collect() } else { self.group_by @@ -749,21 +770,11 @@ impl DisplayAs for AggregateExec { .enumerate() .map(|(idx, is_null)| { if *is_null { - let (e, alias) = &self.group_by.null_expr[idx]; - let e = e.to_string(); - if &e != alias { - format!("{e} as {alias}") - } else { - e - } + format_expr_with_alias( + &self.group_by.null_expr[idx], + ) } else { - let (e, alias) = &self.group_by.expr[idx]; - let e = e.to_string(); - if &e != alias { - format!("{e} as {alias}") - } else { - e - } + format_expr_with_alias(&self.group_by.expr[idx]) } }) .collect::>() @@ -789,6 +800,47 @@ impl DisplayAs for AggregateExec { write!(f, ", ordering_mode={:?}", self.input_order_mode)?; } } + DisplayFormatType::TreeRender => { + let g: Vec = if self.group_by.is_single() { + self.group_by + .expr + .iter() + .map(format_expr_with_alias) + .collect() + } else { + self.group_by + .groups + .iter() + .map(|group| { + let terms = group + .iter() + .enumerate() + .map(|(idx, is_null)| { + if *is_null { + format_expr_with_alias( + &self.group_by.null_expr[idx], + ) + } else { + format_expr_with_alias(&self.group_by.expr[idx]) + } + }) + .collect::>() + .join(", "); + format!("({terms})") + }) + .collect() + }; + let a: Vec = self + .aggr_expr + .iter() + .map(|agg| agg.name().to_string()) + .collect(); + writeln!(f, "mode={:?}", self.mode)?; + if !g.is_empty() { + writeln!(f, "group_by={}", g.join(", "))?; + } + writeln!(f, "aggr={}", a.join(", "))?; + } } Ok(()) } @@ -1797,6 +1849,10 @@ mod tests { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "TestYieldingExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 708f006b0d39..ea14ce676c1a 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -108,6 +108,10 @@ impl DisplayAs for AnalyzeExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "AnalyzeExec verbose={}", self.verbose) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index fa8d125d62d1..0eb95bb66598 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -122,6 +122,10 @@ impl DisplayAs for CoalesceBatchesExec { Ok(()) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 9a955155c01e..8fb40640dcc0 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -92,6 +92,10 @@ impl DisplayAs for CoalescePartitionsExec { } None => write!(f, "CoalescePartitionsExec"), }, + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/common.rs b/datafusion/physical-plan/src/common.rs index b83641acf2ce..a8d4a3ddf3d1 100644 --- a/datafusion/physical-plan/src/common.rs +++ b/datafusion/physical-plan/src/common.rs @@ -180,7 +180,7 @@ pub fn compute_record_batch_statistics( } } -/// Write in Arrow IPC format. +/// Write in Arrow IPC File format. pub struct IPCWriter { /// Path pub path: PathBuf, diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 0cc1cb02438a..564f7ac45928 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -18,24 +18,80 @@ //! Implementation of physical plan display. See //! [`crate::displayable`] for examples of how to format -use std::fmt; +use std::collections::{BTreeMap, HashMap}; use std::fmt::Formatter; +use std::{fmt, str::FromStr}; use arrow::datatypes::SchemaRef; use datafusion_common::display::{GraphvizBuilder, PlanType, StringifiedPlan}; +use datafusion_common::DataFusionError; use datafusion_expr::display_schema; use datafusion_physical_expr::LexOrdering; +use crate::render_tree::RenderTree; + use super::{accept, ExecutionPlan, ExecutionPlanVisitor}; /// Options for controlling how each [`ExecutionPlan`] should format itself -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq)] pub enum DisplayFormatType { /// Default, compact format. Example: `FilterExec: c12 < 10.0` + /// + /// This format is designed to provide a detailed textual description + /// of all rele Default, - /// Verbose, showing all available details + /// Verbose, showing all available details. + /// + /// This form is even more detailed than [`Self::Default`] Verbose, + /// TreeRender, displayed in the `tree` explain type. + /// + /// This format is inspired by DuckDB's explain plans. The information + /// presented should be "user friendly", and contain only the most relevant + /// information for understanding a plan. It should NOT contain the same level + /// of detail information as the [`Self::Default`] format. + /// + /// In this mode, each line has one of two formats: + /// + /// 1. A string without a `=`, which is printed in its own line + /// + /// 2. A string with a `=` that is treated as a `key=value pair`. Everything + /// before the first `=` is treated as the key, and everything after the + /// first `=` is treated as the value. + /// + /// For example, if the output of `TreeRender` is this: + /// ```text + /// Parquet + /// partition_sizes=[1] + /// ``` + /// + /// It is rendered in the center of a box in the following way: + /// + /// ```text + /// ┌───────────────────────────┐ + /// │ DataSourceExec │ + /// │ -------------------- │ + /// │ partition_sizes: [1] │ + /// │ Parquet │ + /// └───────────────────────────┘ + /// ``` + TreeRender, +} + +impl FromStr for DisplayFormatType { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "indent" => Ok(Self::Default), + "tree" => Ok(Self::TreeRender), + _ => Err(DataFusionError::Configuration(format!( + "Invalid explain format: {}", + s + ))), + } + } } /// Wraps an `ExecutionPlan` with various methods for formatting @@ -224,6 +280,19 @@ impl<'a> DisplayableExecutionPlan<'a> { } } + pub fn tree_render(&self) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = TreeRenderVisitor { f }; + visitor.visit(self.plan) + } + } + Wrapper { plan: self.inner } + } + /// Return a single-line summary of the root of the plan /// Example: `ProjectionExec: expr=[a@0 as a]`. pub fn one_line(&self) -> impl fmt::Display + 'a { @@ -258,8 +327,18 @@ impl<'a> DisplayableExecutionPlan<'a> { } /// format as a `StringifiedPlan` - pub fn to_stringified(&self, verbose: bool, plan_type: PlanType) -> StringifiedPlan { - StringifiedPlan::new(plan_type, self.indent(verbose).to_string()) + pub fn to_stringified( + &self, + verbose: bool, + plan_type: PlanType, + explain_format: DisplayFormatType, + ) -> StringifiedPlan { + match (&explain_format, &plan_type) { + (DisplayFormatType::TreeRender, PlanType::FinalPhysicalPlan) => { + StringifiedPlan::new(plan_type, self.tree_render().to_string()) + } + _ => StringifiedPlan::new(plan_type, self.indent(verbose).to_string()), + } } } @@ -448,6 +527,482 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { } } +/// This module implements a tree-like art renderer for execution plans, +/// based on DuckDB's implementation: +/// +/// +/// The rendered output looks like this: +/// ```text +/// ┌───────────────────────────┐ +/// │ CoalesceBatchesExec │ +/// └─────────────┬─────────────┘ +/// ┌─────────────┴─────────────┐ +/// │ HashJoinExec ├──────────────┐ +/// └─────────────┬─────────────┘ │ +/// ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +/// │ DataSourceExec ││ DataSourceExec │ +/// └───────────────────────────┘└───────────────────────────┘ +/// ``` +/// +/// The renderer uses a three-layer approach for each node: +/// 1. Top layer: renders the top borders and connections +/// 2. Content layer: renders the node content and vertical connections +/// 3. Bottom layer: renders the bottom borders and connections +/// +/// Each node is rendered in a box of fixed width (NODE_RENDER_WIDTH). +struct TreeRenderVisitor<'a, 'b> { + /// Write to this formatter + f: &'a mut Formatter<'b>, +} + +impl TreeRenderVisitor<'_, '_> { + // Unicode box-drawing characters for creating borders and connections. + const LTCORNER: &'static str = "┌"; // Left top corner + const RTCORNER: &'static str = "┐"; // Right top corner + const LDCORNER: &'static str = "└"; // Left bottom corner + const RDCORNER: &'static str = "┘"; // Right bottom corner + + const TMIDDLE: &'static str = "┬"; // Top T-junction (connects down) + const LMIDDLE: &'static str = "├"; // Left T-junction (connects right) + const DMIDDLE: &'static str = "┴"; // Bottom T-junction (connects up) + + const VERTICAL: &'static str = "│"; // Vertical line + const HORIZONTAL: &'static str = "─"; // Horizontal line + + // TODO: Make these variables configurable. + const MAXIMUM_RENDER_WIDTH: usize = 240; // Maximum total width of the rendered tree + const NODE_RENDER_WIDTH: usize = 29; // Width of each node's box + const MAX_EXTRA_LINES: usize = 30; // Maximum number of extra info lines per node + + /// Main entry point for rendering an execution plan as a tree. + /// The rendering process happens in three stages for each level of the tree: + /// 1. Render top borders and connections + /// 2. Render node content and vertical connections + /// 3. Render bottom borders and connections + pub fn visit(&mut self, plan: &dyn ExecutionPlan) -> Result<(), fmt::Error> { + let root = RenderTree::create_tree(plan); + + for y in 0..root.height { + // Start by rendering the top layer. + self.render_top_layer(&root, y)?; + // Now we render the content of the boxes + self.render_box_content(&root, y)?; + // Render the bottom layer of each of the boxes + self.render_bottom_layer(&root, y)?; + } + + Ok(()) + } + + /// Renders the top layer of boxes at the given y-level of the tree. + /// This includes: + /// - Top corners (┌─┐) for nodes + /// - Horizontal connections between nodes + /// - Vertical connections to parent nodes + fn render_top_layer( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + for x in 0..root.width { + if root.has_node(x, y) { + write!(self.f, "{}", Self::LTCORNER)?; + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + if y == 0 { + // top level node: no node above this one + write!(self.f, "{}", Self::HORIZONTAL)?; + } else { + // render connection to node above this one + write!(self.f, "{}", Self::DMIDDLE)?; + } + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + write!(self.f, "{}", Self::RTCORNER)?; + } else { + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + if !has_adjacent_nodes { + // There are no nodes to the right side of this position + // no need to fill the empty space + continue; + } + // there are nodes next to this, fill the space + write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + + Ok(()) + } + + /// Renders the content layer of boxes at the given y-level of the tree. + /// This includes: + /// - Node names and extra information + /// - Vertical borders (│) for boxes + /// - Vertical connections between nodes + fn render_box_content( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + let mut extra_info: Vec> = vec![vec![]; root.width]; + let mut extra_height = 0; + + for (x, extra_info_item) in extra_info.iter_mut().enumerate().take(root.width) { + if let Some(node) = root.get_node(x, y) { + Self::split_up_extra_info( + &node.extra_text, + extra_info_item, + Self::MAX_EXTRA_LINES, + ); + if extra_info_item.len() > extra_height { + extra_height = extra_info_item.len(); + } + } + } + + let halfway_point = (extra_height + 1) / 2; + + // Render the actual node. + for render_y in 0..=extra_height { + for (x, _) in root.nodes.iter().enumerate().take(root.width) { + if x * Self::NODE_RENDER_WIDTH >= Self::MAXIMUM_RENDER_WIDTH { + break; + } + + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + + if let Some(node) = root.get_node(x, y) { + write!(self.f, "{}", Self::VERTICAL)?; + + // Rigure out what to render. + let mut render_text = String::new(); + if render_y == 0 { + render_text = node.name.clone(); + } else if render_y <= extra_info[x].len() { + render_text = extra_info[x][render_y - 1].clone(); + } + + render_text = Self::adjust_text_for_rendering( + &render_text, + Self::NODE_RENDER_WIDTH - 2, + ); + write!(self.f, "{}", render_text)?; + + if render_y == halfway_point && node.child_positions.len() > 1 { + write!(self.f, "{}", Self::LMIDDLE)?; + } else { + write!(self.f, "{}", Self::VERTICAL)?; + } + } else if render_y == halfway_point { + let has_child_to_the_right = + Self::should_render_whitespace(root, x, y); + if root.has_node(x, y + 1) { + // Node right below this one. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + if has_child_to_the_right { + write!(self.f, "{}", Self::TMIDDLE)?; + // Have another child to the right, Keep rendering the line. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } else { + write!(self.f, "{}", Self::RTCORNER)?; + if has_adjacent_nodes { + // Only a child below this one: fill the reset with spaces. + write!( + self.f, + "{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } + } + } else if has_child_to_the_right { + // Child to the right, but no child right below this one: render a full + // line. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH) + )?; + } else if has_adjacent_nodes { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } else if render_y >= halfway_point { + if root.has_node(x, y + 1) { + // Have a node below this empty spot: render a vertical line. + write!( + self.f, + "{}{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2), + Self::VERTICAL + )?; + if has_adjacent_nodes + || Self::should_render_whitespace(root, x, y) + { + write!( + self.f, + "{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } + } else if has_adjacent_nodes + || Self::should_render_whitespace(root, x, y) + { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } else if has_adjacent_nodes { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + } + + Ok(()) + } + + /// Renders the bottom layer of boxes at the given y-level of the tree. + /// This includes: + /// - Bottom corners (└─┘) for nodes + /// - Horizontal connections between nodes + /// - Vertical connections to child nodes + fn render_bottom_layer( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + for x in 0..=root.width { + if x * Self::NODE_RENDER_WIDTH >= Self::MAXIMUM_RENDER_WIDTH { + break; + } + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + if root.get_node(x, y).is_some() { + write!(self.f, "{}", Self::LDCORNER)?; + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + if root.has_node(x, y + 1) { + // node below this one: connect to that one + write!(self.f, "{}", Self::TMIDDLE)?; + } else { + // no node below this one: end the box + write!(self.f, "{}", Self::HORIZONTAL)?; + } + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + write!(self.f, "{}", Self::RDCORNER)?; + } else if root.has_node(x, y + 1) { + write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + write!(self.f, "{}", Self::VERTICAL)?; + if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { + write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + } + } else if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { + write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + + Ok(()) + } + + fn extra_info_separator() -> String { + "-".repeat(Self::NODE_RENDER_WIDTH - 9) + } + + fn remove_padding(s: &str) -> String { + s.trim().to_string() + } + + pub fn split_up_extra_info( + extra_info: &HashMap, + result: &mut Vec, + max_lines: usize, + ) { + if extra_info.is_empty() { + return; + } + + result.push(Self::extra_info_separator()); + + let mut requires_padding = false; + let mut was_inlined = false; + + // use BTreeMap for repeatable key order + let sorted_extra_info: BTreeMap<_, _> = extra_info.iter().collect(); + for (key, value) in sorted_extra_info { + let mut str = Self::remove_padding(value); + let mut is_inlined = false; + let available_width = Self::NODE_RENDER_WIDTH - 7; + let total_size = key.len() + str.len() + 2; + let is_multiline = str.contains('\n'); + + if str.is_empty() { + str = key.to_string(); + } else if !is_multiline && total_size < available_width { + str = format!("{}: {}", key, str); + is_inlined = true; + } else { + str = format!("{}:\n{}", key, str); + } + + if is_inlined && was_inlined { + requires_padding = false; + } + + if requires_padding { + result.push(String::new()); + } + + let mut splits: Vec = str.split('\n').map(String::from).collect(); + if splits.len() > max_lines { + let mut truncated_splits = Vec::new(); + for split in splits.iter().take(max_lines / 2) { + truncated_splits.push(split.clone()); + } + truncated_splits.push("...".to_string()); + for split in splits.iter().skip(splits.len() - max_lines / 2) { + truncated_splits.push(split.clone()); + } + splits = truncated_splits; + } + for split in splits { + Self::split_string_buffer(&split, result); + } + if result.len() > max_lines { + result.truncate(max_lines); + result.push("...".to_string()); + } + + requires_padding = true; + was_inlined = is_inlined; + } + } + + /// Adjusts text to fit within the specified width by: + /// 1. Truncating with ellipsis if too long + /// 2. Center-aligning within the available space if shorter + fn adjust_text_for_rendering(source: &str, max_render_width: usize) -> String { + let render_width = source.chars().count(); + if render_width > max_render_width { + let truncated = &source[..max_render_width - 3]; + format!("{}...", truncated) + } else { + let total_spaces = max_render_width - render_width; + let half_spaces = total_spaces / 2; + let extra_left_space = if total_spaces % 2 == 0 { 0 } else { 1 }; + format!( + "{}{}{}", + " ".repeat(half_spaces + extra_left_space), + source, + " ".repeat(half_spaces) + ) + } + } + + /// Determines if whitespace should be rendered at a given position. + /// This is important for: + /// 1. Maintaining proper spacing between sibling nodes + /// 2. Ensuring correct alignment of connections between parents and children + /// 3. Preserving the tree structure's visual clarity + fn should_render_whitespace(root: &RenderTree, x: usize, y: usize) -> bool { + let mut found_children = 0; + + for i in (0..=x).rev() { + let node = root.get_node(i, y); + if root.has_node(i, y + 1) { + found_children += 1; + } + if let Some(node) = node { + if node.child_positions.len() > 1 + && found_children < node.child_positions.len() + { + return true; + } + + return false; + } + } + + false + } + + fn split_string_buffer(source: &str, result: &mut Vec) { + let mut character_pos = 0; + let mut start_pos = 0; + let mut render_width = 0; + let mut last_possible_split = 0; + + let chars: Vec = source.chars().collect(); + + while character_pos < chars.len() { + // Treating each char as width 1 for simplification + let char_width = 1; + + // Does the next character make us exceed the line length? + if render_width + char_width > Self::NODE_RENDER_WIDTH - 2 { + if start_pos + 8 > last_possible_split { + // The last character we can split on is one of the first 8 characters of the line + // to not create very small lines we instead split on the current character + last_possible_split = character_pos; + } + + result.push(source[start_pos..last_possible_split].to_string()); + render_width = character_pos - last_possible_split; + start_pos = last_possible_split; + character_pos = last_possible_split; + } + + // check if we can split on this character + if Self::can_split_on_this_char(chars[character_pos]) { + last_possible_split = character_pos; + } + + character_pos += 1; + render_width += char_width; + } + + if source.len() > start_pos { + // append the remainder of the input + result.push(source[start_pos..].to_string()); + } + } + + fn can_split_on_this_char(c: char) -> bool { + (!c.is_ascii_digit() && !c.is_ascii_uppercase() && !c.is_ascii_lowercase()) + && c != '_' + } +} + /// Trait for types which could have additional details when formatted in `Verbose` mode pub trait DisplayAs { /// Format according to `DisplayFormatType`, used when verbose representation looks diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index c4e738cb3ad1..3fdde39df6f1 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -94,6 +94,10 @@ impl DisplayAs for EmptyExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "EmptyExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 851e504b69af..d7556bc07c45 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -27,7 +27,7 @@ pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; pub use datafusion_expr::{Accumulator, ColumnarValue}; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ - expressions, udf, Distribution, Partitioning, PhysicalExpr, + expressions, Distribution, Partitioning, PhysicalExpr, }; use std::any::Any; diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index cb00958cec4c..bf488ccfae56 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -94,6 +94,10 @@ impl DisplayAs for ExplainExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "ExplainExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index a66873bc6576..ffcda1d888b0 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -329,6 +329,9 @@ impl DisplayAs for FilterExec { }; write!(f, "FilterExec: {}{}", self.predicate, display_projections) } + DisplayFormatType::TreeRender => { + write!(f, "predicate={}", self.predicate) + } } } } diff --git a/datafusion/physical-plan/src/insert.rs b/datafusion/physical-plan/src/insert.rs index 63c9c9921248..5272f0ab1867 100644 --- a/datafusion/physical-plan/src/insert.rs +++ b/datafusion/physical-plan/src/insert.rs @@ -153,6 +153,7 @@ impl DisplayAs for DataSinkExec { write!(f, "DataSinkExec: sink=")?; self.sink.fmt_as(t, f) } + DisplayFormatType::TreeRender => self.sink().fmt_as(t, f), } } } diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index ca4c26251de0..35c8961065a5 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -240,6 +240,10 @@ impl DisplayAs for CrossJoinExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "CrossJoinExec") } + DisplayFormatType::TreeRender => { + // no extra info to display + Ok(()) + } } } } diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs index b2e9b37655f1..39a15037260d 100644 --- a/datafusion/physical-plan/src/joins/hash_join.rs +++ b/datafusion/physical-plan/src/joins/hash_join.rs @@ -668,6 +668,19 @@ impl DisplayAs for HashJoinExec { self.mode, self.join_type, on, display_filter, display_projections ) } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({} = {})", c1, c2)) + .collect::>() + .join(", "); + + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={}", on) + } } } } diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 64dfc8219b64..f6fa8878e033 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -424,6 +424,13 @@ impl DisplayAs for NestedLoopJoinExec { self.join_type, display_filter, display_projections ) } + DisplayFormatType::TreeRender => { + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type) + } else { + Ok(()) + } + } } } } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join.rs b/datafusion/physical-plan/src/joins/sort_merge_join.rs index 6c933ca21807..d8446fb332b1 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join.rs @@ -59,7 +59,7 @@ use arrow::compute::{ }; use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; use arrow::error::ArrowError; -use arrow::ipc::reader::FileReader; +use arrow::ipc::reader::StreamReader; use datafusion_common::{ exec_err, internal_err, not_impl_err, plan_err, DataFusionError, HashSet, JoinSide, JoinType, Result, @@ -369,6 +369,19 @@ impl DisplayAs for SortMergeJoinExec { )) ) } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({} = {})", c1, c2)) + .collect::>() + .join(", "); + + if self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={}", on) + } } } } @@ -1394,7 +1407,7 @@ impl SortMergeJoinStream { if let Some(batch) = buffered_batch.batch { spill_record_batches( - vec![batch], + &[batch], spill_file.path().into(), Arc::clone(&self.buffered_schema), )?; @@ -2270,7 +2283,7 @@ fn fetch_right_columns_from_batch_by_idxs( Vec::with_capacity(buffered_indices.len()); let file = BufReader::new(File::open(spill_file.path())?); - let reader = FileReader::try_new(file, None)?; + let reader = StreamReader::try_new(file, None)?; for batch in reader { batch?.columns().iter().for_each(|column| { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 47af4ab9a765..63e95c7a3018 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -380,6 +380,20 @@ impl DisplayAs for SymmetricHashJoinExec { self.mode, self.join_type, on, display_filter ) } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({} = {})", c1, c2)) + .collect::>() + .join(", "); + + writeln!(f, "mode={:?}", self.mode)?; + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={}", on) + } } } } diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6ddaef1a2d28..7c6cac0a36c8 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -35,7 +35,7 @@ pub use datafusion_expr::{Accumulator, ColumnarValue}; pub use datafusion_physical_expr::window::WindowExpr; use datafusion_physical_expr::PhysicalSortExpr; pub use datafusion_physical_expr::{ - expressions, udf, Distribution, Partitioning, PhysicalExpr, + expressions, Distribution, Partitioning, PhysicalExpr, }; pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; @@ -51,6 +51,7 @@ pub use crate::topk::TopK; pub use crate::visitor::{accept, visit_execution_plan, ExecutionPlanVisitor}; mod ordering; +mod render_tree; mod topk; mod visitor; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index f720294c7ad9..b9464e3a88fb 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -108,6 +108,12 @@ impl DisplayAs for GlobalLimitExec { self.fetch.map_or("None".to_string(), |x| x.to_string()) ) } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.fetch { + writeln!(f, "limit={}", fetch)?; + } + write!(f, "skip={}", self.skip) + } } } } @@ -261,6 +267,10 @@ impl DisplayAs for LocalLimitExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "LocalLimitExec: fetch={}", self.fetch) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index fd338cc91353..ae878fdab37b 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -192,6 +192,10 @@ impl DisplayAs for LazyMemoryExec { .join(", ") ) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 6e31f601e152..72f2c13d2040 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -112,6 +112,11 @@ impl DisplayAs for PlaceholderRowExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "PlaceholderRowExec") } + + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 08c4d24f4c7f..3f901311a053 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -167,6 +167,17 @@ impl DisplayAs for ProjectionExec { write!(f, "ProjectionExec: expr=[{}]", expr.join(", ")) } + DisplayFormatType::TreeRender => { + for (i, (e, alias)) in self.expr().iter().enumerate() { + let e = e.to_string(); + if &e == alias { + writeln!(f, "expr{i}={e}")?; + } else { + writeln!(f, "{alias}={e}")?; + } + } + Ok(()) + } } } } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 05b78e4e1da4..7268735ea457 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -223,6 +223,10 @@ impl DisplayAs for RecursiveQueryExec { self.name, self.is_distinct ) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/render_tree.rs b/datafusion/physical-plan/src/render_tree.rs new file mode 100644 index 000000000000..f86e4c55e7b0 --- /dev/null +++ b/datafusion/physical-plan/src/render_tree.rs @@ -0,0 +1,230 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// This code is based on the DuckDB’s implementation: +// + +//! This module provides functionality for rendering an execution plan as a tree structure. +//! It helps in visualizing how different operations in a query are connected and organized. + +use std::collections::HashMap; +use std::fmt::Formatter; +use std::sync::Arc; +use std::{cmp, fmt}; + +use crate::{DisplayFormatType, ExecutionPlan}; + +// TODO: It's never used. +/// Represents a 2D coordinate in the rendered tree. +/// Used to track positions of nodes and their connections. +#[allow(dead_code)] +pub struct Coordinate { + /// Horizontal position in the tree + pub x: usize, + /// Vertical position in the tree + pub y: usize, +} + +impl Coordinate { + pub fn new(x: usize, y: usize) -> Self { + Coordinate { x, y } + } +} + +/// Represents a node in the render tree, containing information about an execution plan operator +/// and its relationships to other operators. +pub struct RenderTreeNode { + /// The name of physical `ExecutionPlan`. + pub name: String, + /// Execution info collected from `ExecutionPlan`. + pub extra_text: HashMap, + /// Positions of child nodes in the rendered tree. + pub child_positions: Vec, +} + +impl RenderTreeNode { + pub fn new(name: String, extra_text: HashMap) -> Self { + RenderTreeNode { + name, + extra_text, + child_positions: vec![], + } + } + + fn add_child_position(&mut self, x: usize, y: usize) { + self.child_positions.push(Coordinate::new(x, y)); + } +} + +/// Main structure for rendering an execution plan as a tree. +/// Manages a 2D grid of nodes and their layout information. +pub struct RenderTree { + /// Storage for tree nodes in a flattened 2D grid + pub nodes: Vec>>, + /// Total width of the rendered tree + pub width: usize, + /// Total height of the rendered tree + pub height: usize, +} + +impl RenderTree { + /// Creates a new render tree from an execution plan. + pub fn create_tree(plan: &dyn ExecutionPlan) -> Self { + let (width, height) = get_tree_width_height(plan); + + let mut result = Self::new(width, height); + + create_tree_recursive(&mut result, plan, 0, 0); + + result + } + + fn new(width: usize, height: usize) -> Self { + RenderTree { + nodes: vec![None; (width + 1) * (height + 1)], + width, + height, + } + } + + pub fn get_node(&self, x: usize, y: usize) -> Option> { + if x >= self.width || y >= self.height { + return None; + } + + let pos = self.get_position(x, y); + self.nodes.get(pos).and_then(|node| node.clone()) + } + + pub fn set_node(&mut self, x: usize, y: usize, node: Arc) { + let pos = self.get_position(x, y); + if let Some(slot) = self.nodes.get_mut(pos) { + *slot = Some(node); + } + } + + pub fn has_node(&self, x: usize, y: usize) -> bool { + if x >= self.width || y >= self.height { + return false; + } + + let pos = self.get_position(x, y); + self.nodes.get(pos).is_some_and(|node| node.is_some()) + } + + fn get_position(&self, x: usize, y: usize) -> usize { + y * self.width + x + } +} + +/// Calculates the required dimensions of the tree. +/// This ensures we allocate enough space for the entire tree structure. +/// +/// # Arguments +/// * `plan` - The execution plan to measure +/// +/// # Returns +/// * A tuple of (width, height) representing the dimensions needed for the tree +fn get_tree_width_height(plan: &dyn ExecutionPlan) -> (usize, usize) { + let children = plan.children(); + + // Leaf nodes take up 1x1 space + if children.is_empty() { + return (1, 1); + } + + let mut width = 0; + let mut height = 0; + + for child in children { + let (child_width, child_height) = get_tree_width_height(child.as_ref()); + width += child_width; + height = cmp::max(height, child_height); + } + + height += 1; + + (width, height) +} + +fn fmt_display(plan: &dyn ExecutionPlan) -> impl fmt::Display + '_ { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + } + + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.plan.fmt_as(DisplayFormatType::TreeRender, f)?; + Ok(()) + } + } + + Wrapper { plan } +} + +/// Recursively builds the render tree structure. +/// Traverses the execution plan and creates corresponding render nodes while +/// maintaining proper positioning and parent-child relationships. +/// +/// # Arguments +/// * `result` - The render tree being constructed +/// * `plan` - Current execution plan node being processed +/// * `x` - Horizontal position in the tree +/// * `y` - Vertical position in the tree +/// +/// # Returns +/// * The width of the subtree rooted at the current node +fn create_tree_recursive( + result: &mut RenderTree, + plan: &dyn ExecutionPlan, + x: usize, + y: usize, +) -> usize { + let display_info = fmt_display(plan).to_string(); + let mut extra_info = HashMap::new(); + + // Parse the key-value pairs from the formatted string. + // See DisplayFormatType::TreeRender for details + for line in display_info.lines() { + if let Some((key, value)) = line.split_once('=') { + extra_info.insert(key.to_string(), value.to_string()); + } else { + extra_info.insert(line.to_string(), "".to_string()); + } + } + + let mut node = RenderTreeNode::new(plan.name().to_string(), extra_info); + + let children = plan.children(); + + if children.is_empty() { + result.set_node(x, y, Arc::new(node)); + return 1; + } + + let mut width = 0; + for child in children { + let child_x = x + width; + let child_y = y + 1; + node.add_child_position(child_x, child_y); + width += create_tree_recursive(result, child.as_ref(), child_x, child_y); + } + + result.set_node(x, y, Arc::new(node)); + + width +} diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 40e68cfcae83..2b2548c8723c 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -506,6 +506,18 @@ impl DisplayAs for RepartitionExec { } Ok(()) } + DisplayFormatType::TreeRender => { + writeln!(f, "partitioning_scheme={}", self.partitioning(),)?; + writeln!( + f, + "output_partition_count={}", + self.input.output_partitioning().partition_count() + )?; + if self.preserve_order { + writeln!(f, "preserve_order={}", self.preserve_order)?; + } + Ok(()) + } } } } diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index dc03c012d9be..5277a50b85ca 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -226,6 +226,15 @@ impl DisplayAs for PartialSortExec { None => write!(f, "PartialSortExec: expr=[{}], common_prefix_length=[{common_prefix_length}]", self.expr), } } + DisplayFormatType::TreeRender => match self.fetch { + Some(fetch) => { + writeln!(f, "{}", self.expr)?; + writeln!(f, "limit={fetch}") + } + None => { + writeln!(f, "{}", self.expr) + } + }, } } } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index d84068527a64..3d2323eb4336 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -307,7 +307,7 @@ impl ExternalSorter { let size = get_reserved_byte_for_record_batch(&input); if self.reservation.try_grow(size).is_err() { - self.sort_or_spill_in_mem_batches().await?; + self.sort_or_spill_in_mem_batches(false).await?; // We've already freed more than half of reserved memory, // so we can grow the reservation again. There's nothing we can do // if this try_grow fails. @@ -332,7 +332,7 @@ impl ExternalSorter { /// /// 2. A combined streaming merge incorporating both in-memory /// batches and data from spill files on disk. - fn sort(&mut self) -> Result { + async fn sort(&mut self) -> Result { // Release the memory reserved for merge back to the pool so // there is some left when `in_mem_sort_stream` requests an // allocation. @@ -340,10 +340,12 @@ impl ExternalSorter { if self.spilled_before() { let mut streams = vec![]; + + // Sort `in_mem_batches` and spill it first. If there are many + // `in_mem_batches` and the memory limit is almost reached, merging + // them with the spilled files at the same time might cause OOM. if !self.in_mem_batches.is_empty() { - let in_mem_stream = - self.in_mem_sort_stream(self.metrics.baseline.intermediate())?; - streams.push(in_mem_stream); + self.sort_or_spill_in_mem_batches(true).await?; } for spill in self.spills.drain(..) { @@ -407,7 +409,7 @@ impl ExternalSorter { let spill_file = self.runtime.disk_manager.create_tmp_file("Sorting")?; let batches = std::mem::take(&mut self.in_mem_batches); let (spilled_rows, spilled_bytes) = spill_record_batches( - batches, + &batches, spill_file.path().into(), Arc::clone(&self.schema), )?; @@ -488,11 +490,17 @@ impl ExternalSorter { /// the memory usage has dropped by a factor of 2, then we don't have /// to spill. Otherwise, we spill to free up memory for inserting /// more batches. - /// /// The factor of 2 aims to avoid a degenerate case where the /// memory required for `fetch` is just under the memory available, - // causing repeated re-sorting of data - async fn sort_or_spill_in_mem_batches(&mut self) -> Result<()> { + /// causing repeated re-sorting of data + /// + /// # Arguments + /// + /// * `force_spill` - If true, the method will spill the in-memory batches + /// even if the memory usage has not dropped by a factor of 2. Otherwise it will + /// only spill when the memory usage has dropped by the pre-defined factor. + /// + async fn sort_or_spill_in_mem_batches(&mut self, force_spill: bool) -> Result<()> { // Release the memory reserved for merge back to the pool so // there is some left when `in_mem_sort_stream` requests an // allocation. At the end of this function, memory will be @@ -529,7 +537,7 @@ impl ExternalSorter { // Sorting may free up some memory especially when fetch is `Some`. If we have // not freed more than 50% of the memory, then we have to spill to free up more // memory for inserting more batches. - if self.reservation.size() > before / 2 { + if (self.reservation.size() > before / 2) || force_spill { // We have not freed more than 50% of the memory, so we have to spill to // free up more memory self.spill().await?; @@ -997,6 +1005,15 @@ impl DisplayAs for SortExec { None => write!(f, "SortExec: expr=[{}], preserve_partitioning=[{preserve_partitioning}]", self.expr), } } + DisplayFormatType::TreeRender => match self.fetch { + Some(fetch) => { + writeln!(f, "{}", self.expr)?; + writeln!(f, "limit={fetch}") + } + None => { + writeln!(f, "{}", self.expr) + } + }, } } } @@ -1110,7 +1127,7 @@ impl ExecutionPlan for SortExec { let batch = batch?; sorter.insert_batch(batch).await?; } - sorter.sort() + sorter.sort().await }) .try_flatten(), ))) @@ -1213,9 +1230,9 @@ mod tests { impl DisplayAs for SortedUnboundedExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "UnboundableExec",).unwrap() - } + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => write!(f, "UnboundableExec",).unwrap(), } Ok(()) } @@ -1401,7 +1418,7 @@ mod tests { // bytes. We leave a little wiggle room for the actual numbers. assert!((3..=10).contains(&spill_count)); assert!((9000..=10000).contains(&spilled_rows)); - assert!((36000..=40000).contains(&spilled_bytes)); + assert!((38000..=42000).contains(&spilled_bytes)); let columns = result[0].columns(); @@ -1474,7 +1491,7 @@ mod tests { // `number_of_batches / (sort_spill_reservation_bytes / batch_size)` assert!((12..=18).contains(&spill_count)); assert!((15000..=20000).contains(&spilled_rows)); - assert!((700000..=900000).contains(&spilled_bytes)); + assert!((900000..=1000000).contains(&spilled_bytes)); // Verify that the result is sorted let concated_result = concat_batches(&schema, &result)?; diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 454a06855175..68593fe6b05d 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -183,6 +183,21 @@ impl DisplayAs for SortPreservingMergeExec { write!(f, ", fetch={fetch}")?; }; + Ok(()) + } + DisplayFormatType::TreeRender => { + for (i, e) in self.expr().iter().enumerate() { + let e = e.to_string(); + if i == self.expr().len() - 1 { + writeln!(f, "{e}")?; + } else { + write!(f, "{e}, ")?; + } + } + if let Some(fetch) = self.fetch { + writeln!(f, "limit={fetch}")?; + }; + Ok(()) } } @@ -1383,6 +1398,10 @@ mod tests { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "CongestedExec",).unwrap() } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "").unwrap() + } } Ok(()) } diff --git a/datafusion/physical-plan/src/spill.rs b/datafusion/physical-plan/src/spill.rs index b45353ae13f0..fa1b8a91cec7 100644 --- a/datafusion/physical-plan/src/spill.rs +++ b/datafusion/physical-plan/src/spill.rs @@ -23,8 +23,8 @@ use std::path::{Path, PathBuf}; use std::ptr::NonNull; use arrow::array::ArrayData; -use arrow::datatypes::SchemaRef; -use arrow::ipc::reader::FileReader; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::ipc::{reader::StreamReader, writer::StreamWriter}; use arrow::record_batch::RecordBatch; use log::debug; use tokio::sync::mpsc::Sender; @@ -34,7 +34,6 @@ use datafusion_execution::disk_manager::RefCountedTempFile; use datafusion_execution::memory_pool::human_readable_size; use datafusion_execution::SendableRecordBatchStream; -use crate::common::IPCWriter; use crate::stream::RecordBatchReceiverStream; /// Read spilled batches from the disk @@ -59,13 +58,13 @@ pub(crate) fn read_spill_as_stream( /// /// Returns total number of the rows spilled to disk. pub(crate) fn spill_record_batches( - batches: Vec, + batches: &[RecordBatch], path: PathBuf, schema: SchemaRef, ) -> Result<(usize, usize)> { - let mut writer = IPCWriter::new(path.as_ref(), schema.as_ref())?; + let mut writer = IPCStreamWriter::new(path.as_ref(), schema.as_ref())?; for batch in batches { - writer.write(&batch)?; + writer.write(batch)?; } writer.finish()?; debug!( @@ -79,7 +78,7 @@ pub(crate) fn spill_record_batches( fn read_spill(sender: Sender>, path: &Path) -> Result<()> { let file = BufReader::new(File::open(path)?); - let reader = FileReader::try_new(file, None)?; + let reader = StreamReader::try_new(file, None)?; for batch in reader { sender .blocking_send(batch.map_err(Into::into)) @@ -98,7 +97,7 @@ pub fn spill_record_batch_by_size( ) -> Result<()> { let mut offset = 0; let total_rows = batch.num_rows(); - let mut writer = IPCWriter::new(&path, schema.as_ref())?; + let mut writer = IPCStreamWriter::new(&path, schema.as_ref())?; while offset < total_rows { let length = std::cmp::min(total_rows - offset, batch_size_rows); @@ -130,7 +129,7 @@ pub fn spill_record_batch_by_size( /// {xxxxxxxxxxxxxxxxxxx} <--- buffer /// ^ ^ ^ ^ /// | | | | -/// col1->{ } | | +/// col1->{ } | | /// col2--------->{ } /// /// In the above case, `get_record_batch_memory_size` will return the size of @@ -179,17 +178,64 @@ fn count_array_data_memory_size( } } +/// Write in Arrow IPC Stream format to a file. +/// +/// Stream format is used for spill because it supports dictionary replacement, and the random +/// access of IPC File format is not needed (IPC File format doesn't support dictionary replacement). +struct IPCStreamWriter { + /// Inner writer + pub writer: StreamWriter, + /// Batches written + pub num_batches: usize, + /// Rows written + pub num_rows: usize, + /// Bytes written + pub num_bytes: usize, +} + +impl IPCStreamWriter { + /// Create new writer + pub fn new(path: &Path, schema: &Schema) -> Result { + let file = File::create(path).map_err(|e| { + exec_datafusion_err!("Failed to create partition file at {path:?}: {e:?}") + })?; + Ok(Self { + num_batches: 0, + num_rows: 0, + num_bytes: 0, + writer: StreamWriter::try_new(file, schema)?, + }) + } + + /// Write one single batch + pub fn write(&mut self, batch: &RecordBatch) -> Result<()> { + self.writer.write(batch)?; + self.num_batches += 1; + self.num_rows += batch.num_rows(); + let num_bytes: usize = batch.get_array_memory_size(); + self.num_bytes += num_bytes; + Ok(()) + } + + /// Finish the writer + pub fn finish(&mut self) -> Result<()> { + self.writer.finish().map_err(Into::into) + } +} + #[cfg(test)] mod tests { use super::*; use crate::spill::{spill_record_batch_by_size, spill_record_batches}; use crate::test::build_table_i32; use arrow::array::{Float64Array, Int32Array, ListArray}; + use arrow::compute::cast; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_execution::disk_manager::DiskManagerConfig; use datafusion_execution::DiskManager; + use itertools::Itertools; use std::fs::File; use std::io::BufReader; use std::sync::Arc; @@ -214,18 +260,85 @@ mod tests { let schema = batch1.schema(); let num_rows = batch1.num_rows() + batch2.num_rows(); let (spilled_rows, _) = spill_record_batches( - vec![batch1, batch2], + &[batch1, batch2], spill_file.path().into(), Arc::clone(&schema), )?; assert_eq!(spilled_rows, num_rows); let file = BufReader::new(File::open(spill_file.path())?); - let reader = FileReader::try_new(file, None)?; + let reader = StreamReader::try_new(file, None)?; - assert_eq!(reader.num_batches(), 2); assert_eq!(reader.schema(), schema); + let batches = reader.collect_vec(); + assert!(batches.len() == 2); + + Ok(()) + } + + #[test] + fn test_batch_spill_and_read_dictionary_arrays() -> Result<()> { + // See https://github.com/apache/datafusion/issues/4658 + + let batch1 = build_table_i32( + ("a2", &vec![0, 1, 2]), + ("b2", &vec![3, 4, 5]), + ("c2", &vec![4, 5, 6]), + ); + + let batch2 = build_table_i32( + ("a2", &vec![10, 11, 12]), + ("b2", &vec![13, 14, 15]), + ("c2", &vec![14, 15, 16]), + ); + + // Dictionary encode the arrays + let dict_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)); + let dict_schema = Arc::new(Schema::new(vec![ + Field::new("a2", dict_type.clone(), true), + Field::new("b2", dict_type.clone(), true), + Field::new("c2", dict_type.clone(), true), + ])); + + let batch1 = RecordBatch::try_new( + Arc::clone(&dict_schema), + batch1 + .columns() + .iter() + .map(|array| cast(array, &dict_type)) + .collect::>()?, + )?; + + let batch2 = RecordBatch::try_new( + Arc::clone(&dict_schema), + batch2 + .columns() + .iter() + .map(|array| cast(array, &dict_type)) + .collect::>()?, + )?; + + let disk_manager = DiskManager::try_new(DiskManagerConfig::NewOs)?; + + let spill_file = disk_manager.create_tmp_file("Test Spill")?; + let num_rows = batch1.num_rows() + batch2.num_rows(); + let (spilled_rows, _) = spill_record_batches( + &[batch1, batch2], + spill_file.path().into(), + Arc::clone(&dict_schema), + )?; + assert_eq!(spilled_rows, num_rows); + + let file = BufReader::new(File::open(spill_file.path())?); + let reader = StreamReader::try_new(file, None)?; + + assert_eq!(reader.schema(), dict_schema); + + let batches = reader.collect_vec(); + assert!(batches.len() == 2); + Ok(()) } @@ -249,11 +362,13 @@ mod tests { )?; let file = BufReader::new(File::open(spill_file.path())?); - let reader = FileReader::try_new(file, None)?; + let reader = StreamReader::try_new(file, None)?; - assert_eq!(reader.num_batches(), 4); assert_eq!(reader.schema(), schema); + let batches = reader.collect_vec(); + assert!(batches.len() == 4); + Ok(()) } diff --git a/datafusion/physical-plan/src/stream.rs b/datafusion/physical-plan/src/stream.rs index 23cbb1ce49c1..6b5d8f46fe55 100644 --- a/datafusion/physical-plan/src/stream.rs +++ b/datafusion/physical-plan/src/stream.rs @@ -27,7 +27,7 @@ use super::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; use crate::displayable; use arrow::{datatypes::SchemaRef, record_batch::RecordBatch}; -use datafusion_common::{internal_err, Result}; +use datafusion_common::{exec_err, Result}; use datafusion_execution::TaskContext; use futures::stream::BoxStream; @@ -128,7 +128,7 @@ impl ReceiverStreamBuilder { // the JoinSet were aborted, which in turn // would imply that the receiver has been // dropped and this code is not running - return Some(internal_err!("Non Panic Task error: {e}")); + return Some(exec_err!("Non Panic Task error: {e}")); } } } @@ -223,6 +223,10 @@ impl RecordBatchReceiverStreamBuilder { } /// Get a handle for sending [`RecordBatch`] to the output + /// + /// If the stream is dropped / canceled, the sender will be closed and + /// calling `tx().send()` will return an error. Producers should stop + /// producing in this case and return control. pub fn tx(&self) -> Sender> { self.inner.tx() } @@ -241,8 +245,21 @@ impl RecordBatchReceiverStreamBuilder { self.inner.spawn(task) } - /// Spawn a blocking task that will be aborted if this builder (or the stream - /// built from it) are dropped + /// Spawn a blocking task tied to the builder and stream. + /// + /// # Drop / Cancel Behavior + /// + /// If this builder (or the stream built from it) is dropped **before** the + /// task starts, the task is also dropped and will never start execute. + /// + /// **Note:** Once the blocking task has started, it **will not** be + /// forcibly stopped on drop as Rust does not allow forcing a running thread + /// to terminate. The task will continue running until it completes or + /// encounters an error. + /// + /// Users should ensure that their blocking function periodically checks for + /// errors calling `tx.blocking_send`. An error signals that the stream has + /// been dropped / cancelled and the blocking task should exit. /// /// This is often used to spawn tasks that write to the sender /// retrieved from [`Self::tx`], for examples, see the document diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index 8bdfca2a8907..18c472a7e187 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -206,6 +206,18 @@ impl DisplayAs for StreamingTableExec { display_orderings(f, &self.projected_output_ordering)?; + Ok(()) + } + DisplayFormatType::TreeRender => { + if self.infinite { + writeln!(f, "infinite={}", self.infinite)?; + } + if let Some(limit) = self.limit { + write!(f, "limit={limit}")?; + } else { + write!(f, "limit=None")?; + } + Ok(()) } } diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index 7d0e3778452f..a2dc1d778436 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -105,10 +105,10 @@ impl DisplayAs for TestMemoryExec { .map_or(String::new(), |limit| format!(", fetch={}", limit)); if self.show_sizes { write!( - f, - "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", - partition_sizes.len(), - ) + f, + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) } else { write!( f, @@ -117,6 +117,10 @@ impl DisplayAs for TestMemoryExec { ) } } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index f0149faa8433..d0a0d25779cc 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -175,6 +175,10 @@ impl DisplayAs for MockExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "MockExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -337,6 +341,10 @@ impl DisplayAs for BarrierExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "BarrierExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -449,6 +457,10 @@ impl DisplayAs for ErrorExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "ErrorExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -535,6 +547,10 @@ impl DisplayAs for StatisticsExec { self.stats.num_rows, ) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -630,6 +646,10 @@ impl DisplayAs for BlockingExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "BlockingExec",) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -772,6 +792,10 @@ impl DisplayAs for PanicExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "PanicExec",) } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 68d1803b7133..791370917523 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -157,6 +157,10 @@ impl DisplayAs for UnionExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "UnionExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } @@ -387,6 +391,10 @@ impl DisplayAs for InterleaveExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "InterleaveExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 430391de5922..4e70d2dc4ee5 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -137,6 +137,9 @@ impl DisplayAs for UnnestExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "UnnestExec") } + DisplayFormatType::TreeRender => { + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/values.rs b/datafusion/physical-plan/src/values.rs index b90c50510cb0..6cb64bcb5d86 100644 --- a/datafusion/physical-plan/src/values.rs +++ b/datafusion/physical-plan/src/values.rs @@ -162,6 +162,10 @@ impl DisplayAs for ValuesExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "ValuesExec") } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } } } } diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 0d9c58b3bf49..f9f4b78686db 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -252,6 +252,17 @@ impl DisplayAs for BoundedWindowAggExec { let mode = &self.input_order_mode; write!(f, "wdw=[{}], mode=[{:?}]", g.join(", "), mode)?; } + DisplayFormatType::TreeRender => { + let g: Vec = self + .window_expr + .iter() + .map(|e| e.name().to_owned().to_string()) + .collect(); + writeln!(f, "select_list={}", g.join(", "))?; + + let mode = &self.input_order_mode; + writeln!(f, "mode={:?}", mode)?; + } } Ok(()) } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index d31fd66ca1f1..3c42d3032ed5 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -181,6 +181,14 @@ impl DisplayAs for WindowAggExec { .collect(); write!(f, "wdw=[{}]", g.join(", "))?; } + DisplayFormatType::TreeRender => { + let g: Vec = self + .window_expr + .iter() + .map(|e| e.name().to_owned().to_string()) + .collect(); + writeln!(f, "select_list={}", g.join(", "))?; + } } Ok(()) } diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index d3d29bfad7ce..126a7d0bba29 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -162,6 +162,9 @@ impl DisplayAs for WorkTableExec { DisplayFormatType::Default | DisplayFormatType::Verbose => { write!(f, "WorkTableExec: name={}", self.name) } + DisplayFormatType::TreeRender => { + write!(f, "name={}", self.name) + } } } } diff --git a/datafusion/proto-common/gen/src/main.rs b/datafusion/proto-common/gen/src/main.rs index 2cbe2afa5488..02e1ecf00bab 100644 --- a/datafusion/proto-common/gen/src/main.rs +++ b/datafusion/proto-common/gen/src/main.rs @@ -17,9 +17,6 @@ use std::path::Path; -type Error = Box; -type Result = std::result::Result; - fn main() -> Result<(), String> { let proto_dir = Path::new("proto"); let proto_path = Path::new("proto/datafusion_common.proto"); diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 00d4969182cf..39897cfcf2de 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -41,6 +41,7 @@ name = "datafusion_proto" default = ["parquet"] json = ["pbjson", "serde", "serde_json"] parquet = ["datafusion/parquet", "datafusion-common/parquet"] +avro = ["datafusion/avro", "datafusion-common/avro"] [dependencies] arrow = { workspace = true } diff --git a/datafusion/proto/gen/src/main.rs b/datafusion/proto/gen/src/main.rs index be61ff58fa8d..7f163162035c 100644 --- a/datafusion/proto/gen/src/main.rs +++ b/datafusion/proto/gen/src/main.rs @@ -17,9 +17,6 @@ use std::path::Path; -type Error = Box; -type Result = std::result::Result; - fn main() -> Result<(), String> { let proto_dir = Path::new("datafusion/proto"); let proto_path = Path::new("datafusion/proto/proto/datafusion.proto"); @@ -29,7 +26,6 @@ fn main() -> Result<(), String> { let descriptor_path = proto_dir.join("proto/proto_descriptor.bin"); prost_build::Config::new() - .protoc_arg("--experimental_allow_proto3_optional") .file_descriptor_set_path(&descriptor_path) .out_dir(out_dir) .compile_well_known_types() diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index e04a89a03dae..cac2f9db1645 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -527,6 +527,7 @@ pub fn parse_expr( ))), ExprType::Wildcard(protobuf::Wildcard { qualifier }) => { let qualifier = qualifier.to_owned().map(|x| x.try_into()).transpose()?; + #[expect(deprecated)] Ok(Expr::Wildcard { qualifier, options: Box::new(WildcardOptions::default()), diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 641dfe7b5fb8..148856cd103c 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -35,6 +35,8 @@ use crate::{ use crate::protobuf::{proto_error, ToProtoError}; use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion::datasource::cte_worktable::CteWorkTable; +#[cfg(feature = "avro")] +use datafusion::datasource::file_format::avro::AvroFormat; #[cfg(feature = "parquet")] use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::file_format::{ @@ -43,8 +45,7 @@ use datafusion::datasource::file_format::{ use datafusion::{ datasource::{ file_format::{ - avro::AvroFormat, csv::CsvFormat, json::JsonFormat as OtherNdJsonFormat, - FileFormat, + csv::CsvFormat, json::JsonFormat as OtherNdJsonFormat, FileFormat, }, listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl}, view::ViewTable, @@ -440,7 +441,15 @@ impl AsLogicalPlan for LogicalPlanNode { } Arc::new(json) } - FileFormatType::Avro(..) => Arc::new(AvroFormat), + #[cfg_attr(not(feature = "avro"), allow(unused_variables))] + FileFormatType::Avro(..) => { + #[cfg(feature = "avro")] + { + Arc::new(AvroFormat) + } + #[cfg(not(feature = "avro"))] + panic!("Unable to process avro file since `avro` feature is not enabled"); + } }; let table_paths = &scan @@ -1072,6 +1081,7 @@ impl AsLogicalPlan for LogicalPlanNode { })) } + #[cfg(feature = "avro")] if any.is::() { maybe_some_type = Some(FileFormatType::Avro(protobuf::AvroFormat {})) diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 5785bc0c4966..5bb0cdb20c9c 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -560,6 +560,7 @@ pub fn serialize_expr( expr_type: Some(ExprType::InList(expr)), } } + #[expect(deprecated)] Expr::Wildcard { qualifier, .. } => protobuf::LogicalExprNode { expr_type: Some(ExprType::Wildcard(protobuf::Wildcard { qualifier: qualifier.to_owned().map(|x| x.into()), diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index d0a31097b5cd..60972ac54ba7 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -29,9 +29,11 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT use datafusion::datasource::file_format::json::JsonSink; #[cfg(feature = "parquet")] use datafusion::datasource::file_format::parquet::ParquetSink; +#[cfg(feature = "avro")] +use datafusion::datasource::physical_plan::AvroSource; #[cfg(feature = "parquet")] use datafusion::datasource::physical_plan::ParquetSource; -use datafusion::datasource::physical_plan::{AvroSource, CsvSource, FileScanConfig}; +use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig}; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::FunctionRegistry; @@ -285,14 +287,20 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode { #[cfg(not(feature = "parquet"))] panic!("Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled") } + #[cfg_attr(not(feature = "avro"), allow(unused_variables))] PhysicalPlanType::AvroScan(scan) => { - let conf = parse_protobuf_file_scan_config( - scan.base_conf.as_ref().unwrap(), - registry, - extension_codec, - Arc::new(AvroSource::new()), - )?; - Ok(conf.build()) + #[cfg(feature = "avro")] + { + let conf = parse_protobuf_file_scan_config( + scan.base_conf.as_ref().unwrap(), + registry, + extension_codec, + Arc::new(AvroSource::new()), + )?; + Ok(conf.build()) + } + #[cfg(not(feature = "avro"))] + panic!("Unable to process a Avro PhysicalPlan when `avro` feature is not enabled") } PhysicalPlanType::CoalesceBatches(coalesce_batches) => { let input: Arc = into_physical_plan( @@ -1706,6 +1714,7 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode { } } + #[cfg(feature = "avro")] if let Some(data_source_exec) = plan.downcast_ref::() { let data_source = data_source_exec.data_source(); if let Some(maybe_avro) = diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index a8ee21365308..b5bfef99a6f3 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -41,7 +41,6 @@ use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema}; use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::file_format::csv::CsvSink; -use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::file_format::json::JsonSink; use datafusion::datasource::file_format::parquet::ParquetSink; use datafusion::datasource::listing::{ListingTableUrl, PartitionedFile}; @@ -95,7 +94,7 @@ use datafusion_common::file_options::json_writer::JsonWriterOptions; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::stats::Precision; use datafusion_common::{ - internal_err, not_impl_err, Constraints, DataFusionError, Result, UnnestOptions, + internal_err, not_impl_err, DataFusionError, Result, UnnestOptions, }; use datafusion_expr::{ Accumulator, AccumulatorFactoryFunction, AggregateUDF, ColumnarValue, ScalarUDF, @@ -738,33 +737,23 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { let mut options = TableParquetOptions::new(); options.global.pushdown_filters = true; - let source = Arc::new( + let file_source = Arc::new( ParquetSource::new(options).with_predicate(Arc::clone(&file_schema), predicate), ); - let scan_config = FileScanConfig { - object_store_url: ObjectStoreUrl::local_filesystem(), - file_schema, - file_groups: vec![vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )]], - constraints: Constraints::empty(), - statistics: Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&Arc::new(Schema::new(vec![ - Field::new("col", DataType::Utf8, false), - ]))), - }, - projection: None, - limit: None, - table_partition_cols: vec![], - output_ordering: vec![], - file_compression_type: FileCompressionType::UNCOMPRESSED, - new_lines_in_values: false, - file_source: source, - }; + let scan_config = + FileScanConfig::new(ObjectStoreUrl::local_filesystem(), file_schema, file_source) + .with_file_groups(vec![vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )]]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( + vec![Field::new("col", DataType::Utf8, false)], + ))), + }); roundtrip_test(scan_config.build()) } @@ -777,9 +766,9 @@ async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { vec![wrap_partition_value_in_dict(ScalarValue::Int64(Some(0)))]; let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - let source = Arc::new(ParquetSource::default()); + let file_source = Arc::new(ParquetSource::default()); let scan_config = - FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, source) + FileScanConfig::new(ObjectStoreUrl::local_filesystem(), schema, file_source) .with_projection(Some(vec![0, 1])) .with_file_group(vec![file_group]) .with_table_partition_cols(vec![Field::new( @@ -801,34 +790,24 @@ fn roundtrip_parquet_exec_with_custom_predicate_expr() -> Result<()> { inner: Arc::new(Column::new("col", 1)), }); - let source = Arc::new( + let file_source = Arc::new( ParquetSource::default() .with_predicate(Arc::clone(&file_schema), custom_predicate_expr), ); - let scan_config = FileScanConfig { - object_store_url: ObjectStoreUrl::local_filesystem(), - file_schema, - file_groups: vec![vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )]], - constraints: Constraints::empty(), - statistics: Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&Arc::new(Schema::new(vec![ - Field::new("col", DataType::Utf8, false), - ]))), - }, - projection: None, - limit: None, - table_partition_cols: vec![], - output_ordering: vec![], - file_compression_type: FileCompressionType::UNCOMPRESSED, - new_lines_in_values: false, - file_source: source, - }; + let scan_config = + FileScanConfig::new(ObjectStoreUrl::local_filesystem(), file_schema, file_source) + .with_file_groups(vec![vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )]]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( + vec![Field::new("col", DataType::Utf8, false)], + ))), + }); #[derive(Debug, Clone, Eq)] struct CustomPredicateExpr { @@ -1608,24 +1587,18 @@ async fn roundtrip_projection_source() -> Result<()> { let statistics = Statistics::new_unknown(&schema); - let source = ParquetSource::default().with_statistics(statistics.clone()); - let scan_config = FileScanConfig { - object_store_url: ObjectStoreUrl::local_filesystem(), - file_groups: vec![vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )]], - constraints: Constraints::empty(), - statistics, - file_schema: schema.clone(), - projection: Some(vec![0, 1, 2]), - limit: None, - table_partition_cols: vec![], - output_ordering: vec![], - file_compression_type: FileCompressionType::UNCOMPRESSED, - new_lines_in_values: false, - file_source: source, - }; + let file_source = ParquetSource::default().with_statistics(statistics.clone()); + let scan_config = FileScanConfig::new( + ObjectStoreUrl::local_filesystem(), + schema.clone(), + file_source, + ) + .with_file_groups(vec![vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )]]) + .with_statistics(statistics) + .with_projection(Some(vec![0, 1, 2])); let filter = Arc::new( FilterExec::try_new( diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index 596be3527b22..e7a04052ed51 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -592,10 +592,12 @@ impl SqlToRel<'_, S> { } not_impl_err!("AnyOp not supported by ExprPlanner: {binary_expr:?}") } + #[expect(deprecated)] SQLExpr::Wildcard(_token) => Ok(Expr::Wildcard { qualifier: None, options: Box::new(WildcardOptions::default()), }), + #[expect(deprecated)] SQLExpr::QualifiedWildcard(object_name, _token) => Ok(Expr::Wildcard { qualifier: Some(self.object_name_to_table_reference(object_name)?), options: Box::new(WildcardOptions::default()), diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs index 1e6e5621b8fe..a238bdb10184 100644 --- a/datafusion/sql/src/parser.rs +++ b/datafusion/sql/src/parser.rs @@ -269,33 +269,107 @@ pub struct DFParser<'a> { pub parser: Parser<'a>, } -impl<'a> DFParser<'a> { - /// Create a new parser for the specified tokens using the +/// Same as `sqlparser` +const DEFAULT_RECURSION_LIMIT: usize = 50; +const DEFAULT_DIALECT: GenericDialect = GenericDialect {}; + +/// Builder for [`DFParser`] +/// +/// # Example: Create and Parse SQL statements +/// ``` +/// # use datafusion_sql::parser::DFParserBuilder; +/// # use datafusion_common::Result; +/// # fn test() -> Result<()> { +/// let mut parser = DFParserBuilder::new("SELECT * FROM foo; SELECT 1 + 2") +/// .build()?; +/// // parse the SQL into DFStatements +/// let statements = parser.parse_statements()?; +/// assert_eq!(statements.len(), 2); +/// # Ok(()) +/// # } +/// ``` +/// +/// # Example: Create and Parse expression with a different dialect +/// ``` +/// # use datafusion_sql::parser::DFParserBuilder; +/// # use datafusion_common::Result; +/// # use datafusion_sql::sqlparser::dialect::MySqlDialect; +/// # use datafusion_sql::sqlparser::ast::Expr; +/// # fn test() -> Result<()> { +/// let dialect = MySqlDialect{}; // Parse using MySQL dialect +/// let mut parser = DFParserBuilder::new("1 + 2") +/// .with_dialect(&dialect) +/// .build()?; +/// // parse 1+2 into an sqlparser::ast::Expr +/// let res = parser.parse_expr()?; +/// assert!(matches!(res.expr, Expr::BinaryOp {..})); +/// # Ok(()) +/// # } +/// ``` +pub struct DFParserBuilder<'a> { + /// The SQL string to parse + sql: &'a str, + /// The Dialect to use (defaults to [`GenericDialect`] + dialect: &'a dyn Dialect, + /// The recursion limit while parsing + recursion_limit: usize, +} + +impl<'a> DFParserBuilder<'a> { + /// Create a new parser builder for the specified tokens using the /// [`GenericDialect`]. - pub fn new(sql: &str) -> Result { - let dialect = &GenericDialect {}; - DFParser::new_with_dialect(sql, dialect) + pub fn new(sql: &'a str) -> Self { + Self { + sql, + dialect: &DEFAULT_DIALECT, + recursion_limit: DEFAULT_RECURSION_LIMIT, + } } - /// Create a new parser for the specified tokens with the - /// specified dialect. - pub fn new_with_dialect( - sql: &str, - dialect: &'a dyn Dialect, - ) -> Result { - let mut tokenizer = Tokenizer::new(dialect, sql); + /// Adjust the parser builder's dialect. Defaults to [`GenericDialect`] + pub fn with_dialect(mut self, dialect: &'a dyn Dialect) -> Self { + self.dialect = dialect; + self + } + + /// Adjust the recursion limit of sql parsing. Defaults to 50 + pub fn with_recursion_limit(mut self, recursion_limit: usize) -> Self { + self.recursion_limit = recursion_limit; + self + } + + pub fn build(self) -> Result, ParserError> { + let mut tokenizer = Tokenizer::new(self.dialect, self.sql); let tokens = tokenizer.tokenize_with_location()?; Ok(DFParser { - parser: Parser::new(dialect).with_tokens_with_locations(tokens), + parser: Parser::new(self.dialect) + .with_tokens_with_locations(tokens) + .with_recursion_limit(self.recursion_limit), }) } +} + +impl<'a> DFParser<'a> { + #[deprecated(since = "46.0.0", note = "DFParserBuilder")] + pub fn new(sql: &'a str) -> Result { + DFParserBuilder::new(sql).build() + } + + #[deprecated(since = "46.0.0", note = "DFParserBuilder")] + pub fn new_with_dialect( + sql: &'a str, + dialect: &'a dyn Dialect, + ) -> Result { + DFParserBuilder::new(sql).with_dialect(dialect).build() + } /// Parse a sql string into one or [`Statement`]s using the /// [`GenericDialect`]. - pub fn parse_sql(sql: &str) -> Result, ParserError> { - let dialect = &GenericDialect {}; - DFParser::parse_sql_with_dialect(sql, dialect) + pub fn parse_sql(sql: &'a str) -> Result, ParserError> { + let mut parser = DFParserBuilder::new(sql).build()?; + + parser.parse_statements() } /// Parse a SQL string and produce one or more [`Statement`]s with @@ -304,37 +378,43 @@ impl<'a> DFParser<'a> { sql: &str, dialect: &dyn Dialect, ) -> Result, ParserError> { - let mut parser = DFParser::new_with_dialect(sql, dialect)?; + let mut parser = DFParserBuilder::new(sql).with_dialect(dialect).build()?; + parser.parse_statements() + } + + pub fn parse_sql_into_expr_with_dialect( + sql: &str, + dialect: &dyn Dialect, + ) -> Result { + let mut parser = DFParserBuilder::new(sql).with_dialect(dialect).build()?; + + parser.parse_expr() + } + + /// Parse a sql string into one or [`Statement`]s + pub fn parse_statements(&mut self) -> Result, ParserError> { let mut stmts = VecDeque::new(); let mut expecting_statement_delimiter = false; loop { // ignore empty statements (between successive statement delimiters) - while parser.parser.consume_token(&Token::SemiColon) { + while self.parser.consume_token(&Token::SemiColon) { expecting_statement_delimiter = false; } - if parser.parser.peek_token() == Token::EOF { + if self.parser.peek_token() == Token::EOF { break; } if expecting_statement_delimiter { - return parser.expected("end of statement", parser.parser.peek_token()); + return self.expected("end of statement", self.parser.peek_token()); } - let statement = parser.parse_statement()?; + let statement = self.parse_statement()?; stmts.push_back(statement); expecting_statement_delimiter = true; } Ok(stmts) } - pub fn parse_sql_into_expr_with_dialect( - sql: &str, - dialect: &dyn Dialect, - ) -> Result { - let mut parser = DFParser::new_with_dialect(sql, dialect)?; - parser.parse_expr() - } - /// Report an unexpected token fn expected( &self, @@ -876,6 +956,7 @@ impl<'a> DFParser<'a> { #[cfg(test)] mod tests { use super::*; + use datafusion_common::assert_contains; use sqlparser::ast::Expr::Identifier; use sqlparser::ast::{BinaryOperator, DataType, Expr, Ident}; use sqlparser::dialect::SnowflakeDialect; @@ -1612,4 +1693,30 @@ mod tests { fn verified_stmt(sql: &str) -> Statement { one_statement_parses_to(sql, sql) } + + #[test] + /// Checks the recursion limit works for sql queries + /// Recursion can happen easily with binary exprs (i.e, AND or OR) + fn test_recursion_limit() { + let sql = "SELECT 1 OR 2"; + + // Expect parse to succeed + DFParserBuilder::new(sql) + .build() + .unwrap() + .parse_statements() + .unwrap(); + + let err = DFParserBuilder::new(sql) + .with_recursion_limit(1) + .build() + .unwrap() + .parse_statements() + .unwrap_err(); + + assert_contains!( + err.to_string(), + "sql parser error: recursion limit exceeded" + ); + } } diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 7af754ed5625..d9ade929ad54 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -52,6 +52,8 @@ pub struct ParserOptions { pub enable_options_value_normalization: bool, /// Whether to collect spans pub collect_spans: bool, + /// Whether `VARCHAR` is mapped to `Utf8View` during SQL planning. + pub map_varchar_to_utf8view: bool, } impl ParserOptions { @@ -70,6 +72,7 @@ impl ParserOptions { parse_float_as_decimal: false, enable_ident_normalization: true, support_varchar_with_length: true, + map_varchar_to_utf8view: false, enable_options_value_normalization: false, collect_spans: false, } @@ -109,6 +112,12 @@ impl ParserOptions { self } + /// Sets the `map_varchar_to_utf8view` option. + pub fn with_map_varchar_to_utf8view(mut self, value: bool) -> Self { + self.map_varchar_to_utf8view = value; + self + } + /// Sets the `enable_options_value_normalization` option. pub fn with_enable_options_value_normalization(mut self, value: bool) -> Self { self.enable_options_value_normalization = value; @@ -134,6 +143,7 @@ impl From<&SqlParserOptions> for ParserOptions { parse_float_as_decimal: options.parse_float_as_decimal, enable_ident_normalization: options.enable_ident_normalization, support_varchar_with_length: options.support_varchar_with_length, + map_varchar_to_utf8view: options.map_varchar_to_utf8view, enable_options_value_normalization: options .enable_options_value_normalization, collect_spans: options.collect_spans, @@ -556,7 +566,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { SQLDataType::Varchar(length) => { match (length, self.options.support_varchar_with_length) { (Some(_), false) => plan_err!("does not support Varchar with length, please set `support_varchar_with_length` to be true"), - _ => Ok(DataType::Utf8), + _ => { + if self.options.map_varchar_to_utf8view { + Ok(DataType::Utf8View) + } else { + Ok(DataType::Utf8) + } + } } } SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 8385682d42bc..eee27cb3013c 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -28,19 +28,19 @@ use crate::utils::{ use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; -use datafusion_common::{not_impl_err, plan_err, Result}; +use datafusion_common::{not_impl_err, plan_err, Column, Result}; use datafusion_common::{RecursionUnnestOption, UnnestOptions}; use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts, }; use datafusion_expr::utils::{ - expr_as_column_expr, expr_to_columns, find_aggregate_exprs, find_window_exprs, + expand_qualified_wildcard, expand_wildcard, expr_as_column_expr, expr_to_columns, + find_aggregate_exprs, find_window_exprs, }; use datafusion_expr::{ - qualified_wildcard_with_options, wildcard_with_options, Aggregate, Expr, Filter, - GroupingSet, LogicalPlan, LogicalPlanBuilder, LogicalPlanBuilderOptions, - Partitioning, + Aggregate, Expr, Filter, GroupingSet, LogicalPlan, LogicalPlanBuilder, + LogicalPlanBuilderOptions, Partitioning, }; use indexmap::IndexMap; @@ -93,6 +93,24 @@ impl SqlToRel<'_, S> { planner_context, )?; + // TOOD: remove this after Expr::Wildcard is removed + #[allow(deprecated)] + for expr in &select_exprs { + debug_assert!(!matches!(expr, Expr::Wildcard { .. })); + } + + // TOOD: remove this after Expr::Wildcard is removed + #[allow(deprecated)] + for expr in &select_exprs { + debug_assert!(!matches!(expr, Expr::Wildcard { .. })); + } + + // TOOD: remove this after Expr::Wildcard is removed + #[allow(deprecated)] + for expr in &select_exprs { + debug_assert!(!matches!(expr, Expr::Wildcard { .. })); + } + let order_by = to_order_by_exprs_with_select(query_order_by, Some(select_exprs.clone()))?; @@ -587,7 +605,7 @@ impl SqlToRel<'_, S> { let mut error_builder = DataFusionErrorBuilder::new(); for expr in projection { match self.sql_select_to_rex(expr, plan, empty_from, planner_context) { - Ok(expr) => prepared_select_exprs.push(expr), + Ok(expr) => prepared_select_exprs.extend(expr), Err(err) => error_builder.add_error(err), } } @@ -601,7 +619,7 @@ impl SqlToRel<'_, S> { plan: &LogicalPlan, empty_from: bool, planner_context: &mut PlannerContext, - ) -> Result { + ) -> Result> { match sql { SelectItem::UnnamedExpr(expr) => { let expr = self.sql_to_expr(expr, plan.schema(), planner_context)?; @@ -610,7 +628,7 @@ impl SqlToRel<'_, S> { &[&[plan.schema()]], &plan.using_columns()?, )?; - Ok(col) + Ok(vec![col]) } SelectItem::ExprWithAlias { expr, alias } => { let select_expr = @@ -626,7 +644,7 @@ impl SqlToRel<'_, S> { Expr::Column(column) if column.name.eq(&name) => col, _ => col.alias(name), }; - Ok(expr) + Ok(vec![expr]) } SelectItem::Wildcard(options) => { Self::check_wildcard_options(&options)?; @@ -639,7 +657,17 @@ impl SqlToRel<'_, S> { planner_context, options, )?; - Ok(wildcard_with_options(planned_options)) + + let expanded = + expand_wildcard(plan.schema(), plan, Some(&planned_options))?; + + // If there is a REPLACE statement, replace that column with the given + // replace expression. Column name remains the same. + if let Some(replace) = planned_options.replace { + replace_columns(expanded, &replace) + } else { + Ok(expanded) + } } SelectItem::QualifiedWildcard(object_name, options) => { Self::check_wildcard_options(&options)?; @@ -660,7 +688,19 @@ impl SqlToRel<'_, S> { planner_context, options, )?; - Ok(qualified_wildcard_with_options(qualifier, planned_options)) + + let expanded = expand_qualified_wildcard( + &qualifier, + plan.schema(), + Some(&planned_options), + )?; + // If there is a REPLACE statement, replace that column with the given + // replace expression. Column name remains the same. + if let Some(replace) = planned_options.replace { + replace_columns(expanded, &replace) + } else { + Ok(expanded) + } } } } @@ -712,7 +752,10 @@ impl SqlToRel<'_, S> { planner_context, ) }) - .collect::>>()?; + .collect::>>()? + .into_iter() + .flatten() + .collect(); let planned_replace = PlannedReplaceSelectItem { items: replace.items.into_iter().map(|i| *i).collect(), planned_expressions: replace_expr, @@ -898,3 +941,26 @@ fn match_window_definitions( } Ok(()) } + +/// If there is a REPLACE statement in the projected expression in the form of +/// "REPLACE (some_column_within_an_expr AS some_column)", this function replaces +/// that column with the given replace expression. Column name remains the same. +/// Multiple REPLACEs are also possible with comma separations. +fn replace_columns( + mut exprs: Vec, + replace: &PlannedReplaceSelectItem, +) -> Result> { + for expr in exprs.iter_mut() { + if let Expr::Column(Column { name, .. }) = expr { + if let Some((_, new_expr)) = replace + .items() + .iter() + .zip(replace.expressions().iter()) + .find(|(item, _)| item.column_name.value == *name) + { + *expr = new_expr.clone().alias(name.clone()) + } + } + } + Ok(exprs) +} diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 21a78b5a569d..ba0668eb06c3 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -430,6 +430,7 @@ impl Unparser<'_> { }) } // TODO: unparsing wildcard addition options + #[expect(deprecated)] Expr::Wildcard { qualifier, .. } => { let attached_token = AttachedToken::empty(); if let Some(qualifier) = qualifier { @@ -740,6 +741,7 @@ impl Unparser<'_> { ) -> Result> { args.iter() .map(|e| { + #[expect(deprecated)] if matches!( e, Expr::Wildcard { @@ -957,8 +959,8 @@ impl Unparser<'_> { Operator::BitwiseShiftRight => Ok(BinaryOperator::PGBitwiseShiftRight), Operator::BitwiseShiftLeft => Ok(BinaryOperator::PGBitwiseShiftLeft), Operator::StringConcat => Ok(BinaryOperator::StringConcat), - Operator::AtArrow => not_impl_err!("unsupported operation: {op:?}"), - Operator::ArrowAt => not_impl_err!("unsupported operation: {op:?}"), + Operator::AtArrow => Ok(BinaryOperator::AtArrow), + Operator::ArrowAt => Ok(BinaryOperator::ArrowAt), } } @@ -1726,6 +1728,7 @@ mod tests { #[test] fn expr_to_sql_ok() -> Result<()> { let dummy_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + #[expect(deprecated)] let dummy_logical_plan = table_scan(Some("t"), &dummy_schema, None)? .project(vec![Expr::Wildcard { qualifier: None, @@ -2121,6 +2124,22 @@ mod tests { ))), "[1, 2, 3]", ), + ( + Expr::BinaryExpr(BinaryExpr { + left: Box::new(col("a")), + op: Operator::ArrowAt, + right: Box::new(col("b")), + }), + "(a <@ b)", + ), + ( + Expr::BinaryExpr(BinaryExpr { + left: Box::new(col("a")), + op: Operator::AtArrow, + right: Box::new(col("b")), + }), + "(a @> b)", + ), ]; for (expr, expected) in tests { diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 447d42d1794d..bec2b4990125 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -984,11 +984,18 @@ impl Unparser<'_> { Ok(Some(builder.build()?)) } LogicalPlan::SubqueryAlias(subquery_alias) => { - Self::unparse_table_scan_pushdown( + let ret = Self::unparse_table_scan_pushdown( &subquery_alias.input, Some(subquery_alias.alias.clone()), already_projected, - ) + )?; + if let Some(alias) = alias { + if let Some(plan) = ret { + let plan = LogicalPlanBuilder::new(plan).alias(alias)?.build()?; + return Ok(Some(plan)); + } + } + Ok(ret) } // SubqueryAlias could be rewritten to a plan with a projection as the top node by [rewrite::subquery_alias_inner_query_and_columns]. // The inner table scan could be a scan with pushdown operations. diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 3f093afaf26a..4a248de101dc 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -632,6 +632,8 @@ pub(crate) fn rewrite_recursive_unnest_bottom_up( } = original_expr.clone().rewrite(&mut rewriter)?; if !transformed { + // TODO: remove the next line after `Expr::Wildcard` is removed + #[expect(deprecated)] if matches!(&transformed_expr, Expr::Column(_)) || matches!(&transformed_expr, Expr::Wildcard { .. }) { diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 79f17a330bfe..3a84a8347512 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -21,7 +21,7 @@ use datafusion_expr::test::function_stub::{ count_udaf, max_udaf, min_udaf, sum, sum_udaf, }; use datafusion_expr::{ - col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, + cast, col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; use datafusion_functions::unicode; @@ -37,6 +37,7 @@ use datafusion_sql::unparser::dialect::{ use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser}; use sqlparser::ast::Statement; use std::hash::Hash; +use std::ops::Add; use std::sync::Arc; use std::{fmt, vec}; @@ -324,9 +325,9 @@ fn roundtrip_statement_with_dialect() -> Result<()> { unparser_dialect: Box::new(UnparserMySqlDialect {}), }, TestStatementWithDialect { - sql: "select * from (select * from j1 limit 10);", + sql: "select j1_id from (select j1_id from j1 limit 10);", expected: - "SELECT * FROM (SELECT * FROM `j1` LIMIT 10) AS `derived_limit`", + "SELECT `j1`.`j1_id` FROM (SELECT `j1`.`j1_id` FROM `j1` LIMIT 10) AS `derived_limit`", parser_dialect: Box::new(MySqlDialect {}), unparser_dialect: Box::new(UnparserMySqlDialect {}), }, @@ -526,85 +527,79 @@ fn roundtrip_statement_with_dialect() -> Result<()> { }, TestStatementWithDialect { sql: "SELECT * FROM (SELECT j1_id + 1 FROM j1) AS temp_j(id2)", - expected: r#"SELECT * FROM (SELECT (`j1`.`j1_id` + 1) AS `id2` FROM `j1`) AS `temp_j`"#, + expected: r#"SELECT `temp_j`.`id2` FROM (SELECT (`j1`.`j1_id` + 1) AS `id2` FROM `j1`) AS `temp_j`"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(SqliteDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM (SELECT j1_id FROM j1 LIMIT 1) AS temp_j(id2)", - expected: r#"SELECT * FROM (SELECT `j1`.`j1_id` AS `id2` FROM `j1` LIMIT 1) AS `temp_j`"#, + expected: r#"SELECT `temp_j`.`id2` FROM (SELECT `j1`.`j1_id` AS `id2` FROM `j1` LIMIT 1) AS `temp_j`"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(SqliteDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3])", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))")"#, + expected: r#"SELECT "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))" FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))")"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) AS t1 (c1)", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS t1 (c1)"#, - parser_dialect: Box::new(GenericDialect {}), - unparser_dialect: Box::new(UnparserDefaultDialect {}), - }, - TestStatementWithDialect { - sql: "SELECT * FROM UNNEST([1,2,3]) AS t1 (c1)", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS t1 (c1)"#, + expected: r#"SELECT t1.c1 FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS t1 (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]), j1", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") CROSS JOIN j1"#, + expected: r#"SELECT "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))", j1.j1_id, j1.j1_string FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") CROSS JOIN j1"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) JOIN j1 ON u.c1 = j1.j1_id", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, + expected: r#"SELECT u.c1, j1.j1_id, j1.j1_string FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) UNION ALL SELECT * FROM UNNEST([4,5,6]) u(c1)", - expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) UNION ALL SELECT * FROM (SELECT UNNEST([4, 5, 6]) AS "UNNEST(make_array(Int64(4),Int64(5),Int64(6)))") AS u (c1)"#, + expected: r#"SELECT u.c1 FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))") AS u (c1) UNION ALL SELECT u.c1 FROM (SELECT UNNEST([4, 5, 6]) AS "UNNEST(make_array(Int64(4),Int64(5),Int64(6)))") AS u (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3])", - expected: r#"SELECT * FROM UNNEST([1, 2, 3])"#, + expected: r#"SELECT UNNEST(make_array(Int64(1),Int64(2),Int64(3))) FROM UNNEST([1, 2, 3])"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) AS t1 (c1)", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS t1 (c1)"#, + expected: r#"SELECT t1.c1 FROM UNNEST([1, 2, 3]) AS t1 (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) AS t1 (c1)", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS t1 (c1)"#, + expected: r#"SELECT t1.c1 FROM UNNEST([1, 2, 3]) AS t1 (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]), j1", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) CROSS JOIN j1"#, + expected: r#"SELECT UNNEST(make_array(Int64(1),Int64(2),Int64(3))), j1.j1_id, j1.j1_string FROM UNNEST([1, 2, 3]) CROSS JOIN j1"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) JOIN j1 ON u.c1 = j1.j1_id", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, + expected: r#"SELECT u.c1, j1.j1_id, j1.j1_string FROM UNNEST([1, 2, 3]) AS u (c1) INNER JOIN j1 ON (u.c1 = j1.j1_id)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) UNION ALL SELECT * FROM UNNEST([4,5,6]) u(c1)", - expected: r#"SELECT * FROM UNNEST([1, 2, 3]) AS u (c1) UNION ALL SELECT * FROM UNNEST([4, 5, 6]) AS u (c1)"#, + expected: r#"SELECT u.c1 FROM UNNEST([1, 2, 3]) AS u (c1) UNION ALL SELECT u.c1 FROM UNNEST([4, 5, 6]) AS u (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, @@ -628,25 +623,25 @@ fn roundtrip_statement_with_dialect() -> Result<()> { }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col)", - expected: r#"SELECT * FROM unnest_table AS u CROSS JOIN UNNEST(u.array_col)"#, + expected: r#"SELECT u.array_col, u.struct_col, UNNEST(outer_ref(u.array_col)) FROM unnest_table AS u CROSS JOIN UNNEST(u.array_col)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col) AS t1 (c1)", - expected: r#"SELECT * FROM unnest_table AS u CROSS JOIN UNNEST(u.array_col) AS t1 (c1)"#, + expected: r#"SELECT u.array_col, u.struct_col, t1.c1 FROM unnest_table AS u CROSS JOIN UNNEST(u.array_col) AS t1 (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(CustomDialectBuilder::default().with_unnest_as_table_factor(true).build()), }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col)", - expected: r#"SELECT * FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))")"#, + expected: r#"SELECT u.array_col, u.struct_col, "UNNEST(outer_ref(u.array_col))" FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))")"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, TestStatementWithDialect { sql: "SELECT * FROM unnest_table u, UNNEST(u.array_col) AS t1 (c1)", - expected: r#"SELECT * FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))") AS t1 (c1)"#, + expected: r#"SELECT u.array_col, u.struct_col, t1.c1 FROM unnest_table AS u CROSS JOIN LATERAL (SELECT UNNEST(u.array_col) AS "UNNEST(outer_ref(u.array_col))") AS t1 (c1)"#, parser_dialect: Box::new(GenericDialect {}), unparser_dialect: Box::new(UnparserDefaultDialect {}), }, @@ -1456,13 +1451,13 @@ fn test_unnest_to_sql() { fn test_join_with_no_conditions() { sql_round_trip( GenericDialect {}, - "SELECT * FROM j1 JOIN j2", - "SELECT * FROM j1 CROSS JOIN j2", + "SELECT j1.j1_id, j1.j1_string FROM j1 JOIN j2", + "SELECT j1.j1_id, j1.j1_string FROM j1 CROSS JOIN j2", ); sql_round_trip( GenericDialect {}, - "SELECT * FROM j1 CROSS JOIN j2", - "SELECT * FROM j1 CROSS JOIN j2", + "SELECT j1.j1_id, j1.j1_string FROM j1 CROSS JOIN j2", + "SELECT j1.j1_id, j1.j1_string FROM j1 CROSS JOIN j2", ); } @@ -1563,7 +1558,7 @@ fn test_unparse_extension_to_statement() -> Result<()> { Arc::new(UnusedUnparser {}), ]); let sql = unparser.plan_to_sql(&extension)?; - let expected = "SELECT * FROM j1"; + let expected = "SELECT j1.j1_id, j1.j1_string FROM j1"; assert_eq!(sql.to_string(), expected); if let Some(err) = plan_to_sql(&extension).err() { @@ -1626,7 +1621,8 @@ fn test_unparse_extension_to_sql() -> Result<()> { Arc::new(UnusedUnparser {}), ]); let sql = unparser.plan_to_sql(&plan)?; - let expected = "SELECT j1.j1_id AS user_id FROM (SELECT * FROM j1)"; + let expected = + "SELECT j1.j1_id AS user_id FROM (SELECT j1.j1_id, j1.j1_string FROM j1)"; assert_eq!(sql.to_string(), expected); if let Some(err) = plan_to_sql(&plan).err() { @@ -1687,3 +1683,66 @@ fn test_unparse_optimized_multi_union() -> Result<()> { Ok(()) } + +/// Test unparse the optimized plan from the following SQL: +/// ``` +/// SELECT +/// customer_view.c_custkey, +/// customer_view.c_name, +/// customer_view.custkey_plus +/// FROM +/// ( +/// SELECT +/// customer.c_custkey, +/// customer.c_name, +/// customer.custkey_plus +/// FROM +/// ( +/// SELECT +/// customer.c_custkey, +/// CAST(customer.c_custkey AS BIGINT) + 1 AS custkey_plus, +/// customer.c_name +/// FROM +/// ( +/// SELECT +/// customer.c_custkey AS c_custkey, +/// customer.c_name AS c_name +/// FROM +/// customer +/// ) AS customer +/// ) AS customer +/// ) AS customer_view +/// ``` +#[test] +fn test_unparse_subquery_alias_with_table_pushdown() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c_custkey", DataType::Int32, false), + Field::new("c_name", DataType::Utf8, false), + ]); + + let table_scan = table_scan(Some("customer"), &schema, Some(vec![0, 1]))?.build()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .alias("customer")? + .project(vec![ + col("customer.c_custkey"), + cast(col("customer.c_custkey"), DataType::Int64) + .add(lit(1)) + .alias("custkey_plus"), + col("customer.c_name"), + ])? + .alias("customer")? + .project(vec![ + col("customer.c_custkey"), + col("customer.c_name"), + col("customer.custkey_plus"), + ])? + .alias("customer_view")? + .build()?; + + let unparser = Unparser::default(); + let sql = unparser.plan_to_sql(&plan)?; + let expected = "SELECT customer_view.c_custkey, customer_view.c_name, customer_view.custkey_plus FROM (SELECT customer.c_custkey, (CAST(customer.c_custkey AS BIGINT) + 1) AS custkey_plus, customer.c_name FROM (SELECT customer.c_custkey, customer.c_name FROM customer AS customer) AS customer) AS customer_view"; + assert_eq!(sql.to_string(), expected); + Ok(()) +} diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 1df18302687e..023ea88cb55f 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -54,15 +54,6 @@ use sqlparser::dialect::{Dialect, GenericDialect, HiveDialect, MySqlDialect}; mod cases; mod common; -#[test] -fn test_schema_support() { - quick_test( - "SELECT * FROM s1.test", - "Projection: *\ - \n TableScan: s1.test", - ); -} - #[test] fn parse_decimals() { let test_data = [ @@ -92,6 +83,7 @@ fn parse_decimals() { parse_float_as_decimal: true, enable_ident_normalization: false, support_varchar_with_length: false, + map_varchar_to_utf8view: false, enable_options_value_normalization: false, collect_spans: false, }, @@ -148,6 +140,7 @@ fn parse_ident_normalization() { parse_float_as_decimal: false, enable_ident_normalization, support_varchar_with_length: false, + map_varchar_to_utf8view: false, enable_options_value_normalization: false, collect_spans: false, }, @@ -449,19 +442,6 @@ Explain quick_test(sql, plan); } -#[test] -fn plan_copy_to_query() { - let sql = "COPY (select * from test_decimal limit 10) to 'output.csv'"; - let plan = r#" -CopyTo: format=csv output_url=output.csv options: () - Limit: skip=0, fetch=10 - Projection: * - TableScan: test_decimal - "# - .trim(); - quick_test(sql, plan); -} - #[test] fn plan_insert() { let sql = @@ -585,15 +565,6 @@ fn select_repeated_column() { ); } -#[test] -fn select_wildcard_with_repeated_column_but_is_aliased() { - quick_test( - "SELECT *, first_name AS fn from person", - "Projection: *, person.first_name AS fn\ - \n TableScan: person", - ); -} - #[test] fn select_scalar_func_with_literal_no_relation() { quick_test( @@ -793,30 +764,6 @@ fn join_with_ambiguous_column() { quick_test(sql, expected); } -#[test] -fn where_selection_with_ambiguous_column() { - let sql = "SELECT * FROM person a, person b WHERE id = id + 1"; - let err = logical_plan(sql) - .expect_err("query should have failed") - .strip_backtrace(); - assert_eq!( - "\"Schema error: Ambiguous reference to unqualified field id\"", - format!("{err:?}") - ); -} - -#[test] -fn natural_join() { - let sql = "SELECT * FROM lineitem a NATURAL JOIN lineitem b"; - let expected = "Projection: *\ - \n Inner Join: Using a.l_item_id = b.l_item_id, a.l_description = b.l_description, a.price = b.price\ - \n SubqueryAlias: a\ - \n TableScan: lineitem\ - \n SubqueryAlias: b\ - \n TableScan: lineitem"; - quick_test(sql, expected); -} - #[test] fn natural_left_join() { let sql = "SELECT l_item_id FROM lineitem a NATURAL LEFT JOIN lineitem b"; @@ -841,83 +788,6 @@ fn natural_right_join() { quick_test(sql, expected); } -#[test] -fn natural_join_no_common_becomes_cross_join() { - let sql = "SELECT * FROM person a NATURAL JOIN lineitem b"; - let expected = "Projection: *\ - \n Cross Join: \ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: lineitem"; - quick_test(sql, expected); -} - -#[test] -fn using_join_multiple_keys() { - let sql = "SELECT * FROM person a join person b using (id, age)"; - let expected = "Projection: *\ - \n Inner Join: Using a.id = b.id, a.age = b.age\ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: person"; - quick_test(sql, expected); -} - -#[test] -fn using_join_multiple_keys_subquery() { - let sql = - "SELECT age FROM (SELECT * FROM person a join person b using (id, age, state))"; - let expected = "Projection: a.age\ - \n Projection: *\ - \n Inner Join: Using a.id = b.id, a.age = b.age, a.state = b.state\ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: person"; - quick_test(sql, expected); -} - -#[test] -fn using_join_multiple_keys_qualified_wildcard_select() { - let sql = "SELECT a.* FROM person a join person b using (id, age)"; - let expected = "Projection: a.*\ - \n Inner Join: Using a.id = b.id, a.age = b.age\ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: person"; - quick_test(sql, expected); -} - -#[test] -fn using_join_multiple_keys_select_all_columns() { - let sql = "SELECT a.*, b.* FROM person a join person b using (id, age)"; - let expected = "Projection: a.*, b.*\ - \n Inner Join: Using a.id = b.id, a.age = b.age\ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: person"; - quick_test(sql, expected); -} - -#[test] -fn using_join_multiple_keys_multiple_joins() { - let sql = "SELECT * FROM person a join person b using (id, age, state) join person c using (id, age, state)"; - let expected = "Projection: *\ - \n Inner Join: Using a.id = c.id, a.age = c.age, a.state = c.state\ - \n Inner Join: Using a.id = b.id, a.age = b.age, a.state = b.state\ - \n SubqueryAlias: a\ - \n TableScan: person\ - \n SubqueryAlias: b\ - \n TableScan: person\ - \n SubqueryAlias: c\ - \n TableScan: person"; - quick_test(sql, expected); -} - #[test] fn select_with_having() { let sql = "SELECT id, age @@ -1233,24 +1103,6 @@ fn select_binary_expr_nested() { quick_test(sql, expected); } -#[test] -fn select_wildcard_with_groupby() { - quick_test( - r#"SELECT * FROM person GROUP BY id, first_name, last_name, age, state, salary, birth_date, "😀""#, - "Projection: *\ - \n Aggregate: groupBy=[[person.id, person.first_name, person.last_name, person.age, person.state, person.salary, person.birth_date, person.😀]], aggr=[[]]\ - \n TableScan: person", - ); - quick_test( - "SELECT * FROM (SELECT first_name, last_name FROM person) AS a GROUP BY first_name, last_name", - "Projection: *\ - \n Aggregate: groupBy=[[a.first_name, a.last_name]], aggr=[[]]\ - \n SubqueryAlias: a\ - \n Projection: person.first_name, person.last_name\ - \n TableScan: person", - ); -} - #[test] fn select_simple_aggregate() { quick_test( @@ -1397,56 +1249,6 @@ fn select_interval_out_of_range() { ); } -#[test] -fn recursive_ctes() { - let sql = " - WITH RECURSIVE numbers AS ( - select 1 as n - UNION ALL - select n + 1 FROM numbers WHERE N < 10 - ) - select * from numbers;"; - quick_test( - sql, - "Projection: *\ - \n SubqueryAlias: numbers\ - \n RecursiveQuery: is_distinct=false\ - \n Projection: Int64(1) AS n\ - \n EmptyRelation\ - \n Projection: numbers.n + Int64(1)\ - \n Filter: numbers.n < Int64(10)\ - \n TableScan: numbers", - ) -} - -#[test] -fn recursive_ctes_disabled() { - let sql = " - WITH RECURSIVE numbers AS ( - select 1 as n - UNION ALL - select n + 1 FROM numbers WHERE N < 10 - ) - select * from numbers;"; - - // manually setting up test here so that we can disable recursive ctes - let mut state = MockSessionState::default(); - state.config_options.execution.enable_recursive_ctes = false; - let context = MockContextProvider { state }; - - let planner = SqlToRel::new_with_options(&context, ParserOptions::default()); - let result = DFParser::parse_sql_with_dialect(sql, &GenericDialect {}); - let mut ast = result.unwrap(); - - let err = planner - .statement_to_plan(ast.pop_front().unwrap()) - .expect_err("query should have failed"); - assert_eq!( - "This feature is not implemented: Recursive CTEs are not enabled", - err.strip_backtrace() - ); -} - #[test] fn select_simple_aggregate_with_groupby_and_column_is_in_aggregate_and_groupby() { quick_test( @@ -1618,15 +1420,6 @@ fn select_aggregate_with_non_column_inner_expression_with_groupby() { ); } -#[test] -fn test_wildcard() { - quick_test( - "SELECT * from person", - "Projection: *\ - \n TableScan: person", - ); -} - #[test] fn select_count_one() { let sql = "SELECT count(1) FROM person"; @@ -2060,20 +1853,6 @@ fn join_with_using() { quick_test(sql, expected); } -#[test] -fn project_wildcard_on_join_with_using() { - let sql = "SELECT * \ - FROM lineitem \ - JOIN lineitem as lineitem2 \ - USING (l_item_id)"; - let expected = "Projection: *\ - \n Inner Join: Using lineitem.l_item_id = lineitem2.l_item_id\ - \n TableScan: lineitem\ - \n SubqueryAlias: lineitem2\ - \n TableScan: lineitem"; - quick_test(sql, expected); -} - #[test] fn equijoin_explicit_syntax_3_tables() { let sql = "SELECT id, order_id, l_description \ @@ -2867,24 +2646,6 @@ fn exists_subquery_schema_outer_schema_overlap() { quick_test(sql, expected); } -#[test] -fn exists_subquery_wildcard() { - let sql = "SELECT id FROM person p WHERE EXISTS \ - (SELECT * FROM person \ - WHERE last_name = p.last_name \ - AND state = p.state)"; - - let expected = "Projection: p.id\ - \n Filter: EXISTS ()\ - \n Subquery:\ - \n Projection: *\ - \n Filter: person.last_name = outer_ref(p.last_name) AND person.state = outer_ref(p.state)\ - \n TableScan: person\ - \n SubqueryAlias: p\ - \n TableScan: person"; - quick_test(sql, expected); -} - #[test] fn in_subquery_uncorrelated() { let sql = "SELECT id FROM person p WHERE id IN \ @@ -2958,88 +2719,6 @@ fn scalar_subquery_reference_outer_field() { quick_test(sql, expected); } -#[test] -fn subquery_references_cte() { - let sql = "WITH \ - cte AS (SELECT * FROM person) \ - SELECT * FROM person WHERE EXISTS (SELECT * FROM cte WHERE id = person.id)"; - - let expected = "Projection: *\ - \n Filter: EXISTS ()\ - \n Subquery:\ - \n Projection: *\ - \n Filter: cte.id = outer_ref(person.id)\ - \n SubqueryAlias: cte\ - \n Projection: *\ - \n TableScan: person\ - \n TableScan: person"; - - quick_test(sql, expected) -} - -#[test] -fn cte_with_no_column_names() { - let sql = "WITH \ - numbers AS ( \ - SELECT 1 as a, 2 as b, 3 as c \ - ) \ - SELECT * FROM numbers;"; - - let expected = "Projection: *\ - \n SubqueryAlias: numbers\ - \n Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c\ - \n EmptyRelation"; - - quick_test(sql, expected) -} - -#[test] -fn cte_with_column_names() { - let sql = "WITH \ - numbers(a, b, c) AS ( \ - SELECT 1, 2, 3 \ - ) \ - SELECT * FROM numbers;"; - - let expected = "Projection: *\ - \n SubqueryAlias: numbers\ - \n Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c\ - \n Projection: Int64(1), Int64(2), Int64(3)\ - \n EmptyRelation"; - - quick_test(sql, expected) -} - -#[test] -fn cte_with_column_aliases_precedence() { - // The end result should always be what CTE specification says - let sql = "WITH \ - numbers(a, b, c) AS ( \ - SELECT 1 as x, 2 as y, 3 as z \ - ) \ - SELECT * FROM numbers;"; - - let expected = "Projection: *\ - \n SubqueryAlias: numbers\ - \n Projection: x AS a, y AS b, z AS c\ - \n Projection: Int64(1) AS x, Int64(2) AS y, Int64(3) AS z\ - \n EmptyRelation"; - quick_test(sql, expected) -} - -#[test] -fn cte_unbalanced_number_of_columns() { - let sql = "WITH \ - numbers(a) AS ( \ - SELECT 1, 2, 3 \ - ) \ - SELECT * FROM numbers;"; - - let expected = "Error during planning: Source table contains 3 columns but only 1 names given as column alias"; - let result = logical_plan(sql).err().unwrap(); - assert_eq!(result.strip_backtrace(), expected); -} - #[test] fn aggregate_with_rollup() { let sql = @@ -3133,128 +2812,6 @@ fn join_on_complex_condition() { quick_test(sql, expected); } -#[test] -fn lateral_constant() { - let sql = "SELECT * FROM j1, LATERAL (SELECT 1) AS j2"; - let expected = "Projection: *\ - \n Cross Join: \ - \n TableScan: j1\ - \n SubqueryAlias: j2\ - \n Projection: Int64(1)\ - \n EmptyRelation"; - quick_test(sql, expected); -} - -#[test] -fn lateral_comma_join() { - let sql = "SELECT j1_string, j2_string FROM - j1, \ - LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2"; - let expected = "Projection: j1.j1_string, j2.j2_string\ - \n Cross Join: \ - \n TableScan: j1\ - \n SubqueryAlias: j2\ - \n Subquery:\ - \n Projection: *\ - \n Filter: outer_ref(j1.j1_id) < j2.j2_id\ - \n TableScan: j2"; - quick_test(sql, expected); -} - -#[test] -fn lateral_comma_join_referencing_join_rhs() { - let sql = "SELECT * FROM\ - \n j1 JOIN (j2 JOIN j3 ON(j2_id = j3_id - 2)) ON(j1_id = j2_id),\ - \n LATERAL (SELECT * FROM j3 WHERE j3_string = j2_string) as j4;"; - let expected = "Projection: *\ - \n Cross Join: \ - \n Inner Join: Filter: j1.j1_id = j2.j2_id\ - \n TableScan: j1\ - \n Inner Join: Filter: j2.j2_id = j3.j3_id - Int64(2)\ - \n TableScan: j2\ - \n TableScan: j3\ - \n SubqueryAlias: j4\ - \n Subquery:\ - \n Projection: *\ - \n Filter: j3.j3_string = outer_ref(j2.j2_string)\ - \n TableScan: j3"; - quick_test(sql, expected); -} - -#[test] -fn lateral_comma_join_with_shadowing() { - // The j1_id on line 3 references the (closest) j1 definition from line 2. - let sql = "\ - SELECT * FROM j1, LATERAL (\ - SELECT * FROM j1, LATERAL (\ - SELECT * FROM j2 WHERE j1_id = j2_id\ - ) as j2\ - ) as j2;"; - let expected = "Projection: *\ - \n Cross Join: \ - \n TableScan: j1\ - \n SubqueryAlias: j2\ - \n Subquery:\ - \n Projection: *\ - \n Cross Join: \ - \n TableScan: j1\ - \n SubqueryAlias: j2\ - \n Subquery:\ - \n Projection: *\ - \n Filter: outer_ref(j1.j1_id) = j2.j2_id\ - \n TableScan: j2"; - quick_test(sql, expected); -} - -#[test] -fn lateral_left_join() { - let sql = "SELECT j1_string, j2_string FROM \ - j1 \ - LEFT JOIN LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2 ON(true);"; - let expected = "Projection: j1.j1_string, j2.j2_string\ - \n Left Join: Filter: Boolean(true)\ - \n TableScan: j1\ - \n SubqueryAlias: j2\ - \n Subquery:\ - \n Projection: *\ - \n Filter: outer_ref(j1.j1_id) < j2.j2_id\ - \n TableScan: j2"; - quick_test(sql, expected); -} - -#[test] -fn lateral_nested_left_join() { - let sql = "SELECT * FROM - j1, \ - (j2 LEFT JOIN LATERAL (SELECT * FROM j3 WHERE j1_id + j2_id = j3_id) AS j3 ON(true))"; - let expected = "Projection: *\ - \n Cross Join: \ - \n TableScan: j1\ - \n Left Join: Filter: Boolean(true)\ - \n TableScan: j2\ - \n SubqueryAlias: j3\ - \n Subquery:\ - \n Projection: *\ - \n Filter: outer_ref(j1.j1_id) + outer_ref(j2.j2_id) = j3.j3_id\ - \n TableScan: j3"; - quick_test(sql, expected); -} - -#[test] -fn lateral_unnest() { - let sql = "SELECT * from unnest_table u, unnest(u.array_col)"; - let expected = "Projection: *\ - \n Cross Join: \ - \n SubqueryAlias: u\ - \n TableScan: unnest_table\ - \n Subquery:\ - \n Projection: __unnest_placeholder(outer_ref(u.array_col),depth=1) AS UNNEST(outer_ref(u.array_col))\ - \n Unnest: lists[__unnest_placeholder(outer_ref(u.array_col))|depth=1] structs[]\ - \n Projection: outer_ref(u.array_col) AS __unnest_placeholder(outer_ref(u.array_col))\ - \n EmptyRelation"; - quick_test(sql, expected); -} - #[test] fn hive_aggregate_with_filter() -> Result<()> { let dialect = &HiveDialect {}; @@ -3515,20 +3072,6 @@ fn test_one_side_constant_full_join() { quick_test(sql, expected); } -#[test] -fn test_select_all_inner_join() { - let sql = "SELECT * - FROM person \ - INNER JOIN orders \ - ON orders.customer_id * 2 = person.id + 10"; - - let expected = "Projection: *\ - \n Inner Join: Filter: orders.customer_id * Int64(2) = person.id + Int64(10)\ - \n TableScan: person\ - \n TableScan: orders"; - quick_test(sql, expected); -} - #[test] fn test_select_join_key_inner_join() { let sql = "SELECT orders.customer_id * 2, person.id + 10 @@ -4258,34 +3801,6 @@ fn test_prepare_statement_to_plan_limit() { prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan); } -#[test] -fn test_prepare_statement_to_plan_value_list() { - let sql = "PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter);"; - - let expected_plan = "Prepare: \"my_plan\" [Utf8, Utf8] \ - \n Projection: *\ - \n SubqueryAlias: t\ - \n Projection: column1 AS num, column2 AS letter\ - \n Values: (Int64(1), $1), (Int64(2), $2)"; - - let expected_dt = "[Utf8, Utf8]"; - - let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt); - - /////////////////// - // replace params with values - let param_values = vec![ - ScalarValue::from("a".to_string()), - ScalarValue::from("b".to_string()), - ]; - let expected_plan = "Projection: *\ - \n SubqueryAlias: t\ - \n Projection: column1 AS num, column2 AS letter\ - \n Values: (Int64(1), Utf8(\"a\") AS $1), (Int64(2), Utf8(\"b\") AS $2)"; - - prepare_stmt_replace_params_quick_test(plan, param_values, expected_plan); -} - #[test] fn test_prepare_statement_unknown_list_param() { let sql = "SELECT id from person where id = $2"; @@ -4320,44 +3835,6 @@ fn test_prepare_statement_bad_list_idx() { assert_contains!(err.to_string(), "Error during planning: Failed to parse placeholder id: invalid digit found in string"); } -#[test] -fn test_table_alias() { - let sql = "select * from (\ - (select id from person) t1 \ - CROSS JOIN \ - (select age from person) t2 \ - ) as f"; - - let expected = "Projection: *\ - \n SubqueryAlias: f\ - \n Cross Join: \ - \n SubqueryAlias: t1\ - \n Projection: person.id\ - \n TableScan: person\ - \n SubqueryAlias: t2\ - \n Projection: person.age\ - \n TableScan: person"; - quick_test(sql, expected); - - let sql = "select * from (\ - (select id from person) t1 \ - CROSS JOIN \ - (select age from person) t2 \ - ) as f (c1, c2)"; - - let expected = "Projection: *\ - \n SubqueryAlias: f\ - \n Projection: t1.id AS c1, t2.age AS c2\ - \n Cross Join: \ - \n SubqueryAlias: t1\ - \n Projection: person.id\ - \n TableScan: person\ - \n SubqueryAlias: t2\ - \n Projection: person.age\ - \n TableScan: person"; - quick_test(sql, expected); -} - #[test] fn test_inner_join_with_cast_key() { let sql = "SELECT person.id, person.age diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index 504e5164805a..5742f583acf7 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -42,7 +42,7 @@ async-trait = { workspace = true } bigdecimal = { workspace = true } bytes = { workspace = true, optional = true } chrono = { workspace = true, optional = true } -clap = { version = "4.5.30", features = ["derive", "env"] } +clap = { version = "4.5.31", features = ["derive", "env"] } datafusion = { workspace = true, default-features = true, features = ["avro"] } futures = { workspace = true } half = { workspace = true, default-features = true } @@ -55,17 +55,18 @@ postgres-types = { version = "0.2.8", features = ["derive", "with-chrono-0_4"], rust_decimal = { version = "1.36.0", features = ["tokio-pg"] } # When updating the following dependency verify that sqlite test file regeneration works correctly # by running the regenerate_sqlite_files.sh script. -sqllogictest = "0.27.2" +sqllogictest = "0.28.0" sqlparser = { workspace = true } tempfile = { workspace = true } testcontainers = { version = "0.23", features = ["default"], optional = true } testcontainers-modules = { version = "0.11", features = ["postgres"], optional = true } -thiserror = "2.0.0" +thiserror = "2.0.12" tokio = { workspace = true } tokio-postgres = { version = "0.7.12", optional = true } [features] avro = ["datafusion/avro"] +backtrace = ["datafusion/backtrace"] postgres = [ "bytes", "chrono", diff --git a/datafusion/sqllogictest/test_files/alias.slt b/datafusion/sqllogictest/test_files/alias.slt new file mode 100644 index 000000000000..340ffb6078e4 --- /dev/null +++ b/datafusion/sqllogictest/test_files/alias.slt @@ -0,0 +1,59 @@ + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +# test table alias +statement count 0 +create table t1(id int); + +statement count 0 +create table t2(age int); + +query TT +explain select * from ((select id from t1) cross join (select age from t2)) as f; +---- +logical_plan +01)SubqueryAlias: f +02)--Cross Join: +03)----TableScan: t1 projection=[id] +04)----TableScan: t2 projection=[age] +physical_plan +01)CrossJoinExec +02)--DataSourceExec: partitions=1, partition_sizes=[0] +03)--DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain select * from ((select id from t1) cross join (select age from t2)) as f(c1, c2); +---- +logical_plan +01)SubqueryAlias: f +02)--Projection: t1.id AS c1, t2.age AS c2 +03)----Cross Join: +04)------TableScan: t1 projection=[id] +05)------TableScan: t2 projection=[age] +physical_plan +01)ProjectionExec: expr=[id@0 as c1, age@1 as c2] +02)--CrossJoinExec +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +statement count 0 +drop table t1; + +statement count 0 +drop table t2; \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt index 6b5b246aee51..3b7f12960681 100644 --- a/datafusion/sqllogictest/test_files/array.slt +++ b/datafusion/sqllogictest/test_files/array.slt @@ -1204,8 +1204,10 @@ select array_element([1, 2], NULL); ---- NULL -query error +query I select array_element(NULL, 2); +---- +NULL # array_element scalar function #1 (with positive index) query IT @@ -1433,6 +1435,93 @@ NULL 23 NULL 43 5 NULL + +## array_max +# array_max scalar function #1 (with positive index) +query I +select array_max(make_array(5, 3, 6, 4)); +---- +6 + +query I +select array_max(make_array(5, 3, 4, NULL, 6, NULL)); +---- +6 + +query I +select array_max(make_array(NULL, NULL)); +---- +NULL + +query T +select array_max(make_array('h', 'e', 'o', 'l', 'l')); +---- +o + +query T +select array_max(make_array('h', 'e', 'l', NULL, 'l', 'o', NULL)); +---- +o + +query B +select array_max(make_array(false, true, false, true)); +---- +true + +query B +select array_max(make_array(false, true, NULL, false, true)); +---- +true + +query D +select array_max(make_array(DATE '1992-09-01', DATE '1993-03-01', DATE '1999-05-01', DATE '1985-11-01')); +---- +1999-05-01 + +query D +select array_max(make_array(DATE '1995-09-01', DATE '1999-05-01', DATE '1993-03-01', NULL)); +---- +1999-05-01 + +query P +select array_max(make_array(TIMESTAMP '1992-09-01', TIMESTAMP '1995-06-01', TIMESTAMP '1984-10-01')); +---- +1995-06-01T00:00:00 + +query P +select array_max(make_array(NULL, TIMESTAMP '1996-10-01', TIMESTAMP '1995-06-01')); +---- +1996-10-01T00:00:00 + +query R +select array_max(make_array(5.1, -3.2, 6.3, 4.9)); +---- +6.3 + +query ?I +select input, array_max(input) from (select make_array(d - 1, d, d + 1) input from (values (0), (10), (20), (30), (NULL)) t(d)) +---- +[-1, 0, 1] 1 +[9, 10, 11] 11 +[19, 20, 21] 21 +[29, 30, 31] 31 +[NULL, NULL, NULL] NULL + +query II +select array_max(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)')), array_max(arrow_cast(make_array(1), 'FixedSizeList(1, Int64)')); +---- +3 1 + +query I +select array_max(make_array()); +---- +NULL + +# Testing with empty arguments should result in an error +query error DataFusion error: Error during planning: 'array_max' does not support zero arguments +select array_max(); + + ## array_pop_back (aliases: `list_pop_back`) # array_pop_back scalar function with null @@ -2265,6 +2354,52 @@ select array_sort([]); ---- [] +# test with null arguments +query ? +select array_sort(NULL); +---- +NULL + +query ? +select array_sort(column1, NULL) from arrays_values; +---- +NULL +NULL +NULL +NULL +NULL +NULL +NULL +NULL + +query ?? +select array_sort(column1, 'DESC', NULL), array_sort(column1, 'ASC', NULL) from arrays_values; +---- +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL + +query ?? +select array_sort(column1, NULL, 'NULLS FIRST'), array_sort(column1, NULL, 'NULLS LAST') from arrays_values; +---- +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL +NULL NULL + +## test with argument of incorrect types +query error DataFusion error: Execution error: the second parameter of array_sort expects DESC or ASC +select array_sort([1, 3, null, 5, NULL, -5], 1), array_sort([1, 3, null, 5, NULL, -5], 'DESC', 1), array_sort([1, 3, null, 5, NULL, -5], 1, 1); + # test with empty row, the row that does not match the condition has row count 0 statement ok create table t1(a int, b int) as values (100, 1), (101, 2), (102, 3), (101, 2); @@ -2290,8 +2425,10 @@ select list_sort(make_array(1, 3, null, 5, NULL, -5)), list_sort(make_array(1, 3 # array_append with NULLs -query error +query ? select array_append(null, 1); +---- +[1] query error select array_append(null, [2, 3]); @@ -2539,8 +2676,10 @@ select array_append(column1, arrow_cast(make_array(1, 11, 111), 'FixedSizeList(3 # DuckDB: [4] # ClickHouse: Null # Since they dont have the same result, we just follow Postgres, return error -query error +query ? select array_prepend(4, NULL); +---- +[4] query ? select array_prepend(4, []); @@ -2575,11 +2714,10 @@ select array_prepend(null, [[1,2,3]]); query error select array_prepend([], []); -# DuckDB: [null] -# ClickHouse: [null] -# TODO: We may also return [null] -query error +query ? select array_prepend(null, null); +---- +[NULL] query ? select array_append([], null); @@ -5264,9 +5402,11 @@ NULL [3] [5] # array_ndims scalar function #1 #follow PostgreSQL -query error +query I select array_ndims(null); +---- +NULL query I select diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index f39ff56ce449..e2bb23e35732 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -631,3 +631,20 @@ COPY source_table to '/tmp/table.parquet' (row_group_size 55 + 102); # Copy using execution.keep_partition_by_columns with an invalid value query error DataFusion error: Invalid or Unsupported Configuration: provided value for 'execution.keep_partition_by_columns' was not recognized: "invalid_value" COPY source_table to '/tmp/table.parquet' OPTIONS (execution.keep_partition_by_columns invalid_value); + +statement count 0 +create table t; + +query TT +explain COPY (select * from t limit 10) to 'output.csv'; +---- +logical_plan +01)CopyTo: format=csv output_url=output.csv options: () +02)--Limit: skip=0, fetch=10 +03)----TableScan: t projection=[], fetch=10 +physical_plan +01)DataSinkExec: sink=CsvSink(file_groups=[]) +02)--DataSourceExec: partitions=1, partition_sizes=[0], fetch=10 + +statement count 0 +drop table t; diff --git a/datafusion/sqllogictest/test_files/cte.slt b/datafusion/sqllogictest/test_files/cte.slt index 95b9b5a9252e..e019af9775a4 100644 --- a/datafusion/sqllogictest/test_files/cte.slt +++ b/datafusion/sqllogictest/test_files/cte.slt @@ -859,3 +859,149 @@ SELECT * FROM 400 500 1 400 500 2 400 500 3 + +query error DataFusion error: Error during planning: Source table contains 3 columns but only 1 names given as column alias +with numbers(a) as (select 1, 2, 3) select * from numbers; + +query TT +explain with numbers(a,b,c) as (select 1 as x, 2 as y, 3 as z) select * from numbers; +---- +logical_plan +01)SubqueryAlias: numbers +02)--Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c +03)----EmptyRelation +physical_plan +01)ProjectionExec: expr=[1 as a, 2 as b, 3 as c] +02)--PlaceholderRowExec + +query TT +explain with numbers(a,b,c) as (select 1,2,3) select * from numbers; +---- +logical_plan +01)SubqueryAlias: numbers +02)--Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c +03)----EmptyRelation +physical_plan +01)ProjectionExec: expr=[1 as a, 2 as b, 3 as c] +02)--PlaceholderRowExec + +query TT +explain with numbers as (select 1 as a, 2 as b, 3 as c) select * from numbers; +---- +logical_plan +01)SubqueryAlias: numbers +02)--Projection: Int64(1) AS a, Int64(2) AS b, Int64(3) AS c +03)----EmptyRelation +physical_plan +01)ProjectionExec: expr=[1 as a, 2 as b, 3 as c] +02)--PlaceholderRowExec + +statement count 0 +create table person (id int, name string, primary key(id)) + +query TT +explain with cte as (select * from person) SELECT * FROM person WHERE EXISTS (SELECT * FROM cte WHERE id = person.id); +---- +logical_plan +01)LeftSemi Join: person.id = __correlated_sq_1.id +02)--TableScan: person projection=[id, name] +03)--SubqueryAlias: __correlated_sq_1 +04)----SubqueryAlias: cte +05)------TableScan: person projection=[id] +physical_plan +01)CoalesceBatchesExec: target_batch_size=8182 +02)--HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(id@0, id@0)] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +statement count 0 +drop table person; + +statement count 0 +create table j1(a int); + +statement count 0 +create table j2(b int); + +query TT +explain SELECT * FROM j1, LATERAL (SELECT 1) AS j2; +---- +logical_plan +01)Cross Join: +02)--TableScan: j1 projection=[a] +03)--SubqueryAlias: j2 +04)----Projection: Int64(1) +05)------EmptyRelation +physical_plan +01)CrossJoinExec +02)--DataSourceExec: partitions=1, partition_sizes=[0] +03)--ProjectionExec: expr=[1 as Int64(1)] +04)----PlaceholderRowExec + +statement count 0 +drop table j1; + +statement count 0 +drop table j2; + +query TT +explain WITH RECURSIVE numbers AS ( + select 1 as n + UNION ALL + select n + 1 FROM numbers WHERE N < 10 +) select * from numbers; +---- +logical_plan +01)SubqueryAlias: numbers +02)--RecursiveQuery: is_distinct=false +03)----Projection: Int64(1) AS n +04)------EmptyRelation +05)----Projection: numbers.n + Int64(1) +06)------Filter: numbers.n < Int64(10) +07)--------TableScan: numbers +physical_plan +01)RecursiveQueryExec: name=numbers, is_distinct=false +02)--ProjectionExec: expr=[1 as n] +03)----PlaceholderRowExec +04)--CoalescePartitionsExec +05)----ProjectionExec: expr=[n@0 + 1 as numbers.n + Int64(1)] +06)------CoalesceBatchesExec: target_batch_size=8182 +07)--------FilterExec: n@0 < 10 +08)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)------------WorkTableExec: name=numbers + +query TT +explain WITH RECURSIVE numbers AS ( + select 1 as n + UNION ALL + select n + 1 FROM numbers WHERE N < 10 +) select * from numbers; +---- +logical_plan +01)SubqueryAlias: numbers +02)--RecursiveQuery: is_distinct=false +03)----Projection: Int64(1) AS n +04)------EmptyRelation +05)----Projection: numbers.n + Int64(1) +06)------Filter: numbers.n < Int64(10) +07)--------TableScan: numbers +physical_plan +01)RecursiveQueryExec: name=numbers, is_distinct=false +02)--ProjectionExec: expr=[1 as n] +03)----PlaceholderRowExec +04)--CoalescePartitionsExec +05)----ProjectionExec: expr=[n@0 + 1 as numbers.n + Int64(1)] +06)------CoalesceBatchesExec: target_batch_size=8182 +07)--------FilterExec: n@0 < 10 +08)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)------------WorkTableExec: name=numbers + +statement count 0 +set datafusion.execution.enable_recursive_ctes = false; + +query error DataFusion error: This feature is not implemented: Recursive CTEs are not enabled +explain WITH RECURSIVE numbers AS ( + select 1 as n + UNION ALL + select n + 1 FROM numbers WHERE N < 10 +) select * from numbers; diff --git a/datafusion/sqllogictest/test_files/ddl.slt b/datafusion/sqllogictest/test_files/ddl.slt index 6f75a7d7f8fd..bc15f2210380 100644 --- a/datafusion/sqllogictest/test_files/ddl.slt +++ b/datafusion/sqllogictest/test_files/ddl.slt @@ -827,3 +827,31 @@ drop table table_with_pk; statement ok set datafusion.catalog.information_schema = false; + +# Test VARCHAR is mapped to Utf8View during SQL planning when setting map_varchar_to_utf8view to true +statement ok +CREATE TABLE t1(c1 VARCHAR(10) NOT NULL, c2 VARCHAR); + +query TTT +DESCRIBE t1; +---- +c1 Utf8 NO +c2 Utf8 YES + +statement ok +set datafusion.sql_parser.map_varchar_to_utf8view = true; + +statement ok +CREATE TABLE t2(c1 VARCHAR(10) NOT NULL, c2 VARCHAR); + +query TTT +DESCRIBE t2; +---- +c1 Utf8View NO +c2 Utf8View YES + +statement ok +DROP TABLE t1; + +statement ok +DROP TABLE t2; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 16c61a1db6ee..1d63d02bb941 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -175,13 +175,11 @@ initial_logical_plan 01)Projection: simple_explain_test.a, simple_explain_test.b, simple_explain_test.c 02)--TableScan: simple_explain_test logical_plan after inline_table_scan SAME TEXT AS ABOVE -logical_plan after expand_wildcard_rule SAME TEXT AS ABOVE logical_plan after resolve_grouping_function SAME TEXT AS ABOVE logical_plan after type_coercion SAME TEXT AS ABOVE analyzed_logical_plan SAME TEXT AS ABOVE logical_plan after eliminate_nested_union SAME TEXT AS ABOVE logical_plan after simplify_expressions SAME TEXT AS ABOVE -logical_plan after unwrap_cast_in_comparison SAME TEXT AS ABOVE logical_plan after replace_distinct_aggregate SAME TEXT AS ABOVE logical_plan after eliminate_join SAME TEXT AS ABOVE logical_plan after decorrelate_predicate_subquery SAME TEXT AS ABOVE @@ -200,13 +198,11 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after simplify_expressions SAME TEXT AS ABOVE -logical_plan after unwrap_cast_in_comparison SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after optimize_projections TableScan: simple_explain_test projection=[a, b, c] logical_plan after eliminate_nested_union SAME TEXT AS ABOVE logical_plan after simplify_expressions SAME TEXT AS ABOVE -logical_plan after unwrap_cast_in_comparison SAME TEXT AS ABOVE logical_plan after replace_distinct_aggregate SAME TEXT AS ABOVE logical_plan after eliminate_join SAME TEXT AS ABOVE logical_plan after decorrelate_predicate_subquery SAME TEXT AS ABOVE @@ -225,7 +221,6 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after simplify_expressions SAME TEXT AS ABOVE -logical_plan after unwrap_cast_in_comparison SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after optimize_projections SAME TEXT AS ABOVE @@ -247,8 +242,8 @@ physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after coalesce_batches SAME TEXT AS ABOVE physical_plan after OutputRequirements DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after LimitPushdown SAME TEXT AS ABOVE +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan_with_stats DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:),(Col[1]:),(Col[2]:)]] @@ -323,8 +318,8 @@ physical_plan after OutputRequirements 01)GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Exact(671), [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Exact(671), [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Exact(671), [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Exact(671), [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] physical_plan_with_schema DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, schema=[id:Int32;N, bool_col:Boolean;N, tinyint_col:Int32;N, smallint_col:Int32;N, int_col:Int32;N, bigint_col:Int64;N, float_col:Float32;N, double_col:Float64;N, date_string_col:BinaryView;N, string_col:BinaryView;N, timestamp_col:Timestamp(Nanosecond, None);N] @@ -363,8 +358,8 @@ physical_plan after OutputRequirements 01)GlobalLimitExec: skip=0, fetch=10 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet physical_plan after LimitAggregation SAME TEXT AS ABOVE -physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after LimitPushdown DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet +physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after SanityCheckPlan SAME TEXT AS ABOVE physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet physical_plan_with_stats DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/parquet-testing/data/alltypes_plain.parquet]]}, projection=[id, bool_col, tinyint_col, smallint_col, int_col, bigint_col, float_col, double_col, date_string_col, string_col, timestamp_col], limit=10, file_type=parquet, statistics=[Rows=Exact(8), Bytes=Exact(671), [(Col[0]:),(Col[1]:),(Col[2]:),(Col[3]:),(Col[4]:),(Col[5]:),(Col[6]:),(Col[7]:),(Col[8]:),(Col[9]:),(Col[10]:)]] diff --git a/datafusion/sqllogictest/test_files/explain_tree.slt b/datafusion/sqllogictest/test_files/explain_tree.slt new file mode 100644 index 000000000000..bafc4d559b80 --- /dev/null +++ b/datafusion/sqllogictest/test_files/explain_tree.slt @@ -0,0 +1,1643 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Tests for tree explain + + + +statement ok +set datafusion.explain.format = "tree"; + +######## Setup Data Files ####### + +# table1: CSV +query I +COPY (VALUES (1, 'foo', 1, '2023-01-01'), (2, 'bar', 2, '2023-01-02'), (3, 'baz', 3, '2023-01-03')) +TO 'test_files/scratch/explain_tree/table1.csv'; +---- +3 + +statement ok +CREATE EXTERNAL TABLE table1 ( + int_col INT, + string_col TEXT, + bigint_col BIGINT, + date_col DATE +) +STORED AS CSV +LOCATION 'test_files/scratch/explain_tree/table1.csv'; + +# table2: Parquet +query I +COPY (SELECT * from table1) +TO 'test_files/scratch/explain_tree/table2.parquet' +---- +3 + +statement ok +CREATE EXTERNAL TABLE table2 +STORED AS PARQUET +LOCATION 'test_files/scratch/explain_tree/table2.parquet'; + + +# table3: Memory +statement ok +CREATE TABLE table3 as select * from table1; + +# table4: JSON +query I +COPY (SELECT * from table1) +TO 'test_files/scratch/explain_tree/table4.json' +---- +3 + +statement ok +CREATE EXTERNAL TABLE table4 +STORED AS JSON +LOCATION 'test_files/scratch/explain_tree/table4.json'; + +# table5: ARROW +query I +COPY (SELECT * from table1) +TO 'test_files/scratch/explain_tree/table5.arrow' +---- +3 + +statement ok +CREATE EXTERNAL TABLE table5 +STORED AS ARROW +LOCATION 'test_files/scratch/explain_tree/table5.arrow'; + +statement ok +CREATE UNBOUNDED EXTERNAL TABLE annotated_data_infinite2 ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER +) +STORED AS CSV +WITH ORDER (a ASC, b ASC, c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +statement ok +CREATE TABLE hashjoin_datatype_table_t1_source(c1 INT, c2 BIGINT, c3 DECIMAL(5,2), c4 VARCHAR) +AS VALUES +(1, 86400000, 1.23, 'abc'), +(2, 172800000, 456.00, 'def'), +(null, 259200000, 789.000, 'ghi'), +(3, null, -123.12, 'jkl') +; + +statement ok +CREATE TABLE hashjoin_datatype_table_t1 +AS SELECT + arrow_cast(c1, 'Date32') as c1, + arrow_cast(c2, 'Date64') as c2, + c3, + arrow_cast(c4, 'Dictionary(Int32, Utf8)') as c4 +FROM + hashjoin_datatype_table_t1_source + +statement ok +CREATE TABLE hashjoin_datatype_table_t2_source(c1 INT, c2 BIGINT, c3 DECIMAL(10,2), c4 VARCHAR) +AS VALUES +(1, 86400000, -123.12, 'abc'), +(null, null, 100000.00, 'abcdefg'), +(null, 259200000, 0.00, 'qwerty'), +(3, null, 789.000, 'qwe') +; + +statement ok +CREATE TABLE hashjoin_datatype_table_t2 +AS SELECT + arrow_cast(c1, 'Date32') as c1, + arrow_cast(c2, 'Date64') as c2, + c3, + arrow_cast(c4, 'Dictionary(Int32, Utf8)') as c4 +FROM + hashjoin_datatype_table_t2_source + +######## Begin Queries ######## + +# Filter +query TT +explain SELECT int_col FROM table1 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: csv │ +24)└───────────────────────────┘ + +# Aggregate +query TT +explain SELECT string_col, SUM(bigint_col) FROM table1 GROUP BY string_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ AggregateExec │ +03)│ -------------------- │ +04)│ aggr: │ +05)│ sum(table1.bigint_col) │ +06)│ │ +07)│ group_by: │ +08)│ string_col@0 as string_col│ +09)│ │ +10)│ mode: │ +11)│ FinalPartitioned │ +12)└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ +14)│ CoalesceBatchesExec │ +15)└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐ +17)│ RepartitionExec │ +18)│ -------------------- │ +19)│ output_partition_count: │ +20)│ 4 │ +21)│ │ +22)│ partitioning_scheme: │ +23)│ Hash([string_col@0], 4) │ +24)└─────────────┬─────────────┘ +25)┌─────────────┴─────────────┐ +26)│ AggregateExec │ +27)│ -------------------- │ +28)│ aggr: │ +29)│ sum(table1.bigint_col) │ +30)│ │ +31)│ group_by: │ +32)│ string_col@0 as string_col│ +33)│ │ +34)│ mode: Partial │ +35)└─────────────┬─────────────┘ +36)┌─────────────┴─────────────┐ +37)│ RepartitionExec │ +38)│ -------------------- │ +39)│ output_partition_count: │ +40)│ 1 │ +41)│ │ +42)│ partitioning_scheme: │ +43)│ RoundRobinBatch(4) │ +44)└─────────────┬─────────────┘ +45)┌─────────────┴─────────────┐ +46)│ DataSourceExec │ +47)│ -------------------- │ +48)│ files: 1 │ +49)│ format: csv │ +50)└───────────────────────────┘ + + +# Limit +query TT +explain SELECT int_col FROM table1 LIMIT 3,2; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ GlobalLimitExec │ +03)│ -------------------- │ +04)│ limit: 2 │ +05)│ skip: 3 │ +06)└─────────────┬─────────────┘ +07)┌─────────────┴─────────────┐ +08)│ DataSourceExec │ +09)│ -------------------- │ +10)│ files: 1 │ +11)│ format: csv │ +12)└───────────────────────────┘ + +# 2 Joins +query TT +explain SELECT table1.string_col, table2.date_col FROM table1 JOIN table2 ON table1.int_col = table2.int_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ HashJoinExec │ +06)│ -------------------- │ +07)│ on: ├──────────────┐ +08)│ (int_col@0 = int_col@0) │ │ +09)└─────────────┬─────────────┘ │ +10)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +11)│ CoalesceBatchesExec ││ CoalesceBatchesExec │ +12)└─────────────┬─────────────┘└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +14)│ RepartitionExec ││ RepartitionExec │ +15)│ -------------------- ││ -------------------- │ +16)│ output_partition_count: ││ output_partition_count: │ +17)│ 4 ││ 4 │ +18)│ ││ │ +19)│ partitioning_scheme: ││ partitioning_scheme: │ +20)│ Hash([int_col@0], 4) ││ Hash([int_col@0], 4) │ +21)└─────────────┬─────────────┘└─────────────┬─────────────┘ +22)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +23)│ RepartitionExec ││ RepartitionExec │ +24)│ -------------------- ││ -------------------- │ +25)│ output_partition_count: ││ output_partition_count: │ +26)│ 1 ││ 1 │ +27)│ ││ │ +28)│ partitioning_scheme: ││ partitioning_scheme: │ +29)│ RoundRobinBatch(4) ││ RoundRobinBatch(4) │ +30)└─────────────┬─────────────┘└─────────────┬─────────────┘ +31)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +32)│ DataSourceExec ││ DataSourceExec │ +33)│ -------------------- ││ -------------------- │ +34)│ files: 1 ││ files: 1 │ +35)│ format: csv ││ format: parquet │ +36)└───────────────────────────┘└───────────────────────────┘ + +# 3 Joins +query TT +explain SELECT + table1.string_col, + table2.date_col, + table3.date_col +FROM + table1 JOIN table2 ON table1.int_col = table2.int_col + JOIN table3 ON table2.int_col = table3.int_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ HashJoinExec │ +06)│ -------------------- │ +07)│ on: ├───────────────────────────────────────────┐ +08)│ (int_col@1 = int_col@0) │ │ +09)└─────────────┬─────────────┘ │ +10)┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ +11)│ CoalesceBatchesExec │ │ CoalesceBatchesExec │ +12)└─────────────┬─────────────┘ └─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ ┌─────────────┴─────────────┐ +14)│ HashJoinExec │ │ RepartitionExec │ +15)│ -------------------- │ │ -------------------- │ +16)│ on: │ │ output_partition_count: │ +17)│ (int_col@0 = int_col@0) ├──────────────┐ │ 1 │ +18)│ │ │ │ │ +19)│ │ │ │ partitioning_scheme: │ +20)│ │ │ │ Hash([int_col@0], 4) │ +21)└─────────────┬─────────────┘ │ └─────────────┬─────────────┘ +22)┌─────────────┴─────────────┐┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +23)│ CoalesceBatchesExec ││ CoalesceBatchesExec ││ DataSourceExec │ +24)│ ││ ││ -------------------- │ +25)│ ││ ││ bytes: 1560 │ +26)│ ││ ││ format: memory │ +27)│ ││ ││ rows: 1 │ +28)└─────────────┬─────────────┘└─────────────┬─────────────┘└───────────────────────────┘ +29)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +30)│ RepartitionExec ││ RepartitionExec │ +31)│ -------------------- ││ -------------------- │ +32)│ output_partition_count: ││ output_partition_count: │ +33)│ 4 ││ 4 │ +34)│ ││ │ +35)│ partitioning_scheme: ││ partitioning_scheme: │ +36)│ Hash([int_col@0], 4) ││ Hash([int_col@0], 4) │ +37)└─────────────┬─────────────┘└─────────────┬─────────────┘ +38)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +39)│ RepartitionExec ││ RepartitionExec │ +40)│ -------------------- ││ -------------------- │ +41)│ output_partition_count: ││ output_partition_count: │ +42)│ 1 ││ 1 │ +43)│ ││ │ +44)│ partitioning_scheme: ││ partitioning_scheme: │ +45)│ RoundRobinBatch(4) ││ RoundRobinBatch(4) │ +46)└─────────────┬─────────────┘└─────────────┬─────────────┘ +47)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +48)│ DataSourceExec ││ DataSourceExec │ +49)│ -------------------- ││ -------------------- │ +50)│ files: 1 ││ files: 1 │ +51)│ format: csv ││ format: parquet │ +52)└───────────────────────────┘└───────────────────────────┘ + +# Long Filter (demonstrate what happens with wrapping) +query TT +explain SELECT int_col FROM table1 +WHERE string_col != 'foo' AND string_col != 'bar' AND string_col != 'a really long string constant' +; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo AND │ +09)│ string_col@1 != bar │ +10)│ AND string_col@1 != a │ +11)│ really long string │ +12)│ constant │ +13)└─────────────┬─────────────┘ +14)┌─────────────┴─────────────┐ +15)│ RepartitionExec │ +16)│ -------------------- │ +17)│ output_partition_count: │ +18)│ 1 │ +19)│ │ +20)│ partitioning_scheme: │ +21)│ RoundRobinBatch(4) │ +22)└─────────────┬─────────────┘ +23)┌─────────────┴─────────────┐ +24)│ DataSourceExec │ +25)│ -------------------- │ +26)│ files: 1 │ +27)│ format: csv │ +28)└───────────────────────────┘ + +# Check maximum line limit. +query TT +explain SELECT int_col FROM table1 +WHERE string_col != 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != │ +09)│ aaaaaaaaaaaaaa │ +10)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +11)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +12)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +13)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +14)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +15)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +16)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +17)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +18)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +19)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +20)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +21)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +22)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +23)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +24)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +25)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +26)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +27)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +28)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +29)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +30)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +31)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +32)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +33)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +34)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +35)│aaaaaaaaaaaaaaaaaaaaaaaaaaa│ +36)│ ... │ +37)└─────────────┬─────────────┘ +38)┌─────────────┴─────────────┐ +39)│ RepartitionExec │ +40)│ -------------------- │ +41)│ output_partition_count: │ +42)│ 1 │ +43)│ │ +44)│ partitioning_scheme: │ +45)│ RoundRobinBatch(4) │ +46)└─────────────┬─────────────┘ +47)┌─────────────┴─────────────┐ +48)│ DataSourceExec │ +49)│ -------------------- │ +50)│ files: 1 │ +51)│ format: csv │ +52)└───────────────────────────┘ + +# Check exactly the render width. +query TT +explain SELECT int_col FROM table1 +WHERE string_col != 'aaaaaaaaaaa'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│string_col@1 != aaaaaaaaaaa│ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: csv │ +24)└───────────────────────────┘ + +# Check with the render witdth + 1. +query TT +explain SELECT int_col FROM table1 +WHERE string_col != 'aaaaaaaaaaaa'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != │ +09)│ aaaaaaaaaaaa │ +10)└─────────────┬─────────────┘ +11)┌─────────────┴─────────────┐ +12)│ RepartitionExec │ +13)│ -------------------- │ +14)│ output_partition_count: │ +15)│ 1 │ +16)│ │ +17)│ partitioning_scheme: │ +18)│ RoundRobinBatch(4) │ +19)└─────────────┬─────────────┘ +20)┌─────────────┴─────────────┐ +21)│ DataSourceExec │ +22)│ -------------------- │ +23)│ files: 1 │ +24)│ format: csv │ +25)└───────────────────────────┘ + +# Query with filter on csv +query TT +explain SELECT int_col FROM table1 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: csv │ +24)└───────────────────────────┘ + + +# Query with filter on parquet +query TT +explain SELECT int_col FROM table2 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: parquet │ +24)│ │ +25)│ predicate: │ +26)│ string_col@1 != foo │ +27)└───────────────────────────┘ + +# Query with filter on memory +query TT +explain SELECT int_col FROM table3 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ DataSourceExec │ +12)│ -------------------- │ +13)│ bytes: 1560 │ +14)│ format: memory │ +15)│ rows: 1 │ +16)└───────────────────────────┘ + +# Query with filter on json +query TT +explain SELECT int_col FROM table4 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: json │ +24)└───────────────────────────┘ + +# Query with filter on arrow +query TT +explain SELECT int_col FROM table5 WHERE string_col != 'foo'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ FilterExec │ +06)│ -------------------- │ +07)│ predicate: │ +08)│ string_col@1 != foo │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ RepartitionExec │ +12)│ -------------------- │ +13)│ output_partition_count: │ +14)│ 1 │ +15)│ │ +16)│ partitioning_scheme: │ +17)│ RoundRobinBatch(4) │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ DataSourceExec │ +21)│ -------------------- │ +22)│ files: 1 │ +23)│ format: arrow │ +24)└───────────────────────────┘ + + +# Query with window agg. +query TT +explain select count(*) over() from table1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ count(*) ROWS BETWEEN │ +05)│ UNBOUNDED PRECEDING │ +06)│ AND UNBOUNDED FOLLOWING: │ +07)│ count(Int64(1)) ROWS │ +08)│ BETWEEN UNBOUNDED │ +09)│ PRECEDING AND UNBOUNDED │ +10)│ FOLLOWING@0 │ +11)└─────────────┬─────────────┘ +12)┌─────────────┴─────────────┐ +13)│ WindowAggExec │ +14)│ -------------------- │ +15)│ select_list: │ +16)│ count(Int64(1)) ROWS │ +17)│ BETWEEN UNBOUNDED │ +18)│ PRECEDING AND UNBOUNDED │ +19)│ FOLLOWING │ +20)└─────────────┬─────────────┘ +21)┌─────────────┴─────────────┐ +22)│ DataSourceExec │ +23)│ -------------------- │ +24)│ files: 1 │ +25)│ format: csv │ +26)└───────────────────────────┘ + +# Query with bounded window agg. +query TT +explain SELECT + v1, + SUM(v1) OVER (ORDER BY v1 ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS rolling_sum +FROM generate_series(1, 1000) AS t1(v1); +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ rolling_sum: │ +05)│ sum(t1.v1) ORDER BY [t1.v1│ +06)│ ASC NULLS LAST] ROWS │ +07)│ BETWEEN 1 PRECEDING │ +08)│ AND CURRENT ROW@1 │ +09)│ │ +10)│ v1: v1@0 │ +11)└─────────────┬─────────────┘ +12)┌─────────────┴─────────────┐ +13)│ BoundedWindowAggExec │ +14)│ -------------------- │ +15)│ mode: Sorted │ +16)│ │ +17)│ select_list: │ +18)│ sum(t1.v1) ORDER BY [t1.v1│ +19)│ ASC NULLS LAST] ROWS │ +20)│ BETWEEN 1 PRECEDING │ +21)│ AND CURRENT ROW │ +22)└─────────────┬─────────────┘ +23)┌─────────────┴─────────────┐ +24)│ SortExec │ +25)│ -------------------- │ +26)│ v1@0 ASC NULLS LAST │ +27)└─────────────┬─────────────┘ +28)┌─────────────┴─────────────┐ +29)│ ProjectionExec │ +30)│ -------------------- │ +31)│ v1: value@0 │ +32)└─────────────┬─────────────┘ +33)┌─────────────┴─────────────┐ +34)│ LazyMemoryExec │ +35)└───────────────────────────┘ + +query TT +explain select + count(*) over(), + row_number() over () +from table1 +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ count(*) ROWS BETWEEN │ +05)│ UNBOUNDED PRECEDING │ +06)│ AND UNBOUNDED FOLLOWING: │ +07)│ count(Int64(1)) ROWS │ +08)│ BETWEEN UNBOUNDED │ +09)│ PRECEDING AND UNBOUNDED │ +10)│ FOLLOWING@0 │ +11)│ │ +12)│ row_number() ROWS BETWEEN │ +13)│ UNBOUNDED PRECEDING AND │ +14)│ UNBOUNDED FOLLOWING: │ +15)│ row_number() ROWS BETWEEN │ +16)│ UNBOUNDED PRECEDING AND │ +17)│ UNBOUNDED FOLLOWING@1 │ +18)└─────────────┬─────────────┘ +19)┌─────────────┴─────────────┐ +20)│ WindowAggExec │ +21)│ -------------------- │ +22)│ select_list: │ +23)│ count(Int64(1)) ROWS │ +24)│ BETWEEN UNBOUNDED │ +25)│ PRECEDING AND UNBOUNDED │ +26)│ FOLLOWING, row_number() │ +27)│ ROWS BETWEEN UNBOUNDED │ +28)│ PRECEDING AND UNBOUNDED │ +29)│ FOLLOWING │ +30)└─────────────┬─────────────┘ +31)┌─────────────┴─────────────┐ +32)│ DataSourceExec │ +33)│ -------------------- │ +34)│ files: 1 │ +35)│ format: csv │ +36)└───────────────────────────┘ + +# Query for sort. +query TT +explain SELECT * FROM table1 ORDER BY string_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortExec │ +03)│ -------------------- │ +04)│string_col@1 ASC NULLS LAST│ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ DataSourceExec │ +08)│ -------------------- │ +09)│ files: 1 │ +10)│ format: csv │ +11)└───────────────────────────┘ + +# Query for sort with limit. +query TT +explain SELECT * FROM table1 ORDER BY string_col LIMIT 1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortExec │ +03)│ -------------------- │ +04)│ limit: 1 │ +05)│ │ +06)│string_col@1 ASC NULLS LAST│ +07)└─────────────┬─────────────┘ +08)┌─────────────┴─────────────┐ +09)│ DataSourceExec │ +10)│ -------------------- │ +11)│ files: 1 │ +12)│ format: csv │ +13)└───────────────────────────┘ + +# Query with projection on csv +query TT +explain SELECT int_col, bigint_col, int_col+bigint_col AS sum_col FROM table1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ bigint_col: │ +05)│ bigint_col@1 │ +06)│ │ +07)│ int_col: int_col@0 │ +08)│ │ +09)│ sum_col: │ +10)│ CAST(int_col@0 AS Int64) +│ +11)│ bigint_col@1 │ +12)└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ +14)│ RepartitionExec │ +15)│ -------------------- │ +16)│ output_partition_count: │ +17)│ 1 │ +18)│ │ +19)│ partitioning_scheme: │ +20)│ RoundRobinBatch(4) │ +21)└─────────────┬─────────────┘ +22)┌─────────────┴─────────────┐ +23)│ DataSourceExec │ +24)│ -------------------- │ +25)│ files: 1 │ +26)│ format: csv │ +27)└───────────────────────────┘ + +query TT +explain select + rank() over(ORDER BY int_col DESC), + row_number() over (ORDER BY int_col ASC) +from table1 +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ rank() ORDER BY [table1 │ +05)│ .int_col DESC NULLS │ +06)│ FIRST] RANGE BETWEEN │ +07)│ UNBOUNDED PRECEDING AND │ +08)│ CURRENT ROW: │ +09)│ rank() ORDER BY [table1 │ +10)│ .int_col DESC NULLS │ +11)│ FIRST] RANGE BETWEEN │ +12)│ UNBOUNDED PRECEDING AND │ +13)│ CURRENT ROW@1 │ +14)│ │ +15)│ row_number() ORDER BY │ +16)│ [table1.int_col ASC │ +17)│ NULLS LAST] RANGE │ +18)│ BETWEEN UNBOUNDED │ +19)│ PRECEDING AND CURRENT │ +20)│ ROW: │ +21)│ row_number() ORDER BY │ +22)│ [table1.int_col ASC │ +23)│ NULLS LAST] RANGE │ +24)│ BETWEEN UNBOUNDED │ +25)│ PRECEDING AND CURRENT │ +26)│ ROW@2 │ +27)└─────────────┬─────────────┘ +28)┌─────────────┴─────────────┐ +29)│ BoundedWindowAggExec │ +30)│ -------------------- │ +31)│ mode: Sorted │ +32)│ │ +33)│ select_list: │ +34)│ row_number() ORDER BY │ +35)│ [table1.int_col ASC │ +36)│ NULLS LAST] RANGE │ +37)│ BETWEEN UNBOUNDED │ +38)│ PRECEDING AND CURRENT │ +39)│ ROW │ +40)└─────────────┬─────────────┘ +41)┌─────────────┴─────────────┐ +42)│ SortExec │ +43)│ -------------------- │ +44)│ int_col@0 ASC NULLS LAST │ +45)└─────────────┬─────────────┘ +46)┌─────────────┴─────────────┐ +47)│ BoundedWindowAggExec │ +48)│ -------------------- │ +49)│ mode: Sorted │ +50)│ │ +51)│ select_list: │ +52)│ rank() ORDER BY [table1 │ +53)│ .int_col DESC NULLS │ +54)│ FIRST] RANGE BETWEEN │ +55)│ UNBOUNDED PRECEDING AND │ +56)│ CURRENT ROW │ +57)└─────────────┬─────────────┘ +58)┌─────────────┴─────────────┐ +59)│ SortExec │ +60)│ -------------------- │ +61)│ int_col@0 DESC │ +62)└─────────────┬─────────────┘ +63)┌─────────────┴─────────────┐ +64)│ DataSourceExec │ +65)│ -------------------- │ +66)│ files: 1 │ +67)│ format: csv │ +68)└───────────────────────────┘ + +# Query with projection on parquet +query TT +explain SELECT int_col, bigint_col, int_col+bigint_col AS sum_col FROM table2; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ bigint_col: │ +05)│ bigint_col@1 │ +06)│ │ +07)│ int_col: int_col@0 │ +08)│ │ +09)│ sum_col: │ +10)│ CAST(int_col@0 AS Int64) +│ +11)│ bigint_col@1 │ +12)└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ +14)│ RepartitionExec │ +15)│ -------------------- │ +16)│ output_partition_count: │ +17)│ 1 │ +18)│ │ +19)│ partitioning_scheme: │ +20)│ RoundRobinBatch(4) │ +21)└─────────────┬─────────────┘ +22)┌─────────────┴─────────────┐ +23)│ DataSourceExec │ +24)│ -------------------- │ +25)│ files: 1 │ +26)│ format: parquet │ +27)└───────────────────────────┘ + + +# Query with projection on memory +query TT +explain SELECT int_col, bigint_col, int_col+bigint_col AS sum_col FROM table3; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ bigint_col: │ +05)│ bigint_col@1 │ +06)│ │ +07)│ int_col: int_col@0 │ +08)│ │ +09)│ sum_col: │ +10)│ CAST(int_col@0 AS Int64) +│ +11)│ bigint_col@1 │ +12)└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ +14)│ DataSourceExec │ +15)│ -------------------- │ +16)│ bytes: 1560 │ +17)│ format: memory │ +18)│ rows: 1 │ +19)└───────────────────────────┘ + +# Query with projection on json +query TT +explain SELECT int_col, bigint_col, int_col+bigint_col AS sum_col FROM table4; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ bigint_col: │ +05)│ bigint_col@0 │ +06)│ │ +07)│ int_col: int_col@1 │ +08)│ │ +09)│ sum_col: │ +10)│ int_col@1 + bigint_col@0 │ +11)└─────────────┬─────────────┘ +12)┌─────────────┴─────────────┐ +13)│ RepartitionExec │ +14)│ -------------------- │ +15)│ output_partition_count: │ +16)│ 1 │ +17)│ │ +18)│ partitioning_scheme: │ +19)│ RoundRobinBatch(4) │ +20)└─────────────┬─────────────┘ +21)┌─────────────┴─────────────┐ +22)│ DataSourceExec │ +23)│ -------------------- │ +24)│ files: 1 │ +25)│ format: json │ +26)└───────────────────────────┘ + + +# Query with projection on arrow +query TT +explain SELECT int_col, bigint_col, int_col+bigint_col AS sum_col FROM table5; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ bigint_col: │ +05)│ bigint_col@1 │ +06)│ │ +07)│ int_col: int_col@0 │ +08)│ │ +09)│ sum_col: │ +10)│ CAST(int_col@0 AS Int64) +│ +11)│ bigint_col@1 │ +12)└─────────────┬─────────────┘ +13)┌─────────────┴─────────────┐ +14)│ RepartitionExec │ +15)│ -------------------- │ +16)│ output_partition_count: │ +17)│ 1 │ +18)│ │ +19)│ partitioning_scheme: │ +20)│ RoundRobinBatch(4) │ +21)└─────────────┬─────────────┘ +22)┌─────────────┴─────────────┐ +23)│ DataSourceExec │ +24)│ -------------------- │ +25)│ files: 1 │ +26)│ format: arrow │ +27)└───────────────────────────┘ + +# Query with PartialSortExec. +query TT +EXPLAIN SELECT * +FROM annotated_data_infinite2 +ORDER BY a, b, d; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ PartialSortExec │ +03)│ -------------------- │ +04)│ a@1 ASC NULLS LAST, b@2 │ +05)│ ASC NULLS LAST, d@4 │ +06)│ ASC NULLS LAST │ +07)└─────────────┬─────────────┘ +08)┌─────────────┴─────────────┐ +09)│ StreamingTableExec │ +10)│ -------------------- │ +11)│ infinite: true │ +12)│ limit: None │ +13)└───────────────────────────┘ + +query TT +EXPLAIN SELECT * +FROM annotated_data_infinite2 +ORDER BY a, b, d +LIMIT 50; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ PartialSortExec │ +03)│ -------------------- │ +04)│ a@1 ASC NULLS LAST, b@2 │ +05)│ ASC NULLS LAST, d@4 │ +06)│ ASC NULLS LAST │ +07)│ │ +08)│ limit: 50 │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ StreamingTableExec │ +12)│ -------------------- │ +13)│ infinite: true │ +14)│ limit: None │ +15)└───────────────────────────┘ + +# Query with hash join. +query TT +explain select * from table1 inner join table2 on table1.int_col = table2.int_col and table1.string_col = table2.string_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ HashJoinExec │ +06)│ -------------------- │ +07)│ on: │ +08)│ (int_col@0 = int_col@0), ├──────────────┐ +09)│ (CAST(table1.string_col │ │ +10)│ AS Utf8View)@4 = │ │ +11)│ string_col@1) │ │ +12)└─────────────┬─────────────┘ │ +13)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +14)│ CoalesceBatchesExec ││ CoalesceBatchesExec │ +15)└─────────────┬─────────────┘└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +17)│ RepartitionExec ││ RepartitionExec │ +18)│ -------------------- ││ -------------------- │ +19)│ output_partition_count: ││ output_partition_count: │ +20)│ 4 ││ 4 │ +21)│ ││ │ +22)│ partitioning_scheme: ││ partitioning_scheme: │ +23)│ Hash([int_col@0, CAST ││ Hash([int_col@0, │ +24)│ (table1.string_col ││ string_col@1], │ +25)│ AS Utf8View)@4], 4) ││ 4) │ +26)└─────────────┬─────────────┘└─────────────┬─────────────┘ +27)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +28)│ ProjectionExec ││ RepartitionExec │ +29)│ -------------------- ││ -------------------- │ +30)│ CAST(table1.string_col AS ││ output_partition_count: │ +31)│ Utf8View): ││ 1 │ +32)│ CAST(string_col@1 AS ││ │ +33)│ Utf8View) ││ partitioning_scheme: │ +34)│ ││ RoundRobinBatch(4) │ +35)│ bigint_col: ││ │ +36)│ bigint_col@2 ││ │ +37)│ ││ │ +38)│ date_col: date_col@3 ││ │ +39)│ int_col: int_col@0 ││ │ +40)│ ││ │ +41)│ string_col: ││ │ +42)│ string_col@1 ││ │ +43)└─────────────┬─────────────┘└─────────────┬─────────────┘ +44)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +45)│ RepartitionExec ││ DataSourceExec │ +46)│ -------------------- ││ -------------------- │ +47)│ output_partition_count: ││ files: 1 │ +48)│ 1 ││ format: parquet │ +49)│ ││ │ +50)│ partitioning_scheme: ││ │ +51)│ RoundRobinBatch(4) ││ │ +52)└─────────────┬─────────────┘└───────────────────────────┘ +53)┌─────────────┴─────────────┐ +54)│ DataSourceExec │ +55)│ -------------------- │ +56)│ files: 1 │ +57)│ format: csv │ +58)└───────────────────────────┘ + +# Query with outer hash join. +query TT +explain select * from table1 left outer join table2 on table1.int_col = table2.int_col and table1.string_col = table2.string_col; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalesceBatchesExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ HashJoinExec │ +06)│ -------------------- │ +07)│ join_type: Left │ +08)│ │ +09)│ on: ├──────────────┐ +10)│ (int_col@0 = int_col@0), │ │ +11)│ (CAST(table1.string_col │ │ +12)│ AS Utf8View)@4 = │ │ +13)│ string_col@1) │ │ +14)└─────────────┬─────────────┘ │ +15)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +16)│ CoalesceBatchesExec ││ CoalesceBatchesExec │ +17)└─────────────┬─────────────┘└─────────────┬─────────────┘ +18)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +19)│ RepartitionExec ││ RepartitionExec │ +20)│ -------------------- ││ -------------------- │ +21)│ output_partition_count: ││ output_partition_count: │ +22)│ 4 ││ 4 │ +23)│ ││ │ +24)│ partitioning_scheme: ││ partitioning_scheme: │ +25)│ Hash([int_col@0, CAST ││ Hash([int_col@0, │ +26)│ (table1.string_col ││ string_col@1], │ +27)│ AS Utf8View)@4], 4) ││ 4) │ +28)└─────────────┬─────────────┘└─────────────┬─────────────┘ +29)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +30)│ ProjectionExec ││ RepartitionExec │ +31)│ -------------------- ││ -------------------- │ +32)│ CAST(table1.string_col AS ││ output_partition_count: │ +33)│ Utf8View): ││ 1 │ +34)│ CAST(string_col@1 AS ││ │ +35)│ Utf8View) ││ partitioning_scheme: │ +36)│ ││ RoundRobinBatch(4) │ +37)│ bigint_col: ││ │ +38)│ bigint_col@2 ││ │ +39)│ ││ │ +40)│ date_col: date_col@3 ││ │ +41)│ int_col: int_col@0 ││ │ +42)│ ││ │ +43)│ string_col: ││ │ +44)│ string_col@1 ││ │ +45)└─────────────┬─────────────┘└─────────────┬─────────────┘ +46)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +47)│ RepartitionExec ││ DataSourceExec │ +48)│ -------------------- ││ -------------------- │ +49)│ output_partition_count: ││ files: 1 │ +50)│ 1 ││ format: parquet │ +51)│ ││ │ +52)│ partitioning_scheme: ││ │ +53)│ RoundRobinBatch(4) ││ │ +54)└─────────────┬─────────────┘└───────────────────────────┘ +55)┌─────────────┴─────────────┐ +56)│ DataSourceExec │ +57)│ -------------------- │ +58)│ files: 1 │ +59)│ format: csv │ +60)└───────────────────────────┘ + +# Query with nested loop join. +query TT +explain select int_col from table1 where exists (select count(*) from table2); +---- +physical_plan +01)┌───────────────────────────┐ +02)│ NestedLoopJoinExec │ +03)│ -------------------- ├──────────────┐ +04)│ join_type: LeftSemi │ │ +05)└─────────────┬─────────────┘ │ +06)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +07)│ DataSourceExec ││ ProjectionExec │ +08)│ -------------------- ││ │ +09)│ files: 1 ││ │ +10)│ format: csv ││ │ +11)└───────────────────────────┘└─────────────┬─────────────┘ +12)-----------------------------┌─────────────┴─────────────┐ +13)-----------------------------│ AggregateExec │ +14)-----------------------------│ -------------------- │ +15)-----------------------------│ aggr: count(Int64(1)) │ +16)-----------------------------│ mode: Final │ +17)-----------------------------└─────────────┬─────────────┘ +18)-----------------------------┌─────────────┴─────────────┐ +19)-----------------------------│ CoalescePartitionsExec │ +20)-----------------------------└─────────────┬─────────────┘ +21)-----------------------------┌─────────────┴─────────────┐ +22)-----------------------------│ AggregateExec │ +23)-----------------------------│ -------------------- │ +24)-----------------------------│ aggr: count(Int64(1)) │ +25)-----------------------------│ mode: Partial │ +26)-----------------------------└─────────────┬─────────────┘ +27)-----------------------------┌─────────────┴─────────────┐ +28)-----------------------------│ RepartitionExec │ +29)-----------------------------│ -------------------- │ +30)-----------------------------│ output_partition_count: │ +31)-----------------------------│ 1 │ +32)-----------------------------│ │ +33)-----------------------------│ partitioning_scheme: │ +34)-----------------------------│ RoundRobinBatch(4) │ +35)-----------------------------└─────────────┬─────────────┘ +36)-----------------------------┌─────────────┴─────────────┐ +37)-----------------------------│ DataSourceExec │ +38)-----------------------------│ -------------------- │ +39)-----------------------------│ files: 1 │ +40)-----------------------------│ format: parquet │ +41)-----------------------------└───────────────────────────┘ + +# Query with cross join. +query TT +explain select * from table1 cross join table2 ; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CrossJoinExec ├──────────────┐ +03)└─────────────┬─────────────┘ │ +04)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +05)│ DataSourceExec ││ RepartitionExec │ +06)│ -------------------- ││ -------------------- │ +07)│ files: 1 ││ output_partition_count: │ +08)│ format: csv ││ 1 │ +09)│ ││ │ +10)│ ││ partitioning_scheme: │ +11)│ ││ RoundRobinBatch(4) │ +12)└───────────────────────────┘└─────────────┬─────────────┘ +13)-----------------------------┌─────────────┴─────────────┐ +14)-----------------------------│ DataSourceExec │ +15)-----------------------------│ -------------------- │ +16)-----------------------------│ files: 1 │ +17)-----------------------------│ format: parquet │ +18)-----------------------------└───────────────────────────┘ + + +# Query with sort merge join. +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query TT +explain select * from hashjoin_datatype_table_t1 t1 join hashjoin_datatype_table_t2 t2 on t1.c1 = t2.c1 +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortMergeJoinExec │ +03)│ -------------------- ├──────────────┐ +04)│ on: (c1@0 = c1@0) │ │ +05)└─────────────┬─────────────┘ │ +06)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +07)│ SortExec ││ SortExec │ +08)│ -------------------- ││ -------------------- │ +09)│ c1@0 ASC ││ c1@0 ASC │ +10)└─────────────┬─────────────┘└─────────────┬─────────────┘ +11)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +12)│ DataSourceExec ││ DataSourceExec │ +13)│ -------------------- ││ -------------------- │ +14)│ bytes: 6040 ││ bytes: 6040 │ +15)│ format: memory ││ format: memory │ +16)│ rows: 1 ││ rows: 1 │ +17)└───────────────────────────┘└───────────────────────────┘ + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +# cleanup +statement ok +drop table table1; + +statement ok +drop table table2; + +statement ok +drop table table3; + +statement ok +drop table table4; + +statement ok +drop table table5; + +# Test on StreamingTableExec +# prepare table +statement ok +CREATE UNBOUNDED EXTERNAL TABLE data ( + "date" DATE, + "ticker" VARCHAR, + "time" TIMESTAMP, +) STORED AS CSV +WITH ORDER ("date", "ticker", "time") +LOCATION './a.parquet'; + + +# query +query TT +explain SELECT * FROM data +WHERE ticker = 'A' +ORDER BY "date", "time"; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortPreservingMergeExec │ +03)│ -------------------- │ +04)│ date@0 ASC NULLS LAST, │ +05)│ time@2 ASC NULLS LAST │ +06)└─────────────┬─────────────┘ +07)┌─────────────┴─────────────┐ +08)│ CoalesceBatchesExec │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ FilterExec │ +12)│ -------------------- │ +13)│ predicate: │ +14)│ ticker@1 = A │ +15)└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐ +17)│ RepartitionExec │ +18)│ -------------------- │ +19)│ output_partition_count: │ +20)│ 1 │ +21)│ │ +22)│ partitioning_scheme: │ +23)│ RoundRobinBatch(4) │ +24)└─────────────┬─────────────┘ +25)┌─────────────┴─────────────┐ +26)│ StreamingTableExec │ +27)│ -------------------- │ +28)│ infinite: true │ +29)│ limit: None │ +30)└───────────────────────────┘ + + +# constant ticker, CAST(time AS DATE) = time, order by time +query TT +explain SELECT * FROM data +WHERE ticker = 'A' AND CAST(time AS DATE) = date +ORDER BY "time" +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortPreservingMergeExec │ +03)│ -------------------- │ +04)│ time@2 ASC NULLS LAST │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ CoalesceBatchesExec │ +08)└─────────────┬─────────────┘ +09)┌─────────────┴─────────────┐ +10)│ FilterExec │ +11)│ -------------------- │ +12)│ predicate: │ +13)│ ticker@1 = A AND CAST(time│ +14)│ @2 AS Date32) = date@0 │ +15)└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐ +17)│ RepartitionExec │ +18)│ -------------------- │ +19)│ output_partition_count: │ +20)│ 1 │ +21)│ │ +22)│ partitioning_scheme: │ +23)│ RoundRobinBatch(4) │ +24)└─────────────┬─────────────┘ +25)┌─────────────┴─────────────┐ +26)│ StreamingTableExec │ +27)│ -------------------- │ +28)│ infinite: true │ +29)│ limit: None │ +30)└───────────────────────────┘ + +# same thing but order by date +query TT +explain SELECT * FROM data +WHERE ticker = 'A' AND CAST(time AS DATE) = date +ORDER BY "date" +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortPreservingMergeExec │ +03)│ -------------------- │ +04)│ date@0 ASC NULLS LAST │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ CoalesceBatchesExec │ +08)└─────────────┬─────────────┘ +09)┌─────────────┴─────────────┐ +10)│ FilterExec │ +11)│ -------------------- │ +12)│ predicate: │ +13)│ ticker@1 = A AND CAST(time│ +14)│ @2 AS Date32) = date@0 │ +15)└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐ +17)│ RepartitionExec │ +18)│ -------------------- │ +19)│ output_partition_count: │ +20)│ 1 │ +21)│ │ +22)│ partitioning_scheme: │ +23)│ RoundRobinBatch(4) │ +24)└─────────────┬─────────────┘ +25)┌─────────────┴─────────────┐ +26)│ StreamingTableExec │ +27)│ -------------------- │ +28)│ infinite: true │ +29)│ limit: None │ +30)└───────────────────────────┘ + +# same thing but order by ticker +query TT +explain SELECT * FROM data +WHERE ticker = 'A' AND CAST(time AS DATE) = date +ORDER BY "ticker" +---- +physical_plan +01)┌───────────────────────────┐ +02)│ CoalescePartitionsExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ CoalesceBatchesExec │ +06)└─────────────┬─────────────┘ +07)┌─────────────┴─────────────┐ +08)│ FilterExec │ +09)│ -------------------- │ +10)│ predicate: │ +11)│ ticker@1 = A AND CAST(time│ +12)│ @2 AS Date32) = date@0 │ +13)└─────────────┬─────────────┘ +14)┌─────────────┴─────────────┐ +15)│ RepartitionExec │ +16)│ -------------------- │ +17)│ output_partition_count: │ +18)│ 1 │ +19)│ │ +20)│ partitioning_scheme: │ +21)│ RoundRobinBatch(4) │ +22)└─────────────┬─────────────┘ +23)┌─────────────┴─────────────┐ +24)│ StreamingTableExec │ +25)│ -------------------- │ +26)│ infinite: true │ +27)│ limit: None │ +28)└───────────────────────────┘ + + +# same thing but order by time, date +query TT +explain SELECT * FROM data +WHERE ticker = 'A' AND CAST(time AS DATE) = date +ORDER BY "time", "date"; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortPreservingMergeExec │ +03)│ -------------------- │ +04)│ time@2 ASC NULLS LAST, │ +05)│ date@0 ASC NULLS LAST │ +06)└─────────────┬─────────────┘ +07)┌─────────────┴─────────────┐ +08)│ CoalesceBatchesExec │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ FilterExec │ +12)│ -------------------- │ +13)│ predicate: │ +14)│ ticker@1 = A AND CAST(time│ +15)│ @2 AS Date32) = date@0 │ +16)└─────────────┬─────────────┘ +17)┌─────────────┴─────────────┐ +18)│ RepartitionExec │ +19)│ -------------------- │ +20)│ output_partition_count: │ +21)│ 1 │ +22)│ │ +23)│ partitioning_scheme: │ +24)│ RoundRobinBatch(4) │ +25)└─────────────┬─────────────┘ +26)┌─────────────┴─────────────┐ +27)│ StreamingTableExec │ +28)│ -------------------- │ +29)│ infinite: true │ +30)│ limit: None │ +31)└───────────────────────────┘ + + + + +# query +query TT +explain SELECT * FROM data +WHERE date = '2006-01-02' +ORDER BY "ticker", "time"; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ SortPreservingMergeExec │ +03)│ -------------------- │ +04)│ ticker@1 ASC NULLS LAST, │ +05)│ time@2 ASC NULLS LAST │ +06)└─────────────┬─────────────┘ +07)┌─────────────┴─────────────┐ +08)│ CoalesceBatchesExec │ +09)└─────────────┬─────────────┘ +10)┌─────────────┴─────────────┐ +11)│ FilterExec │ +12)│ -------------------- │ +13)│ predicate: │ +14)│ date@0 = 2006-01-02 │ +15)└─────────────┬─────────────┘ +16)┌─────────────┴─────────────┐ +17)│ RepartitionExec │ +18)│ -------------------- │ +19)│ output_partition_count: │ +20)│ 1 │ +21)│ │ +22)│ partitioning_scheme: │ +23)│ RoundRobinBatch(4) │ +24)└─────────────┬─────────────┘ +25)┌─────────────┴─────────────┐ +26)│ StreamingTableExec │ +27)│ -------------------- │ +28)│ infinite: true │ +29)│ limit: None │ +30)└───────────────────────────┘ + + + +# Test explain tree for WorkTableExec +query TT +EXPLAIN WITH RECURSIVE nodes AS ( + SELECT 1 as id + UNION ALL + SELECT id + 1 as id + FROM nodes + WHERE id < 10 +) +SELECT * FROM nodes +---- +physical_plan +01)┌───────────────────────────┐ +02)│ RecursiveQueryExec ├──────────────┐ +03)└─────────────┬─────────────┘ │ +04)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +05)│ ProjectionExec ││ CoalescePartitionsExec │ +06)│ -------------------- ││ │ +07)│ id: 1 ││ │ +08)└─────────────┬─────────────┘└─────────────┬─────────────┘ +09)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +10)│ PlaceholderRowExec ││ ProjectionExec │ +11)│ ││ -------------------- │ +12)│ ││ id: id@0 + 1 │ +13)└───────────────────────────┘└─────────────┬─────────────┘ +14)-----------------------------┌─────────────┴─────────────┐ +15)-----------------------------│ CoalesceBatchesExec │ +16)-----------------------------└─────────────┬─────────────┘ +17)-----------------------------┌─────────────┴─────────────┐ +18)-----------------------------│ FilterExec │ +19)-----------------------------│ -------------------- │ +20)-----------------------------│ predicate: id@0 < 10 │ +21)-----------------------------└─────────────┬─────────────┘ +22)-----------------------------┌─────────────┴─────────────┐ +23)-----------------------------│ RepartitionExec │ +24)-----------------------------│ -------------------- │ +25)-----------------------------│ output_partition_count: │ +26)-----------------------------│ 1 │ +27)-----------------------------│ │ +28)-----------------------------│ partitioning_scheme: │ +29)-----------------------------│ RoundRobinBatch(4) │ +30)-----------------------------└─────────────┬─────────────┘ +31)-----------------------------┌─────────────┴─────────────┐ +32)-----------------------------│ WorkTableExec │ +33)-----------------------------│ -------------------- │ +34)-----------------------------│ name: nodes │ +35)-----------------------------└───────────────────────────┘ + +query TT +explain COPY (VALUES (1, 'foo', 1, '2023-01-01'), (2, 'bar', 2, '2023-01-02'), (3, 'baz', 3, '2023-01-03')) +TO 'test_files/scratch/explain_tree/1.json'; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ DataSinkExec │ +03)└─────────────┬─────────────┘ +04)┌─────────────┴─────────────┐ +05)│ DataSourceExec │ +06)│ -------------------- │ +07)│ bytes: 2672 │ +08)│ format: memory │ +09)│ rows: 1 │ +10)└───────────────────────────┘ diff --git a/datafusion/sqllogictest/test_files/expr.slt b/datafusion/sqllogictest/test_files/expr.slt index 7980b180ae68..74e9fe065a73 100644 --- a/datafusion/sqllogictest/test_files/expr.slt +++ b/datafusion/sqllogictest/test_files/expr.slt @@ -2057,3 +2057,16 @@ select 1 where null between null and 2; query T select 'A' where null between 2 and null; ---- + +### Demonstrate use of E literals for escaping +# should not have literal tab +query T +select 'foo\t\tbar'; +---- +foo\t\tbar + +# should have literal tab +query T +select E'foo\t\tbar'; +---- +foo bar diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 2b3ebcda1520..0cc8045dccd0 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -5537,3 +5537,32 @@ drop view t statement ok drop table source; + + +# test select_wildcard_with_groupby +statement count 0 +create table t(a int, b int, c int, "😀" int); + +query TT +explain select * from t group by a, b, c, "😀"; +---- +logical_plan +01)Aggregate: groupBy=[[t.a, t.b, t.c, t.😀]], aggr=[[]] +02)--TableScan: t projection=[a, b, c, 😀] +physical_plan +01)AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b, c@2 as c, 😀@3 as 😀], aggr=[] +02)--DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain select * from (select a, b from t) as c group by a, b; +---- +logical_plan +01)Aggregate: groupBy=[[c.a, c.b]], aggr=[[]] +02)--SubqueryAlias: c +03)----TableScan: t projection=[a, b] +physical_plan +01)AggregateExec: mode=Single, gby=[a@0 as a, b@1 as b], aggr=[] +02)--DataSourceExec: partitions=1, partition_sizes=[0] + +statement count 0 +drop table t; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index cbfaee9bd6fb..496f24abf6ed 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -232,6 +232,7 @@ datafusion.execution.split_file_groups_by_statistics false datafusion.execution.target_partitions 7 datafusion.execution.time_zone +00:00 datafusion.execution.use_row_number_estimates_to_optimize_partitioning false +datafusion.explain.format indent datafusion.explain.logical_plan_only false datafusion.explain.physical_plan_only false datafusion.explain.show_schema false @@ -262,7 +263,9 @@ datafusion.sql_parser.collect_spans false datafusion.sql_parser.dialect generic datafusion.sql_parser.enable_ident_normalization true datafusion.sql_parser.enable_options_value_normalization false +datafusion.sql_parser.map_varchar_to_utf8view false datafusion.sql_parser.parse_float_as_decimal false +datafusion.sql_parser.recursion_limit 50 datafusion.sql_parser.support_varchar_with_length true # show all variables with verbose @@ -328,6 +331,7 @@ datafusion.execution.split_file_groups_by_statistics false Attempt to eliminate datafusion.execution.target_partitions 7 Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system datafusion.execution.time_zone +00:00 The default time zone Some functions, e.g. `EXTRACT(HOUR from SOME_TIME)`, shift the underlying datetime according to this time zone, and then extract the hour datafusion.execution.use_row_number_estimates_to_optimize_partitioning false Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. +datafusion.explain.format indent Display format of explain. Default is "indent". When set to "tree", it will print the plan in a tree-rendered format. datafusion.explain.logical_plan_only false When set to true, the explain statement will only print logical plans datafusion.explain.physical_plan_only false When set to true, the explain statement will only print physical plans datafusion.explain.show_schema false When set to true, the explain statement will print schema information @@ -354,11 +358,13 @@ datafusion.optimizer.repartition_sorts true Should DataFusion execute sorts in a datafusion.optimizer.repartition_windows true Should DataFusion repartition data using the partitions keys to execute window functions in parallel using the provided `target_partitions` level datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys -datafusion.sql_parser.collect_spans false When set to true, the source locations relative to the original SQL query (i.e. [`Span`](sqlparser::tokenizer::Span)) will be collected and recorded in the logical plan nodes. +datafusion.sql_parser.collect_spans false When set to true, the source locations relative to the original SQL query (i.e. [`Span`](https://docs.rs/sqlparser/latest/sqlparser/tokenizer/struct.Span.html)) will be collected and recorded in the logical plan nodes. datafusion.sql_parser.dialect generic Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB and Databricks. datafusion.sql_parser.enable_ident_normalization true When set to true, SQL parser will normalize ident (convert ident to lowercase when not quoted) datafusion.sql_parser.enable_options_value_normalization false When set to true, SQL parser will normalize options value (convert value to lowercase). Note that this option is ignored and will be removed in the future. All case-insensitive values are normalized automatically. +datafusion.sql_parser.map_varchar_to_utf8view false If true, `VARCHAR` is mapped to `Utf8View` during SQL planning. If false, `VARCHAR` is mapped to `Utf8` during SQL planning. Default is false. datafusion.sql_parser.parse_float_as_decimal false When set to true, SQL parser will parse float as decimal type +datafusion.sql_parser.recursion_limit 50 Specifies the recursion depth limit when parsing complex SQL Queries datafusion.sql_parser.support_varchar_with_length true If true, permit lengths for `VARCHAR` such as `VARCHAR(20)`, but ignore the length. If false, error if a `VARCHAR` with a length is specified. The Arrow type system does not have a notion of maximum string length and thus DataFusion can not enforce such limits. # show_variable_in_config_options diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 0397e0c367b1..50af06dc40fc 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -4541,3 +4541,269 @@ DROP TABLE test statement ok set datafusion.execution.target_partitions = 1; + +# test using_join_multiple_keys_subquery +statement count 0 +create table person(id int, age int, state int); + +statement count 0 +create table lineitem(c1 int); + +query TT +explain SELECT * FROM person a join person b using (id, age); +---- +logical_plan +01)Projection: a.id, a.age, a.state, b.state +02)--Inner Join: a.id = b.id, a.age = b.age +03)----SubqueryAlias: a +04)------TableScan: person projection=[id, age, state] +05)----SubqueryAlias: b +06)------TableScan: person projection=[id, age, state] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1)], projection=[id@0, age@1, state@2, state@5] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT age FROM (SELECT * FROM person a join person b using (id, age, state)); +---- +logical_plan +01)Projection: a.age +02)--Inner Join: a.id = b.id, a.age = b.age, a.state = b.state +03)----SubqueryAlias: a +04)------TableScan: person projection=[id, age, state] +05)----SubqueryAlias: b +06)------TableScan: person projection=[id, age, state] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1), (state@2, state@2)], projection=[age@1] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT a.* FROM person a join person b using (id, age); +---- +logical_plan +01)Projection: a.id, a.age, a.state +02)--Inner Join: a.id = b.id, a.age = b.age +03)----SubqueryAlias: a +04)------TableScan: person projection=[id, age, state] +05)----SubqueryAlias: b +06)------TableScan: person projection=[id, age] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1)], projection=[id@0, age@1, state@2] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT a.*, b.* FROM person a join person b using (id, age); +---- +logical_plan +01)Inner Join: a.id = b.id, a.age = b.age +02)--SubqueryAlias: a +03)----TableScan: person projection=[id, age, state] +04)--SubqueryAlias: b +05)----TableScan: person projection=[id, age, state] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1)] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT * FROM person a join person b using (id, age, state) join person c using (id, age, state); +---- +logical_plan +01)Projection: a.id, a.age, a.state +02)--Inner Join: a.id = c.id, a.age = c.age, a.state = c.state +03)----Projection: a.id, a.age, a.state +04)------Inner Join: a.id = b.id, a.age = b.age, a.state = b.state +05)--------SubqueryAlias: a +06)----------TableScan: person projection=[id, age, state] +07)--------SubqueryAlias: b +08)----------TableScan: person projection=[id, age, state] +09)----SubqueryAlias: c +10)------TableScan: person projection=[id, age, state] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1), (state@2, state@2)], projection=[id@0, age@1, state@2] +03)----CoalesceBatchesExec: target_batch_size=3 +04)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0), (age@1, age@1), (state@2, state@2)], projection=[id@0, age@1, state@2] +05)--------DataSourceExec: partitions=1, partition_sizes=[0] +06)--------DataSourceExec: partitions=1, partition_sizes=[0] +07)----DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT * FROM person a NATURAL JOIN lineitem b; +---- +logical_plan +01)Cross Join: +02)--SubqueryAlias: a +03)----TableScan: person projection=[id, age, state] +04)--SubqueryAlias: b +05)----TableScan: lineitem projection=[c1] +physical_plan +01)CrossJoinExec +02)--DataSourceExec: partitions=1, partition_sizes=[0] +03)--DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain SELECT * FROM lineitem JOIN lineitem as lineitem2 USING (c1) +---- +logical_plan +01)Projection: lineitem.c1 +02)--Inner Join: lineitem.c1 = lineitem2.c1 +03)----TableScan: lineitem projection=[c1] +04)----SubqueryAlias: lineitem2 +05)------TableScan: lineitem projection=[c1] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(c1@0, c1@0)], projection=[c1@0] +03)----DataSourceExec: partitions=1, partition_sizes=[0] +04)----DataSourceExec: partitions=1, partition_sizes=[0] + +statement count 0 +drop table person; + +statement count 0 +drop table lineitem; + +statement count 0 +create table j1(j1_string varchar, j1_id int); + +statement count 0 +create table j2(j2_string varchar, j2_id int); + +statement count 0 +create table j3(j3_string varchar, j3_id int); + +statement count 0 +create table j4(j4_string varchar, j4_id int); + +query TT +explain SELECT j1_string, j2_string FROM j1, LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2; +---- +logical_plan +01)Cross Join: +02)--TableScan: j1 projection=[j1_string] +03)--SubqueryAlias: j2 +04)----Projection: j2.j2_string +05)------Subquery: +06)--------Filter: outer_ref(j1.j1_id) < j2.j2_id +07)----------TableScan: j2 projection=[j2_string, j2_id] +physical_plan_error This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn(Int32, Column { relation: Some(Bare { table: "j1" }), name: "j1_id" }) + +query TT +explain SELECT * FROM j1 JOIN (j2 JOIN j3 ON(j2_id = j3_id - 2)) ON(j1_id = j2_id), LATERAL (SELECT * FROM j3 WHERE j3_string = j2_string) as j4 +---- +logical_plan +01)Cross Join: +02)--Inner Join: CAST(j2.j2_id AS Int64) = CAST(j3.j3_id AS Int64) - Int64(2) +03)----Inner Join: j1.j1_id = j2.j2_id +04)------TableScan: j1 projection=[j1_string, j1_id] +05)------TableScan: j2 projection=[j2_string, j2_id] +06)----TableScan: j3 projection=[j3_string, j3_id] +07)--SubqueryAlias: j4 +08)----Subquery: +09)------Filter: j3.j3_string = outer_ref(j2.j2_string) +10)--------TableScan: j3 projection=[j3_string, j3_id] +physical_plan_error This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn(Utf8, Column { relation: Some(Bare { table: "j2" }), name: "j2_string" }) + +query TT +explain SELECT * FROM j1, LATERAL (SELECT * FROM j1, LATERAL (SELECT * FROM j2 WHERE j1_id = j2_id) as j2) as j2; +---- +logical_plan +01)Cross Join: +02)--TableScan: j1 projection=[j1_string, j1_id] +03)--SubqueryAlias: j2 +04)----Subquery: +05)------Cross Join: +06)--------TableScan: j1 projection=[j1_string, j1_id] +07)--------SubqueryAlias: j2 +08)----------Subquery: +09)------------Filter: outer_ref(j1.j1_id) = j2.j2_id +10)--------------TableScan: j2 projection=[j2_string, j2_id] +physical_plan_error This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn(Int32, Column { relation: Some(Bare { table: "j1" }), name: "j1_id" }) + +query TT +explain SELECT j1_string, j2_string FROM j1 LEFT JOIN LATERAL (SELECT * FROM j2 WHERE j1_id < j2_id) AS j2 ON(true); +---- +logical_plan +01)Left Join: +02)--TableScan: j1 projection=[j1_string] +03)--SubqueryAlias: j2 +04)----Projection: j2.j2_string +05)------Subquery: +06)--------Filter: outer_ref(j1.j1_id) < j2.j2_id +07)----------TableScan: j2 projection=[j2_string, j2_id] +physical_plan_error This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn(Int32, Column { relation: Some(Bare { table: "j1" }), name: "j1_id" }) + +query TT +explain SELECT * FROM j1, (j2 LEFT JOIN LATERAL (SELECT * FROM j3 WHERE j1_id + j2_id = j3_id) AS j3 ON(true)); +---- +logical_plan +01)Cross Join: +02)--TableScan: j1 projection=[j1_string, j1_id] +03)--Left Join: +04)----TableScan: j2 projection=[j2_string, j2_id] +05)----SubqueryAlias: j3 +06)------Subquery: +07)--------Filter: outer_ref(j1.j1_id) + outer_ref(j2.j2_id) = j3.j3_id +08)----------TableScan: j3 projection=[j3_string, j3_id] +physical_plan_error This feature is not implemented: Physical plan does not support logical expression OuterReferenceColumn(Int32, Column { relation: Some(Bare { table: "j1" }), name: "j1_id" }) + +query TT +explain SELECT * FROM j1, LATERAL (SELECT 1) AS j2; +---- +logical_plan +01)Cross Join: +02)--TableScan: j1 projection=[j1_string, j1_id] +03)--SubqueryAlias: j2 +04)----Projection: Int64(1) +05)------EmptyRelation +physical_plan +01)CrossJoinExec +02)--DataSourceExec: partitions=1, partition_sizes=[0] +03)--ProjectionExec: expr=[1 as Int64(1)] +04)----PlaceholderRowExec + +statement count 0 +drop table j1; + +statement count 0 +drop table j2; + +statement count 0 +drop table j3; + +statement count 0 +drop table j4; + +statement count 0 +create table person(id int); + +statement count 0 +create table orders(customer_id int); + +query TT +explain SELECT * FROM person INNER JOIN orders ON orders.customer_id * 2 = person.id + 10 +---- +logical_plan +01)Inner Join: CAST(person.id AS Int64) + Int64(10) = CAST(orders.customer_id AS Int64) * Int64(2) +02)--TableScan: person projection=[id] +03)--TableScan: orders projection=[customer_id] +physical_plan +01)CoalesceBatchesExec: target_batch_size=3 +02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(person.id + Int64(10)@1, orders.customer_id * Int64(2)@1)], projection=[id@0, customer_id@2] +03)----ProjectionExec: expr=[id@0 as id, CAST(id@0 AS Int64) + 10 as person.id + Int64(10)] +04)------DataSourceExec: partitions=1, partition_sizes=[0] +05)----ProjectionExec: expr=[customer_id@0 as customer_id, CAST(customer_id@0 AS Int64) * 2 as orders.customer_id * Int64(2)] +06)------DataSourceExec: partitions=1, partition_sizes=[0] + +statement count 0 +drop table person; + +statement count 0 +drop table orders; diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index d7da21c58ec6..f088e071d7e7 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -986,17 +986,26 @@ statement ok create table t(a0 int, a int, b int, c int) as values (1, 2, 3, 4), (5, 6, 7, 8); # expect this query to run successfully, not error +query IIII +select * from (select c, a, NULL::int as a0, b from t order by a, c) t1 +union all +select * from (select c, NULL::int as a, a0, b from t order by a0, c) t2 +order by c, a, a0, b +limit 2; +---- +4 2 NULL 3 +4 NULL 1 3 + query III select * from (select c, a, NULL::int as a0 from t order by a, c) t1 union all select * from (select c, NULL::int as a, a0 from t order by a0, c) t2 -order by c, a, a0, b +order by c, a, a0 limit 2; ---- 4 2 NULL 4 NULL 1 - # Casting from numeric to string types breaks the ordering statement ok CREATE EXTERNAL TABLE ordered_table ( @@ -1232,43 +1241,41 @@ physical_plan # Test: inputs into union with different orderings query TT -explain select * from (select b, c, a, NULL::int as a0 from ordered_table order by a, c) t1 +explain select * from (select b, c, a, NULL::int as a0, d from ordered_table order by a, c) t1 union all -select * from (select b, c, NULL::int as a, a0 from ordered_table order by a0, c) t2 +select * from (select b, c, NULL::int as a, a0, d from ordered_table order by a0, c) t2 order by d, c, a, a0, b limit 2; ---- logical_plan -01)Projection: t1.b, t1.c, t1.a, t1.a0 -02)--Sort: t1.d ASC NULLS LAST, t1.c ASC NULLS LAST, t1.a ASC NULLS LAST, t1.a0 ASC NULLS LAST, t1.b ASC NULLS LAST, fetch=2 -03)----Union -04)------SubqueryAlias: t1 -05)--------Projection: ordered_table.b, ordered_table.c, ordered_table.a, Int32(NULL) AS a0, ordered_table.d -06)----------TableScan: ordered_table projection=[a, b, c, d] -07)------SubqueryAlias: t2 -08)--------Projection: ordered_table.b, ordered_table.c, Int32(NULL) AS a, ordered_table.a0, ordered_table.d -09)----------TableScan: ordered_table projection=[a0, b, c, d] +01)Sort: t1.d ASC NULLS LAST, t1.c ASC NULLS LAST, t1.a ASC NULLS LAST, t1.a0 ASC NULLS LAST, t1.b ASC NULLS LAST, fetch=2 +02)--Union +03)----SubqueryAlias: t1 +04)------Projection: ordered_table.b, ordered_table.c, ordered_table.a, Int32(NULL) AS a0, ordered_table.d +05)--------TableScan: ordered_table projection=[a, b, c, d] +06)----SubqueryAlias: t2 +07)------Projection: ordered_table.b, ordered_table.c, Int32(NULL) AS a, ordered_table.a0, ordered_table.d +08)--------TableScan: ordered_table projection=[a0, b, c, d] physical_plan -01)ProjectionExec: expr=[b@0 as b, c@1 as c, a@2 as a, a0@3 as a0] -02)--SortPreservingMergeExec: [d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a@2 ASC NULLS LAST, a0@3 ASC NULLS LAST, b@0 ASC NULLS LAST], fetch=2 -03)----UnionExec -04)------SortExec: TopK(fetch=2), expr=[d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a@2 ASC NULLS LAST, b@0 ASC NULLS LAST], preserve_partitioning=[false] -05)--------ProjectionExec: expr=[b@1 as b, c@2 as c, a@0 as a, NULL as a0, d@3 as d] -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c, d], output_ordering=[c@2 ASC NULLS LAST], file_type=csv, has_header=true -07)------SortExec: TopK(fetch=2), expr=[d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a0@3 ASC NULLS LAST, b@0 ASC NULLS LAST], preserve_partitioning=[false] -08)--------ProjectionExec: expr=[b@1 as b, c@2 as c, NULL as a, a0@0 as a0, d@3 as d] -09)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, b, c, d], output_ordering=[c@2 ASC NULLS LAST], file_type=csv, has_header=true +01)SortPreservingMergeExec: [d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a@2 ASC NULLS LAST, a0@3 ASC NULLS LAST, b@0 ASC NULLS LAST], fetch=2 +02)--UnionExec +03)----SortExec: TopK(fetch=2), expr=[d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a@2 ASC NULLS LAST, b@0 ASC NULLS LAST], preserve_partitioning=[false] +04)------ProjectionExec: expr=[b@1 as b, c@2 as c, a@0 as a, NULL as a0, d@3 as d] +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c, d], output_ordering=[c@2 ASC NULLS LAST], file_type=csv, has_header=true +06)----SortExec: TopK(fetch=2), expr=[d@4 ASC NULLS LAST, c@1 ASC NULLS LAST, a0@3 ASC NULLS LAST, b@0 ASC NULLS LAST], preserve_partitioning=[false] +07)------ProjectionExec: expr=[b@1 as b, c@2 as c, NULL as a, a0@0 as a0, d@3 as d] +08)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, b, c, d], output_ordering=[c@2 ASC NULLS LAST], file_type=csv, has_header=true # Test: run the query from above -query IIII -select * from (select b, c, a, NULL::int as a0 from ordered_table order by a, c) t1 +query IIIII +select * from (select b, c, a, NULL::int as a0, d from ordered_table order by a, c) t1 union all -select * from (select b, c, NULL::int as a, a0 from ordered_table order by a0, c) t2 +select * from (select b, c, NULL::int as a, a0, d from ordered_table order by a0, c) t2 order by d, c, a, a0, b limit 2; ---- -0 0 0 NULL -0 0 NULL 1 +0 0 0 NULL 0 +0 0 NULL 1 0 statement ok diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index 5d0f417640ec..33df0d26f361 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -312,3 +312,18 @@ SET datafusion.explain.logical_plan_only=false; statement ok DROP TABLE person; + +statement ok +SET datafusion.explain.logical_plan_only=true; + +statement count 0 +PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter); + +statement count 5 +explain PREPARE my_plan(STRING, STRING) AS SELECT * FROM (VALUES(1, $1), (2, $2)) AS t (num, letter); + +query IT +EXECUTE my_plan('a', 'b'); +---- +1 a +2 b diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt index 66413775b393..f583d659fd4f 100644 --- a/datafusion/sqllogictest/test_files/scalar.slt +++ b/datafusion/sqllogictest/test_files/scalar.slt @@ -1927,7 +1927,7 @@ select position('' in '') ---- 1 -query error DataFusion error: Error during planning: Function 'strpos' expects NativeType::String but received NativeType::Int64 +query error DataFusion error: Error during planning: Internal error: Expect TypeSignatureClass::Native\(LogicalType\(Native\(String\), String\)\) but received NativeType::Int64, DataType: Int64 select position(1 in 1) query I diff --git a/datafusion/sqllogictest/test_files/select.slt b/datafusion/sqllogictest/test_files/select.slt index f1ac0696bff9..d5e0c449762f 100644 --- a/datafusion/sqllogictest/test_files/select.slt +++ b/datafusion/sqllogictest/test_files/select.slt @@ -1820,6 +1820,9 @@ query I select a from t; ---- +statement count 0 +drop table t; + statement ok set datafusion.optimizer.max_passes=3; @@ -1842,3 +1845,13 @@ SELECT t1.v1 FROM (SELECT 1 AS "t1.v1"); # Test issue: https://github.com/apache/datafusion/issues/14124 query error DataFusion error: Arrow error: Arithmetic overflow: Overflow happened on: 10000 \* 100000000000000000000000000000000000 SELECT ('0.54321543215432154321543215432154321'::DECIMAL(35,35) + 10000)::VARCHAR + +# where_selection_with_ambiguous_column +statement ok +CREATE TABLE t(a int, b int, id int); + +query error DataFusion error: Schema error: Ambiguous reference to unqualified field id +select * from t a, t b where id = id + 1; + +statement count 0 +drop table t; diff --git a/datafusion/sqllogictest/test_files/simplify_expr.slt b/datafusion/sqllogictest/test_files/simplify_expr.slt new file mode 100644 index 000000000000..d10e603ea5f3 --- /dev/null +++ b/datafusion/sqllogictest/test_files/simplify_expr.slt @@ -0,0 +1,34 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement count 0 +create table t(a int) as values (1); + +# test between simplification +query TT +explain select a from t where a BETWEEN 3 and 3 +---- +logical_plan +01)Filter: t.a = Int32(3) +02)--TableScan: t projection=[a] +physical_plan +01)CoalesceBatchesExec: target_batch_size=8192 +02)--FilterExec: a@0 = 3 +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +statement count 0 +drop table t; diff --git a/datafusion/sqllogictest/test_files/string/string_view.slt b/datafusion/sqllogictest/test_files/string/string_view.slt index 754937e18f14..69c4b9bfcb4b 100644 --- a/datafusion/sqllogictest/test_files/string/string_view.slt +++ b/datafusion/sqllogictest/test_files/string/string_view.slt @@ -660,7 +660,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: contains(test.column1_utf8view, Utf8View("foo")) AS c1, contains(test.column1_utf8view, test.column2_utf8view) AS c2, contains(test.column1_utf8view, CAST(test.column2_large_utf8 AS Utf8View)) AS c3, contains(CAST(test.column1_utf8 AS Utf8View), test.column2_utf8view) AS c4, contains(test.column1_utf8, test.column2_utf8) AS c5, contains(CAST(test.column1_utf8 AS LargeUtf8), test.column2_large_utf8) AS c6, contains(CAST(test.column1_large_utf8 AS Utf8View), test.column1_utf8view) AS c7, contains(test.column1_large_utf8, CAST(test.column2_utf8 AS LargeUtf8)) AS c8, contains(test.column1_large_utf8, test.column2_large_utf8) AS c9 +01)Projection: contains(test.column1_utf8view, Utf8("foo")) AS c1, contains(test.column1_utf8view, test.column2_utf8view) AS c2, contains(test.column1_utf8view, test.column2_large_utf8) AS c3, contains(test.column1_utf8, test.column2_utf8view) AS c4, contains(test.column1_utf8, test.column2_utf8) AS c5, contains(test.column1_utf8, test.column2_large_utf8) AS c6, contains(test.column1_large_utf8, test.column1_utf8view) AS c7, contains(test.column1_large_utf8, test.column2_utf8) AS c8, contains(test.column1_large_utf8, test.column2_large_utf8) AS c9 02)--TableScan: test projection=[column1_utf8, column2_utf8, column1_large_utf8, column2_large_utf8, column1_utf8view, column2_utf8view] ## Ensure no casts for ENDS_WITH @@ -671,7 +671,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: ends_with(test.column1_utf8view, Utf8View("foo")) AS c1, ends_with(test.column2_utf8view, test.column2_utf8view) AS c2 +01)Projection: ends_with(test.column1_utf8view, Utf8("foo")) AS c1, ends_with(test.column2_utf8view, test.column2_utf8view) AS c2 02)--TableScan: test projection=[column1_utf8view, column2_utf8view] ## Ensure no casts for LEVENSHTEIN @@ -682,7 +682,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: levenshtein(test.column1_utf8view, Utf8View("foo")) AS c1, levenshtein(test.column1_utf8view, test.column2_utf8view) AS c2 +01)Projection: levenshtein(test.column1_utf8view, Utf8("foo")) AS c1, levenshtein(test.column1_utf8view, test.column2_utf8view) AS c2 02)--TableScan: test projection=[column1_utf8view, column2_utf8view] ## Ensure no casts for LOWER @@ -784,7 +784,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: regexp_like(test.column1_utf8view, Utf8View("^https?://(?:www\.)?([^/]+)/.*$")) AS k +01)Projection: regexp_like(test.column1_utf8view, Utf8("^https?://(?:www\.)?([^/]+)/.*$")) AS k 02)--TableScan: test projection=[column1_utf8view] ## Ensure no casts for REGEXP_MATCH @@ -825,7 +825,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: replace(test.column1_utf8view, Utf8View("foo"), Utf8View("bar")) AS c1, replace(test.column1_utf8view, test.column2_utf8view, Utf8View("bar")) AS c2 +01)Projection: replace(test.column1_utf8view, Utf8("foo"), Utf8("bar")) AS c1, replace(test.column1_utf8view, test.column2_utf8view, Utf8("bar")) AS c2 02)--TableScan: test projection=[column1_utf8view, column2_utf8view] ## Ensure no casts for REVERSE @@ -906,7 +906,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: strpos(test.column1_utf8view, Utf8View("f")) AS c, strpos(test.column1_utf8view, test.column2_utf8view) AS c2 +01)Projection: strpos(test.column1_utf8view, Utf8("f")) AS c, strpos(test.column1_utf8view, test.column2_utf8view) AS c2 02)--TableScan: test projection=[column1_utf8view, column2_utf8view] ## Ensure no casts for SUBSTR @@ -1078,7 +1078,7 @@ EXPLAIN SELECT FROM test; ---- logical_plan -01)Projection: string_to_array(test.column1_utf8view, Utf8View(",")) AS c +01)Projection: string_to_array(test.column1_utf8view, Utf8(",")) AS c 02)--TableScan: test projection=[column1_utf8view] ## Ensure no unexpected casts for array_to_string diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 94c9eaf810fb..5a722c2288ac 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -1393,3 +1393,61 @@ item1 1970-01-01T00:00:03 75 statement ok drop table source_table; + +statement count 0 +drop table t1; + +statement count 0 +drop table t2; + +statement count 0 +drop table t3; + +# test count wildcard +statement count 0 +create table t1(a int) as values (1); + +statement count 0 +create table t2(b int) as values (1); + +query I +SELECT a FROM t1 WHERE EXISTS (SELECT count(*) FROM t2) +---- +1 + +query TT +explain SELECT a FROM t1 WHERE EXISTS (SELECT count(*) FROM t2) +---- +logical_plan +01)LeftSemi Join: +02)--TableScan: t1 projection=[a] +03)--SubqueryAlias: __correlated_sq_1 +04)----Projection: +05)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +06)--------TableScan: t2 projection=[] + +statement count 0 +drop table t1; + +statement count 0 +drop table t2; + + +# test exists_subquery_wildcard +statement count 0 +create table person(id int, last_name int, state int); + +query TT +explain SELECT id FROM person p WHERE EXISTS + (SELECT * FROM person WHERE last_name = p.last_name AND state = p.state) +---- +logical_plan +01)Projection: p.id +02)--LeftSemi Join: p.last_name = __correlated_sq_1.last_name, p.state = __correlated_sq_1.state +03)----SubqueryAlias: p +04)------TableScan: person projection=[id, last_name, state] +05)----SubqueryAlias: __correlated_sq_1 +06)------TableScan: person projection=[last_name, state] + +statement count 0 +drop table person; diff --git a/datafusion/sqllogictest/test_files/timestamps.slt b/datafusion/sqllogictest/test_files/timestamps.slt index 7eadb3c89dac..dcbcfbfa439d 100644 --- a/datafusion/sqllogictest/test_files/timestamps.slt +++ b/datafusion/sqllogictest/test_files/timestamps.slt @@ -2847,6 +2847,26 @@ SELECT to_char(null, '%d-%m-%Y'); ---- NULL +query T +SELECT to_char(date_column, '%Y-%m-%d') +FROM (VALUES + (DATE '2020-09-01'), + (NULL) +) AS t(date_column); +---- +2020-09-01 +NULL + +query T +SELECT to_char(date_column, '%Y-%m-%d') +FROM (VALUES + (NULL), + (DATE '2020-09-01') +) AS t(date_column); +---- +NULL +2020-09-01 + query T SELECT to_char(column1, column2) FROM diff --git a/datafusion/sqllogictest/test_files/wildcard.slt b/datafusion/sqllogictest/test_files/wildcard.slt index 7c076f040feb..1a480eac0cc3 100644 --- a/datafusion/sqllogictest/test_files/wildcard.slt +++ b/datafusion/sqllogictest/test_files/wildcard.slt @@ -145,3 +145,18 @@ DROP TABLE t2; statement ok DROP TABLE aggregate_simple; + +statement ok +create table t(a int, b int, c int) as values (1, 2, 3); + +query error DataFusion error: Error during planning: Projections require unique expression names but the expression "t\.a" at position 0 and "t\.a" at position 3 have the same name\. Consider aliasing \("AS"\) one of them\. +select *, a from t; + +# a is aliased to other name so the query is valid +query IIII +select *, a as aka from t; +---- +1 2 3 1 + +statement count 0 +drop table t; diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 61bb2f022709..1a9acc0f531a 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -2833,13 +2833,12 @@ logical_plan 06)----------Projection: CAST(annotated_data_infinite.inc_col AS Int64) AS __common_expr_1, annotated_data_infinite.ts, annotated_data_infinite.inc_col 07)------------TableScan: annotated_data_infinite projection=[ts, inc_col] physical_plan -01)ProjectionExec: expr=[sum1@0 as sum1, sum2@1 as sum2, count1@2 as count1, count2@3 as count2] -02)--ProjectionExec: expr=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@5 as sum1, sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@3 as sum2, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@6 as count1, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@4 as count2, ts@1 as ts] -03)----GlobalLimitExec: skip=0, fetch=5 -04)------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }], mode=[Sorted] -05)--------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }], mode=[Sorted] -06)----------ProjectionExec: expr=[CAST(inc_col@1 AS Int64) as __common_expr_1, ts@0 as ts, inc_col@1 as inc_col] -07)------------StreamingTableExec: partition_sizes=1, projection=[ts, inc_col], infinite_source=true, output_ordering=[ts@0 ASC NULLS LAST] +01)ProjectionExec: expr=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@5 as sum1, sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@3 as sum2, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@6 as count1, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@4 as count2] +02)--GlobalLimitExec: skip=0, fetch=5 +03)----BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }], mode=[Sorted] +04)------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }], mode=[Sorted] +05)--------ProjectionExec: expr=[CAST(inc_col@1 AS Int64) as __common_expr_1, ts@0 as ts, inc_col@1 as inc_col] +06)----------StreamingTableExec: partition_sizes=1, projection=[ts, inc_col], infinite_source=true, output_ordering=[ts@0 ASC NULLS LAST] query IIII SELECT @@ -2879,13 +2878,12 @@ logical_plan 06)----------Projection: CAST(annotated_data_infinite.inc_col AS Int64) AS __common_expr_1, annotated_data_infinite.ts, annotated_data_infinite.inc_col 07)------------TableScan: annotated_data_infinite projection=[ts, inc_col] physical_plan -01)ProjectionExec: expr=[sum1@0 as sum1, sum2@1 as sum2, count1@2 as count1, count2@3 as count2] -02)--ProjectionExec: expr=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@5 as sum1, sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@3 as sum2, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@6 as count1, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@4 as count2, ts@1 as ts] -03)----GlobalLimitExec: skip=0, fetch=5 -04)------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }], mode=[Sorted] -05)--------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }], mode=[Sorted] -06)----------ProjectionExec: expr=[CAST(inc_col@1 AS Int64) as __common_expr_1, ts@0 as ts, inc_col@1 as inc_col] -07)------------StreamingTableExec: partition_sizes=1, projection=[ts, inc_col], infinite_source=true, output_ordering=[ts@0 ASC NULLS LAST] +01)ProjectionExec: expr=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@5 as sum1, sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@3 as sum2, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING@6 as count1, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING@4 as count2] +02)--GlobalLimitExec: skip=0, fetch=5 +03)----BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND 1 FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(1)), is_causal: false }], mode=[Sorted] +04)------BoundedWindowAggExec: wdw=[sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }, count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(annotated_data_infinite.inc_col) ORDER BY [annotated_data_infinite.ts DESC NULLS FIRST] ROWS BETWEEN 3 PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(3)), is_causal: false }], mode=[Sorted] +05)--------ProjectionExec: expr=[CAST(inc_col@1 AS Int64) as __common_expr_1, ts@0 as ts, inc_col@1 as inc_col] +06)----------StreamingTableExec: partition_sizes=1, projection=[ts, inc_col], infinite_source=true, output_ordering=[ts@0 ASC NULLS LAST] query IIII diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index 3e3ea7843ac9..dbb7fd65c1a0 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -39,7 +39,7 @@ itertools = { workspace = true } object_store = { workspace = true } pbjson-types = { workspace = true } prost = { workspace = true } -substrait = { version = "0.53", features = ["serde"] } +substrait = { version = "0.54", features = ["serde"] } url = { workspace = true } tokio = { workspace = true, features = ["fs"] } diff --git a/datafusion/substrait/src/logical_plan/consumer.rs b/datafusion/substrait/src/logical_plan/consumer.rs index ffeff3e9df47..60e71ca39d33 100644 --- a/datafusion/substrait/src/logical_plan/consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer.rs @@ -68,8 +68,8 @@ use datafusion::logical_expr::{ }; use datafusion::prelude::{lit, JoinType}; use datafusion::{ - arrow, error::Result, logical_expr::utils::split_conjunction, prelude::Column, - scalar::ScalarValue, + arrow, error::Result, logical_expr::utils::split_conjunction, + logical_expr::utils::split_conjunction_owned, prelude::Column, scalar::ScalarValue, }; use std::collections::HashSet; use std::sync::Arc; @@ -104,10 +104,11 @@ use substrait::proto::{ rel::RelType, rel_common, sort_field::{SortDirection, SortKind::*}, - AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, ExchangeRel, - Expression, ExtendedExpression, ExtensionLeafRel, ExtensionMultiRel, - ExtensionSingleRel, FetchRel, FilterRel, FunctionArgument, JoinRel, NamedStruct, - Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortField, SortRel, Type, + AggregateFunction, AggregateRel, ConsistentPartitionWindowRel, CrossRel, + DynamicParameter, ExchangeRel, Expression, ExtendedExpression, ExtensionLeafRel, + ExtensionMultiRel, ExtensionSingleRel, FetchRel, FilterRel, FunctionArgument, + JoinRel, NamedStruct, Plan, ProjectRel, ReadRel, Rel, RelCommon, SetRel, SortField, + SortRel, Type, }; #[async_trait] @@ -392,6 +393,14 @@ pub trait SubstraitConsumer: Send + Sync + Sized { not_impl_err!("Enum expression not supported") } + async fn consume_dynamic_parameter( + &self, + _expr: &DynamicParameter, + _input_schema: &DFSchema, + ) -> Result { + not_impl_err!("Dynamic Parameter expression not supported") + } + // User-Defined Functionality // The details of extension relations, and how to handle them, are fully up to users to specify. @@ -1327,19 +1336,28 @@ pub async fn from_read_rel( table_ref: TableReference, schema: DFSchema, projection: &Option, + filter: &Option>, ) -> Result { let schema = schema.replace_qualifier(table_ref.clone()); + let filters = if let Some(f) = filter { + let filter_expr = consumer.consume_expression(f, &schema).await?; + split_conjunction_owned(filter_expr) + } else { + vec![] + }; + let plan = { let provider = match consumer.resolve_table_ref(&table_ref).await? { Some(ref provider) => Arc::clone(provider), _ => return plan_err!("No table named '{table_ref}'"), }; - LogicalPlanBuilder::scan( + LogicalPlanBuilder::scan_with_filters( table_ref, provider_as_source(Arc::clone(&provider)), None, + filters, )? .build()? }; @@ -1382,6 +1400,7 @@ pub async fn from_read_rel( table_reference, substrait_schema, &read.projection, + &read.filter, ) .await } @@ -1464,6 +1483,7 @@ pub async fn from_read_rel( table_reference, substrait_schema, &read.projection, + &read.filter, ) .await } @@ -1988,6 +2008,9 @@ pub async fn from_substrait_rex( } RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await, RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await, + RexType::DynamicParameter(expr) => { + consumer.consume_dynamic_parameter(expr, input_schema).await + } }, None => substrait_err!("Expression must set rex_type: {:?}", expression), } diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs index 9dbb246453be..44baf277786d 100644 --- a/datafusion/substrait/src/logical_plan/producer.rs +++ b/datafusion/substrait/src/logical_plan/producer.rs @@ -15,9 +15,6 @@ // specific language governing permissions and limitations // under the License. -use datafusion::config::ConfigOptions; -use datafusion::optimizer::analyzer::expand_wildcard_rule::ExpandWildcardRule; -use datafusion::optimizer::AnalyzerRule; use std::sync::Arc; use substrait::proto::expression_reference::ExprType; @@ -55,6 +52,7 @@ use datafusion::logical_expr::expr::{ AggregateFunctionParams, Alias, BinaryExpr, Case, Cast, GroupingSet, InList, InSubquery, WindowFunction, WindowFunctionParams, }; +use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::{expr, Between, JoinConstraint, LogicalPlan, Operator}; use datafusion::prelude::Expr; use pbjson_types::Any as ProtoAny; @@ -434,14 +432,10 @@ pub fn to_substrait_plan(plan: &LogicalPlan, state: &SessionState) -> Result Result Result> { let projection = scan.projection.as_ref().map(|p| { @@ -560,11 +555,28 @@ pub fn from_table_scan( let table_schema = scan.source.schema().to_dfschema_ref()?; let base_schema = to_substrait_named_struct(&table_schema)?; + let filter_option = if scan.filters.is_empty() { + None + } else { + let table_schema_qualified = Arc::new( + DFSchema::try_from_qualified_schema( + scan.table_name.clone(), + &(scan.source.schema()), + ) + .unwrap(), + ); + + let combined_expr = conjunction(scan.filters.clone()).unwrap(); + let filter_expr = + producer.handle_expr(&combined_expr, &table_schema_qualified)?; + Some(Box::new(filter_expr)) + }; + Ok(Box::new(Rel { rel_type: Some(RelType::Read(Box::new(ReadRel { common: None, base_schema: Some(base_schema), - filter: None, + filter: filter_option, best_effort_filter: None, projection, advanced_extension: None, @@ -1039,7 +1051,7 @@ fn to_substrait_named_struct(schema: &DFSchemaRef) -> Result { .map(|f| to_substrait_type(f.data_type(), f.is_nullable())) .collect::>()?, type_variation_reference: DEFAULT_TYPE_VARIATION_REF, - nullability: r#type::Nullability::Unspecified as i32, + nullability: r#type::Nullability::Required as i32, }; Ok(NamedStruct { @@ -1366,6 +1378,7 @@ pub fn to_substrait_rex( Expr::ScalarSubquery(expr) => { not_impl_err!("Cannot convert {expr:?} to Substrait") } + #[expect(deprecated)] Expr::Wildcard { .. } => not_impl_err!("Cannot convert {expr:?} to Substrait"), Expr::GroupingSet(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"), Expr::Placeholder(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"), diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index e6b8bdbc047e..f989d05c80dd 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1234,6 +1234,11 @@ async fn roundtrip_repartition_hash() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_read_filter() -> Result<()> { + roundtrip_verify_read_filter_count("SELECT a FROM data where a < 5", 1).await +} + fn check_post_join_filters(rel: &Rel) -> Result<()> { // search for target_rel and field value in proto match &rel.rel_type { @@ -1319,6 +1324,56 @@ async fn verify_post_join_filter_value(proto: Box) -> Result<()> { Ok(()) } +fn count_read_filters(rel: &Rel, filter_count: &mut u32) -> Result<()> { + // search for target_rel and field value in proto + match &rel.rel_type { + Some(RelType::Read(read)) => { + // increment counter for read filter if not None + if read.filter.is_some() { + *filter_count += 1; + } + Ok(()) + } + Some(RelType::Filter(filter)) => { + count_read_filters(filter.input.as_ref().unwrap().as_ref(), filter_count) + } + _ => Ok(()), + } +} + +async fn assert_read_filter_count( + proto: Box, + expected_filter_count: u32, +) -> Result<()> { + let mut filter_count: u32 = 0; + for relation in &proto.relations { + match relation.rel_type.as_ref() { + Some(rt) => match rt { + plan_rel::RelType::Rel(rel) => { + match count_read_filters(rel, &mut filter_count) { + Err(e) => return Err(e), + Ok(_) => continue, + } + } + plan_rel::RelType::Root(root) => { + match count_read_filters( + root.input.as_ref().unwrap(), + &mut filter_count, + ) { + Err(e) => return Err(e), + Ok(_) => continue, + } + } + }, + None => return plan_err!("Cannot parse plan relation: None"), + } + } + + assert_eq!(expected_filter_count, filter_count); + + Ok(()) +} + async fn assert_expected_plan_unoptimized( sql: &str, expected_plan_str: &str, @@ -1489,6 +1544,17 @@ async fn roundtrip_verify_post_join_filter(sql: &str) -> Result<()> { verify_post_join_filter_value(proto).await } +async fn roundtrip_verify_read_filter_count( + sql: &str, + expected_filter_count: u32, +) -> Result<()> { + let ctx = create_context().await?; + let proto = roundtrip_with_ctx(sql, ctx).await?; + + // verify that filter counts in read relations are as expected + assert_read_filter_count(proto, expected_filter_count).await +} + async fn roundtrip_all_types(sql: &str) -> Result<()> { roundtrip_with_ctx(sql, create_all_type_context().await?).await?; Ok(()) diff --git a/datafusion/substrait/tests/utils.rs b/datafusion/substrait/tests/utils.rs index 0034cc27bf6e..e3e3ec3fab01 100644 --- a/datafusion/substrait/tests/utils.rs +++ b/datafusion/substrait/tests/utils.rs @@ -480,6 +480,7 @@ pub mod test { } } } + RexType::DynamicParameter(_) => {} // Enum is deprecated RexType::Enum(_) => {} } diff --git a/datafusion/wasmtest/Cargo.toml b/datafusion/wasmtest/Cargo.toml index 30d5bcaedcb7..94515c6754a7 100644 --- a/datafusion/wasmtest/Cargo.toml +++ b/datafusion/wasmtest/Cargo.toml @@ -45,7 +45,7 @@ chrono = { version = "0.4", features = ["wasmbind"] } # all the `std::fmt` and `std::panicking` infrastructure, so isn't great for # code size when deploying. console_error_panic_hook = { version = "0.1.1", optional = true } -datafusion = { workspace = true } +datafusion = { workspace = true, features = ["parquet"] } datafusion-common = { workspace = true, default-features = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } diff --git a/dev/changelog/46.0.0.md b/dev/changelog/46.0.0.md new file mode 100644 index 000000000000..3734161e032f --- /dev/null +++ b/dev/changelog/46.0.0.md @@ -0,0 +1,421 @@ + + +# Apache DataFusion 46.0.0 Changelog + +This release consists of 288 commits from 79 contributors. See credits at the end of this changelog for more information. + +Please see the [Upgrade Guide] for help updating to DataFusion `46.0.0` + +[upgrade guide]: https://datafusion.apache.org/library-user-guide/upgrading.html#datafusion-46-0-0 + +**Breaking changes:** + +- bug: Fix NULL handling in array_slice, introduce `NullHandling` enum to `Signature` [#14289](https://github.com/apache/datafusion/pull/14289) (jkosh44) +- Update REGEXP_MATCH scalar function to support Utf8View [#14449](https://github.com/apache/datafusion/pull/14449) (Omega359) +- Introduce unified `DataSourceExec` for provided datasources, remove `ParquetExec`, `CsvExec`, etc [#14224](https://github.com/apache/datafusion/pull/14224) (mertak-synnada) +- Fix: Avoid recursive external error wrapping [#14371](https://github.com/apache/datafusion/pull/14371) (getChan) +- Add `DataFusionError::Collection` to return multiple `DataFusionError`s [#14439](https://github.com/apache/datafusion/pull/14439) (eliaperantoni) +- function: Allow more expressive array signatures [#14532](https://github.com/apache/datafusion/pull/14532) (jkosh44) +- feat: add resolved `target` to `DmlStatement` (to eliminate need for table lookup after deserialization) [#14631](https://github.com/apache/datafusion/pull/14631) (milenkovicm) +- Signature::Coercible with user defined implicit casting [#14440](https://github.com/apache/datafusion/pull/14440) (jayzhan211) +- Remove CountWildcardRule in Analyzer and move the functionality in ExprPlanner, add `plan_aggregate` and `plan_window` to planner [#14689](https://github.com/apache/datafusion/pull/14689) (jayzhan211) +- Simplify `FileSource::create_file_opener`'s signature [#14798](https://github.com/apache/datafusion/pull/14798) (AdamGS) +- StatisticsV2: initial statistics framework redesign [#14699](https://github.com/apache/datafusion/pull/14699) (Fly-Style) +- fix(substrait): Do not add implicit groupBy expressions in `LogicalPlanBuilder` or when building logical plans from Substrait [#14860](https://github.com/apache/datafusion/pull/14860) (anlinc) + +**Performance related:** + +- perf: Improve `median` with no grouping by 2X [#14399](https://github.com/apache/datafusion/pull/14399) (2010YOUY01) +- Improve performance 10%-100% in `FIRST_VALUE` / `LAST_VALUE` by not sort rows in `FirstValueAccumulator` [#14402](https://github.com/apache/datafusion/pull/14402) (blaginin) +- Speed up `uuid` UDF (40x faster) [#14675](https://github.com/apache/datafusion/pull/14675) (simonvandel) +- script to export benchmark information as Line Protocol format [#14662](https://github.com/apache/datafusion/pull/14662) (logan-keede) +- perf: Drop RowConverter from GroupOrderingPartial [#14566](https://github.com/apache/datafusion/pull/14566) (ctsk) +- Speedup `to_hex` (~2x faster) [#14686](https://github.com/apache/datafusion/pull/14686) (simonvandel) + +**Implemented enhancements:** + +- feat: Speed up `struct` and `named_struct` using `invoke_with_args` [#14276](https://github.com/apache/datafusion/pull/14276) (pepijnve) +- feat: add hint for missing fields [#14521](https://github.com/apache/datafusion/pull/14521) (Lordworms) +- feat: Add support for --mem-pool-type and --memory-limit options to multiple benchmarks [#14642](https://github.com/apache/datafusion/pull/14642) (Kontinuation) +- feat: Implement UNION ALL BY NAME [#14538](https://github.com/apache/datafusion/pull/14538) (rkrishn7) +- feat: Add ScalarUDF support in FFI crate [#14579](https://github.com/apache/datafusion/pull/14579) (timsaucer) +- feat: Improve datafusion-cli memory usage and considering reserve mem… [#14766](https://github.com/apache/datafusion/pull/14766) (zhuqi-lucas) + +**Fixed bugs:** + +- fix: Limits are not applied correctly [#14418](https://github.com/apache/datafusion/pull/14418) (zhuqi-lucas) +- fix(ci): build error with wasm [#14494](https://github.com/apache/datafusion/pull/14494) (Lordworms) +- fix(doc): remove AWS_PROFILE from supported S3 configuration [#14492](https://github.com/apache/datafusion/pull/14492) (hussein-awala) +- fix: `List` of `FixedSizeList` coercion issue in SQL [#14468](https://github.com/apache/datafusion/pull/14468) (alan910127) +- fix: order by expr rewrite fix [#14486](https://github.com/apache/datafusion/pull/14486) (akoshchiy) +- fix: rewrite fetch, skip of the Limit node in correct order [#14496](https://github.com/apache/datafusion/pull/14496) (evenyag) +- fix: Capture nullability in `Values` node planning [#14472](https://github.com/apache/datafusion/pull/14472) (rkrishn7) +- fix: case-sensitive quoted identifiers in DELETE statements [#14584](https://github.com/apache/datafusion/pull/14584) (nantunes) +- fix: Substrait serializer clippy error: not calling truncate [#14723](https://github.com/apache/datafusion/pull/14723) (niebayes) +- fix: normalize column names in table constraints [#14794](https://github.com/apache/datafusion/pull/14794) (jonahgao) +- fix: we are missing the unlimited case for bounded streaming when usi… [#14815](https://github.com/apache/datafusion/pull/14815) (zhuqi-lucas) +- fix: use `return_type_from_args` and mark nullable if any of the input is nullable [#14841](https://github.com/apache/datafusion/pull/14841) (rluvaton) + +**Documentation updates:** + +- Add related source code locations to errors [#13664](https://github.com/apache/datafusion/pull/13664) (eliaperantoni) +- docs: Fix create_udf examples [#14405](https://github.com/apache/datafusion/pull/14405) (nuno-faria) +- Script and documentation for regenerating sqlite test files [#14290](https://github.com/apache/datafusion/pull/14290) (Omega359) +- Improve documentation about extended tests [#14320](https://github.com/apache/datafusion/pull/14320) (alamb) +- Add `Cargo.lock` [#14483](https://github.com/apache/datafusion/pull/14483) (mbrobbel) +- Test all examples from library-user-guide & user-guide docs [#14544](https://github.com/apache/datafusion/pull/14544) (ugoa) +- Add guideline for GSoC 2025 applicants under Contributor Guide [#14582](https://github.com/apache/datafusion/pull/14582) (oznur-synnada) +- Fix typo in comments [#14605](https://github.com/apache/datafusion/pull/14605) (byte-sourcerer) +- Update Community Events in concepts-readings-events.md [#14629](https://github.com/apache/datafusion/pull/14629) (oznur-synnada) +- Minor: Add docs and examples for `DataFusionErrorBuilder` [#14551](https://github.com/apache/datafusion/pull/14551) (alamb) +- docs: Add Sleeper to list of known users [#14648](https://github.com/apache/datafusion/pull/14648) (m09526) +- Add documentation for prepare statements. [#14639](https://github.com/apache/datafusion/pull/14639) (dhegberg) +- Update features / status documentation page [#14645](https://github.com/apache/datafusion/pull/14645) (alamb) +- Add union_extract scalar function [#12116](https://github.com/apache/datafusion/pull/12116) (gstvg) +- Fix CI doctests on main [#14667](https://github.com/apache/datafusion/pull/14667) (alamb) +- chore: adding Linkedin follow page [#14676](https://github.com/apache/datafusion/pull/14676) (comphead) +- Improve EnforceSorting docs. [#14673](https://github.com/apache/datafusion/pull/14673) (wiedld) +- Specify rust toolchain explicitly, document how to change it [#14655](https://github.com/apache/datafusion/pull/14655) (alamb) +- Create gsoc_project_ideas.md [#14774](https://github.com/apache/datafusion/pull/14774) (oznur-synnada) +- docs: Add additional info about memory reservation to the doc of MemoryPool [#14789](https://github.com/apache/datafusion/pull/14789) (Kontinuation) +- docs: Add instruction to build [#14694](https://github.com/apache/datafusion/pull/14694) (dentiny) +- Update website links [#14846](https://github.com/apache/datafusion/pull/14846) (oznur-synnada) +- Improve benchmark docs [#14820](https://github.com/apache/datafusion/pull/14820) (carols10cents) +- Add polygon.io to user list [#14871](https://github.com/apache/datafusion/pull/14871) (xudong963) +- Update dft in intro "Known Users" [#14875](https://github.com/apache/datafusion/pull/14875) (matthewmturner) +- Add `statistics_truncate_length` parquet writer config [#14782](https://github.com/apache/datafusion/pull/14782) (akoshchiy) +- minor: Update docs and error messages about what SQL dialects are supported [#14893](https://github.com/apache/datafusion/pull/14893) (AdamGS) +- Minor: Add Development Environment to Documentation Index [#14890](https://github.com/apache/datafusion/pull/14890) (alamb) +- Examples: boundary analysis example for `AND/OR` conjunctions [#14735](https://github.com/apache/datafusion/pull/14735) (clflushopt) +- Allow setting the recursion limit for sql parsing [#14756](https://github.com/apache/datafusion/pull/14756) (cetra3) +- Document SQL literal syntax and escaping [#14934](https://github.com/apache/datafusion/pull/14934) (alamb) +- Prepare for 46.0.0 release: Version and Changelog [#14903](https://github.com/apache/datafusion/pull/14903) (xudong963) +- MINOR fix(docs): set the proper link for dev-env setup in contrib guide [#14960](https://github.com/apache/datafusion/pull/14960) (clflushopt) + +**Other:** + +- Fix join type coercion when joining 2 relations with the same name via `DataFrame` API [#14387](https://github.com/apache/datafusion/pull/14387) (alamb) +- Minor: fix typo in test name [#14403](https://github.com/apache/datafusion/pull/14403) (alamb) +- test: add regression test for unnesting dictionary encoded columns [#14395](https://github.com/apache/datafusion/pull/14395) (duongcongtoai) +- chore: Upgrade to `arrow`/`parquet` `54.1.0` and fix clippy/ci [#14415](https://github.com/apache/datafusion/pull/14415) (Weijun-H) +- bump arrow version to 54.1.0 and fix clippy error [#14414](https://github.com/apache/datafusion/pull/14414) (Lordworms) +- Support `array_concat` for `Utf8View` [#14378](https://github.com/apache/datafusion/pull/14378) (alamb) +- Fully support LIKE/ILIKE with Utf8View [#14379](https://github.com/apache/datafusion/pull/14379) (alamb) +- Fix `null` input in `map_keys/values` [#14401](https://github.com/apache/datafusion/pull/14401) (cht42) +- Remove dependency on datafusion_catalog from datafusion-cli [#14398](https://github.com/apache/datafusion/pull/14398) (alamb) +- chore(deps): update substrait requirement from 0.52 to 0.53 [#14419](https://github.com/apache/datafusion/pull/14419) (dependabot[bot]) +- move resolve_table_references`out of`datafusion-catalog` [#14441](https://github.com/apache/datafusion/pull/14441) (logan-keede) +- Fix a clippy warning [#14445](https://github.com/apache/datafusion/pull/14445) (mbrobbel) +- Resolve a todo about using workspace dependencies [#14443](https://github.com/apache/datafusion/pull/14443) (mbrobbel) +- Support `Utf8View` to `numeric` coercion [#14377](https://github.com/apache/datafusion/pull/14377) (alamb) +- Fix regression list Type Coercion List with inner type struct which has large/view types [#14385](https://github.com/apache/datafusion/pull/14385) (alamb) +- Improve error message on unsupported correlation [#14458](https://github.com/apache/datafusion/pull/14458) (findepi) +- Replace `once_cell::Lazy` with `std::sync::LazyLock` [#14480](https://github.com/apache/datafusion/pull/14480) (mbrobbel) +- chore(deps): bump bytes from 1.9.0 to 1.10.0 in /datafusion-cli [#14476](https://github.com/apache/datafusion/pull/14476) (dependabot[bot]) +- chore(deps): bump clap from 4.5.27 to 4.5.28 in /datafusion-cli [#14477](https://github.com/apache/datafusion/pull/14477) (dependabot[bot]) +- chore: Fix link to issue and expand comment [#14473](https://github.com/apache/datafusion/pull/14473) (findepi) +- Make Pushdown Filters Public [#14471](https://github.com/apache/datafusion/pull/14471) (cetra3) +- Minor: `cargo fmt` to fix CI [#14487](https://github.com/apache/datafusion/pull/14487) (alamb) +- chore: clean up dependencies for datafusion cli [#14484](https://github.com/apache/datafusion/pull/14484) (comphead) +- Provide user-defined invariants for logical node extensions. [#14329](https://github.com/apache/datafusion/pull/14329) (wiedld) +- DFParser should skip unsupported COPY INTO [#14382](https://github.com/apache/datafusion/pull/14382) (osipovartem) +- Improve Unparser (scalar_to_sql) to respect dialect timestamp type overrides [#14407](https://github.com/apache/datafusion/pull/14407) (sgrebnov) +- Fix link to volcano parallelism paper [#14497](https://github.com/apache/datafusion/pull/14497) (lewiszlw) +- chore(deps): bump aws-config from 1.5.15 to 1.5.16 in /datafusion-cli [#14500](https://github.com/apache/datafusion/pull/14500) (dependabot[bot]) +- chore: Add more LIKE with escape tests [#14501](https://github.com/apache/datafusion/pull/14501) (findepi) +- Fix a clippy warning in `datafusion-sqllogictest` [#14506](https://github.com/apache/datafusion/pull/14506) (mbrobbel) +- minor: improve PR template [#14507](https://github.com/apache/datafusion/pull/14507) (alamb) +- Support `Dictionary` and `List` types in `scalar_to_sql` [#14346](https://github.com/apache/datafusion/pull/14346) (cetra3) +- Serialize `parquet_options` in `datafusion-proto` [#14465](https://github.com/apache/datafusion/pull/14465) (blaginin) +- make datafusion-catalog-listing and move some implementation of listing out of datafusion/core/datasource/listing [#14464](https://github.com/apache/datafusion/pull/14464) (logan-keede) +- refactor: remove uses of `arrow_buffer` & `arrow_array` and use reexport in arrow instead [#14503](https://github.com/apache/datafusion/pull/14503) (Chen-Yuan-Lai) +- core: Support uncorrelated EXISTS [#14474](https://github.com/apache/datafusion/pull/14474) (findepi) +- chore(deps): Update sqlparser to `0.54.0` [#14255](https://github.com/apache/datafusion/pull/14255) (alamb) +- Validate and unpack function arguments tersely [#14513](https://github.com/apache/datafusion/pull/14513) (findepi) +- bug: Fix edge cases in array_slice [#14489](https://github.com/apache/datafusion/pull/14489) (jkosh44) +- Feat: Add fetch to CoalescePartitionsExec [#14499](https://github.com/apache/datafusion/pull/14499) (mertak-synnada) +- Improve error messages to include the function name. [#14511](https://github.com/apache/datafusion/pull/14511) (Omega359) +- to_unixtime does not support timestamps with a timezone [#14490](https://github.com/apache/datafusion/pull/14490) (Omega359) +- bug: Remove array_slice two arg variant [#14527](https://github.com/apache/datafusion/pull/14527) (jkosh44) +- Minor: deprecate unused index mod [#14534](https://github.com/apache/datafusion/pull/14534) (zhuqi-lucas) +- Fix config_namespace macro symbol usage [#14520](https://github.com/apache/datafusion/pull/14520) (davisp) +- functions: Remove NullHandling from scalar funcs [#14531](https://github.com/apache/datafusion/pull/14531) (jkosh44) +- Relax physical schema validation [#14519](https://github.com/apache/datafusion/pull/14519) (findepi) +- Minor: Update changelog for `45.0.0` and tweak `CHANGELOG` docs [#14545](https://github.com/apache/datafusion/pull/14545) (alamb) +- minor: polish MemoryStream related code [#14537](https://github.com/apache/datafusion/pull/14537) (zjregee) +- refactor: switch BooleanBufferBuilder to NullBufferBuilder in MaybeNullBufferBuilder [#14504](https://github.com/apache/datafusion/pull/14504) (Chen-Yuan-Lai) +- Allow constructing ScalarUDF from shared implementation [#14541](https://github.com/apache/datafusion/pull/14541) (findepi) +- some dependency removals and setup for refactor of `FileScanConfig` [#14543](https://github.com/apache/datafusion/pull/14543) (logan-keede) +- Always use `StringViewArray` as output of `substr` [#14498](https://github.com/apache/datafusion/pull/14498) (Kev1n8) +- refactor: remove remaining uses of `arrow_array` and use reexport in `arrow` instead [#14528](https://github.com/apache/datafusion/pull/14528) (Chen-Yuan-Lai) +- chore: update datafusion-testing pin to fix extended tests [#14556](https://github.com/apache/datafusion/pull/14556) (alamb) +- chore: remove partition_keys from (Bounded)WindowAggExec [#14526](https://github.com/apache/datafusion/pull/14526) (irenjj) +- chore(deps): bump nix from 0.28.0 to 0.29.0 [#14559](https://github.com/apache/datafusion/pull/14559) (dependabot[bot]) +- use a single row_count column during predicate pruning instead of one per column [#14295](https://github.com/apache/datafusion/pull/14295) (adriangb) +- Update proto to support to/from json with an extension codec [#14561](https://github.com/apache/datafusion/pull/14561) (Omega359) +- Remove useless test util [#14570](https://github.com/apache/datafusion/pull/14570) (xudong963) +- minor: Move file compression to `datafusion-catalog-listing` [#14555](https://github.com/apache/datafusion/pull/14555) (logan-keede) +- chore(deps): bump strum from 0.26.3 to 0.27.0 [#14573](https://github.com/apache/datafusion/pull/14573) (dependabot[bot]) +- Minor: remove unnecessary dependencies in `datafusion-sqllogictest` [#14578](https://github.com/apache/datafusion/pull/14578) (alamb) +- Fix: limit is missing after removing SPM [#14569](https://github.com/apache/datafusion/pull/14569) (xudong963) +- Adding cargo clean at the end of every step [#14592](https://github.com/apache/datafusion/pull/14592) (Omega359) +- Make it easier to create a ScalarValure representing typed null (#14548) [#14558](https://github.com/apache/datafusion/pull/14558) (cj-zhukov) +- chore(deps): bump substrait from 0.53.0 to 0.53.1 [#14599](https://github.com/apache/datafusion/pull/14599) (dependabot[bot]) +- refactor: remove uses of arrow_schema and use reexport in arrow instead [#14597](https://github.com/apache/datafusion/pull/14597) (Chen-Yuan-Lai) +- Benchmark showcasing with_column and with_column_renamed function performance [#14564](https://github.com/apache/datafusion/pull/14564) (Omega359) +- Remove use of deprecated dict_id in datafusion-proto (#14173) [#14227](https://github.com/apache/datafusion/pull/14227) (cj-zhukov) +- refactor: Move FileSinkConfig out of Core [#14585](https://github.com/apache/datafusion/pull/14585) (logan-keede) +- Revert modification of build dependency [#14606](https://github.com/apache/datafusion/pull/14606) (ugoa) +- chore(deps): bump serialize-javascript and copy-webpack-plugin in /datafusion/wasmtest/datafusion-wasm-app [#14594](https://github.com/apache/datafusion/pull/14594) (dependabot[bot]) +- cli: Add nested expressions [#14614](https://github.com/apache/datafusion/pull/14614) (jkosh44) +- Minor: remove some unnecessary dependencies [#14615](https://github.com/apache/datafusion/pull/14615) (logan-keede) +- Disable extended tests (`extended_tests`) that are failing on runner [#14604](https://github.com/apache/datafusion/pull/14604) (alamb) +- minor: check size overflow before string repeat build [#14575](https://github.com/apache/datafusion/pull/14575) (wForget) +- Speedup `date_trunc` (~20% time reduction) [#14593](https://github.com/apache/datafusion/pull/14593) (simonvandel) +- equivalence classes: use normalized mapping for projection [#14327](https://github.com/apache/datafusion/pull/14327) (askalt) +- chore(deps): bump prost-build from 0.13.4 to 0.13.5 [#14623](https://github.com/apache/datafusion/pull/14623) (dependabot[bot]) +- chore(deps): bump bzip2 from 0.5.0 to 0.5.1 [#14620](https://github.com/apache/datafusion/pull/14620) (dependabot[bot]) +- chore(deps): bump clap from 4.5.28 to 4.5.29 [#14619](https://github.com/apache/datafusion/pull/14619) (dependabot[bot]) +- chore(deps): bump sqllogictest from 0.26.4 to 0.27.0 [#14598](https://github.com/apache/datafusion/pull/14598) (dependabot[bot]) +- Fix ci test [#14625](https://github.com/apache/datafusion/pull/14625) (xudong963) +- chore(deps): group `prost` and `pbjson` dependabot updates [#14626](https://github.com/apache/datafusion/pull/14626) (mbrobbel) +- chore(deps): bump substrait from 0.53.1 to 0.53.2 [#14627](https://github.com/apache/datafusion/pull/14627) (dependabot[bot]) +- refactor: Move various parts of datasource out of core [#14616](https://github.com/apache/datafusion/pull/14616) (logan-keede) +- Use ` take_function_args` in more places [#14525](https://github.com/apache/datafusion/pull/14525) (lgingerich) +- Minor: remove unused `AutoFinishBzEncoder` [#14630](https://github.com/apache/datafusion/pull/14630) (jonahgao) +- Add test for nullable doesn't work when create memory table [#14624](https://github.com/apache/datafusion/pull/14624) (xudong963) +- Fallback to Utf8View for `Dict(_, Utf8View)` in `type_union_resolution_coercion` [#14602](https://github.com/apache/datafusion/pull/14602) (jayzhan211) +- refactor: Make catalog datasource [#14643](https://github.com/apache/datafusion/pull/14643) (logan-keede) +- Implement predicate pruning for not like expressions [#14567](https://github.com/apache/datafusion/pull/14567) (UBarney) +- Migrate math functions to implement invoke_with_args [#14658](https://github.com/apache/datafusion/pull/14658) (lewiszlw) +- bug: fix offset type mismatch when prepending lists [#14672](https://github.com/apache/datafusion/pull/14672) (friendlymatthew) +- Minor: remove confusing `update_plan_from_children` call from `EnforceSorting` [#14650](https://github.com/apache/datafusion/pull/14650) (xudong963) +- Improve UX Rename `FileScanConfig::new_exec` to `FileScanConfig::build` [#14670](https://github.com/apache/datafusion/pull/14670) (alamb) +- Consolidate and expand ident normalization tests [#14374](https://github.com/apache/datafusion/pull/14374) (alamb) +- Update GitHub CI run image for license check job [#14674](https://github.com/apache/datafusion/pull/14674) (findepi) +- Allow extensions_options to accept Option field [#14664](https://github.com/apache/datafusion/pull/14664) (goldmedal) +- Minor: Re-export `datafusion_expr_common` crate [#14696](https://github.com/apache/datafusion/pull/14696) (jayzhan211) +- Migrate the internal and testing functions to invoke_with_args [#14693](https://github.com/apache/datafusion/pull/14693) (goldmedal) +- Improve docs `TableSource` and `DefaultTableSource` [#14665](https://github.com/apache/datafusion/pull/14665) (alamb) +- Improve SQL Planner docs [#14669](https://github.com/apache/datafusion/pull/14669) (alamb) +- MIgrate math function macro to implement invoke_with_args [#14690](https://github.com/apache/datafusion/pull/14690) (goldmedal) +- Improve `downcast_value!` macro [#14683](https://github.com/apache/datafusion/pull/14683) (findepi) +- chore(deps): bump tempfile from 3.16.0 to 3.17.0 [#14713](https://github.com/apache/datafusion/pull/14713) (dependabot[bot]) +- bug: improve schema checking for `insert into` cases [#14572](https://github.com/apache/datafusion/pull/14572) (zhuqi-lucas) +- Early exit on column normalisation to improve DataFrame performance [#14636](https://github.com/apache/datafusion/pull/14636) (blaginin) +- Add example for `LogicalPlanBuilder::insert_into` [#14663](https://github.com/apache/datafusion/pull/14663) (alamb) +- optimize performance of the repeat function (up to 50% faster) [#14697](https://github.com/apache/datafusion/pull/14697) (zjregee) +- `AggregateUDFImpl::schema_name` and `AggregateUDFImpl::display_name` for customizable name [#14695](https://github.com/apache/datafusion/pull/14695) (jayzhan211) +- Add an example of boundary analysis simple expressions. [#14688](https://github.com/apache/datafusion/pull/14688) (clflushopt) +- chore(deps): bump arrow-ipc from 54.1.0 to 54.2.0 [#14719](https://github.com/apache/datafusion/pull/14719) (dependabot[bot]) +- chore(deps): bump strum from 0.27.0 to 0.27.1 [#14718](https://github.com/apache/datafusion/pull/14718) (dependabot[bot]) +- minor: enable decimal dictionary sbbf pruning test [#14711](https://github.com/apache/datafusion/pull/14711) (korowa) +- chore(deps): bump sqllogictest from 0.27.0 to 0.27.1 [#14717](https://github.com/apache/datafusion/pull/14717) (dependabot[bot]) +- minor: simplify `union_extract` code [#14640](https://github.com/apache/datafusion/pull/14640) (alamb) +- make DefaultSubstraitProducer public [#14721](https://github.com/apache/datafusion/pull/14721) (gabotechs) +- chore: Migrate Encoding functions to invoke_with_args [#14727](https://github.com/apache/datafusion/pull/14727) (irenjj) +- chore: Migrate Core Functions to invoke_with_args [#14725](https://github.com/apache/datafusion/pull/14725) (niebayes) +- Fix off by 1 in decimal cast to lower precision [#14731](https://github.com/apache/datafusion/pull/14731) (findepi) +- migrate string functions to `inovke_with_args` [#14722](https://github.com/apache/datafusion/pull/14722) (zjregee) +- chore: Migrate Array Functions to invoke_with_args [#14726](https://github.com/apache/datafusion/pull/14726) (irenjj) +- chore: Migrate Regex function to invoke_with_args [#14728](https://github.com/apache/datafusion/pull/14728) (irenjj) +- bug: Fix memory reservation and allocation problems for SortExec [#14644](https://github.com/apache/datafusion/pull/14644) (Kontinuation) +- Skip target in taplo checks [#14747](https://github.com/apache/datafusion/pull/14747) (findepi) +- chore(deps): bump uuid from 1.13.1 to 1.13.2 [#14739](https://github.com/apache/datafusion/pull/14739) (dependabot[bot]) +- chore(deps): bump blake3 from 1.5.5 to 1.6.0 [#14741](https://github.com/apache/datafusion/pull/14741) (dependabot[bot]) +- chore(deps): bump tempfile from 3.17.0 to 3.17.1 [#14742](https://github.com/apache/datafusion/pull/14742) (dependabot[bot]) +- chore(deps): bump clap from 4.5.29 to 4.5.30 [#14743](https://github.com/apache/datafusion/pull/14743) (dependabot[bot]) +- chore(deps): bump parquet from 54.1.0 to 54.2.0 [#14744](https://github.com/apache/datafusion/pull/14744) (dependabot[bot]) +- Speed up `chr` UDF (~4x faster) [#14700](https://github.com/apache/datafusion/pull/14700) (simonvandel) +- Support aliases in ConstEvaluator [#14734](https://github.com/apache/datafusion/pull/14734) (joroKr21) +- `AggregateUDFImpl::window_function_schema_name` and `AggregateUDFImpl::window_function_display_name` for window aggregate function [#14750](https://github.com/apache/datafusion/pull/14750) (jayzhan211) +- chore: migrate crypto functions to invoke_with_args [#14764](https://github.com/apache/datafusion/pull/14764) (Chen-Yuan-Lai) +- minor: remove custom extract_ok! macro [#14733](https://github.com/apache/datafusion/pull/14733) (ctsk) +- Minor: Further Clean-up in Enforce Sorting [#14732](https://github.com/apache/datafusion/pull/14732) (berkaysynnada) +- chore(deps): bump arrow-flight from 54.1.0 to 54.2.0 [#14786](https://github.com/apache/datafusion/pull/14786) (dependabot[bot]) +- chore(deps): bump serde_json from 1.0.138 to 1.0.139 [#14784](https://github.com/apache/datafusion/pull/14784) (dependabot[bot]) +- dependabot: group arrow/parquet minor/patch bumps, remove limit [#14730](https://github.com/apache/datafusion/pull/14730) (mbrobbel) +- Map access supports constant-resolvable expressions [#14712](https://github.com/apache/datafusion/pull/14712) (Lordworms) +- Fix build after logical conflict [#14791](https://github.com/apache/datafusion/pull/14791) (alamb) +- Fix CI job test-datafusion-pyarrow [#14790](https://github.com/apache/datafusion/pull/14790) (Owen-CH-Leung) +- Use `doc_auto_cfg`, logo and favicon for docs.rs [#14746](https://github.com/apache/datafusion/pull/14746) (mbrobbel) +- chore(deps): bump sqllogictest from 0.27.1 to 0.27.2 [#14785](https://github.com/apache/datafusion/pull/14785) (dependabot[bot]) +- Fix CI fail for extended test (by freeing up more disk space in CI runner) [#14745](https://github.com/apache/datafusion/pull/14745) (2010YOUY01) +- chore: Benchmark deps cleanup [#14793](https://github.com/apache/datafusion/pull/14793) (findepi) +- chore: Fix test not to litter in repository [#14795](https://github.com/apache/datafusion/pull/14795) (findepi) +- chore(deps): bump testcontainers from 0.23.2 to 0.23.3 [#14787](https://github.com/apache/datafusion/pull/14787) (dependabot[bot]) +- chore(deps): bump serde from 1.0.217 to 1.0.218 [#14788](https://github.com/apache/datafusion/pull/14788) (dependabot[bot]) +- refactor: move `DataSource` to `datafusion-datasource` [#14671](https://github.com/apache/datafusion/pull/14671) (logan-keede) +- Fix Clippy 1.85 warnings [#14800](https://github.com/apache/datafusion/pull/14800) (mbrobbel) +- Allow `FileSource`-specific repartitioning [#14754](https://github.com/apache/datafusion/pull/14754) (AdamGS) +- Bump MSRV to 1.82, toolchain to 1.85 [#14811](https://github.com/apache/datafusion/pull/14811) (mbrobbel) +- Chore/Add additional FFI unit tests [#14802](https://github.com/apache/datafusion/pull/14802) (timsaucer) +- Minor: comment in Cargo.toml about MSRV [#14809](https://github.com/apache/datafusion/pull/14809) (alamb) +- Remove unused crate dependencies [#14827](https://github.com/apache/datafusion/pull/14827) (findepi) +- fix(physical-expr): Remove empty constants check when ordering is satisfied [#14829](https://github.com/apache/datafusion/pull/14829) (rkrishn7) +- chore(deps): bump log from 0.4.25 to 0.4.26 [#14847](https://github.com/apache/datafusion/pull/14847) (dependabot[bot]) +- Minor: Ignore examples output directory [#14840](https://github.com/apache/datafusion/pull/14840) (AdamGS) +- Add support for `Dictionary` to AST datatype in unparser [#14783](https://github.com/apache/datafusion/pull/14783) (cetra3) +- Add `range` table function [#14830](https://github.com/apache/datafusion/pull/14830) (simonvandel) +- chore: migrate invoke_batch to invoke_with_args for unicode function [#14856](https://github.com/apache/datafusion/pull/14856) (onlyjackfrost) +- test: change test_function macro to use `return_type_from_args` instead of `return_type` [#14852](https://github.com/apache/datafusion/pull/14852) (rluvaton) +- Move `FileSourceConfig` and `FileStream` to the new `datafusion-datasource` [#14838](https://github.com/apache/datafusion/pull/14838) (AdamGS) +- Minor: Counting elapsed_compute in BoundedWindowAggExec [#14869](https://github.com/apache/datafusion/pull/14869) (2010YOUY01) +- Optimize `gcd` for array and scalar case by avoiding `make_scalar_function` where has unnecessary conversion between scalar and array [#14834](https://github.com/apache/datafusion/pull/14834) (jayzhan211) +- refactor: replace OnceLock with LazyLock [#14870](https://github.com/apache/datafusion/pull/14870) (AmosAidoo) +- Workaround for compilation error due to rkyv#434. [#14863](https://github.com/apache/datafusion/pull/14863) (ryzhyk) +- chore(deps): bump uuid from 1.13.2 to 1.14.0 [#14866](https://github.com/apache/datafusion/pull/14866) (dependabot[bot]) +- refactor: replace OnceLock with LazyLock [#14880](https://github.com/apache/datafusion/pull/14880) (AmosAidoo) +- chore: migrate to `invoke_with_args` for datetime functions [#14876](https://github.com/apache/datafusion/pull/14876) (onlyjackfrost) +- Fix `regenerate_sqlite_files.sh` due to changes in sqllogictests [#14881](https://github.com/apache/datafusion/pull/14881) (alamb) +- Move `FileFormat` and related pieces to `datafusion-datasource` [#14873](https://github.com/apache/datafusion/pull/14873) (AdamGS) +- fix duplicated schema name error from count wildcard [#14824](https://github.com/apache/datafusion/pull/14824) (jayzhan211) +- replace TypeSignature::String with TypeSignature::Coercible for trim functions [#14865](https://github.com/apache/datafusion/pull/14865) (zjregee) +- Window Functions Order Conservation -- Follow-up On Set Monotonicity [#14813](https://github.com/apache/datafusion/pull/14813) (berkaysynnada) +- Implement builder style API for ParserOptions [#14887](https://github.com/apache/datafusion/pull/14887) (kosiew) +- chore: Attach Diagnostic to "function x does not exist" error [#14849](https://github.com/apache/datafusion/pull/14849) (onlyjackfrost) +- Fix: External sort failing on `StringView` due to shared buffers [#14823](https://github.com/apache/datafusion/pull/14823) (2010YOUY01) +- refactor: make SqlToRel::new derive the parser options from the context provider [#14822](https://github.com/apache/datafusion/pull/14822) (niebayes) +- Datafusion-cli: Redesign the datafusion-cli execution and print, make it totally streaming printing without memory overhead [#14877](https://github.com/apache/datafusion/pull/14877) (zhuqi-lucas) +- chore: Strip debuginfo symbols for release [#14843](https://github.com/apache/datafusion/pull/14843) (comphead) +- chore(deps): bump zstd from 0.13.2 to 0.13.3 [#14889](https://github.com/apache/datafusion/pull/14889) (dependabot[bot]) +- Add DataFrame fill_null [#14769](https://github.com/apache/datafusion/pull/14769) (kosiew) +- Cancellation benchmark [#14818](https://github.com/apache/datafusion/pull/14818) (carols10cents) +- Include struct name on FileScanConfig debug impl [#14883](https://github.com/apache/datafusion/pull/14883) (alamb) +- Preserve the name of grouping sets in SimplifyExpressions [#14888](https://github.com/apache/datafusion/pull/14888) (joroKr21) +- Require `Debug` for `DataSource` [#14882](https://github.com/apache/datafusion/pull/14882) (alamb) +- Update regenerate sql dep, revert runner changes. [#14901](https://github.com/apache/datafusion/pull/14901) (Omega359) +- chore(deps): bump flate2 from 1.0.35 to 1.1.0 [#14848](https://github.com/apache/datafusion/pull/14848) (dependabot[bot]) +- replace TypeSignature::String with TypeSignature::Coercible for starts_with [#14812](https://github.com/apache/datafusion/pull/14812) (zjregee) +- Dataframe with_column and with_column_renamed performance improvements [#14653](https://github.com/apache/datafusion/pull/14653) (Omega359) +- chore(deps): bump uuid from 1.14.0 to 1.15.1 [#14911](https://github.com/apache/datafusion/pull/14911) (dependabot[bot]) +- chore(deps): bump libc from 0.2.169 to 0.2.170 [#14912](https://github.com/apache/datafusion/pull/14912) (dependabot[bot]) +- Move HashJoin from `RawTable` to `HashTable` [#14904](https://github.com/apache/datafusion/pull/14904) (Dandandan) +- Rename `DataSource` and `FileSource` fields for consistency [#14898](https://github.com/apache/datafusion/pull/14898) (alamb) +- Fix the null handling for to_char function [#14908](https://github.com/apache/datafusion/pull/14908) (kosiew) +- Add tests for Demonstrate EnforceSorting can remove a needed coalesce [#14919](https://github.com/apache/datafusion/pull/14919) (wiedld) +- Fix: New Datafusion-cli streaming printing way should handle corner case for only one small batch which lines are less than max_rows [#14921](https://github.com/apache/datafusion/pull/14921) (zhuqi-lucas) +- Add docs to `update_coalesce_ctx_children`. [#14907](https://github.com/apache/datafusion/pull/14907) (wiedld) +- chore(deps): bump the arrow-parquet group with 7 updates [#14930](https://github.com/apache/datafusion/pull/14930) (dependabot[bot]) +- chore(deps): bump aws-config from 1.5.16 to 1.5.17 [#14931](https://github.com/apache/datafusion/pull/14931) (dependabot[bot]) +- Add additional protobuf tests for plans that read parquet with projections [#14924](https://github.com/apache/datafusion/pull/14924) (alamb) +- Fix link in datasource readme [#14928](https://github.com/apache/datafusion/pull/14928) (lewiszlw) +- Expose `build_row_filter` method [#14933](https://github.com/apache/datafusion/pull/14933) (xudong963) +- Do not unescape backslashes in datafusion-cli [#14844](https://github.com/apache/datafusion/pull/14844) (Lordworms) +- Set projection before configuring the source [#14685](https://github.com/apache/datafusion/pull/14685) (blaginin) +- Add H2O.ai Database-like Ops benchmark to dfbench (join support) [#14902](https://github.com/apache/datafusion/pull/14902) (zhuqi-lucas) +- Use arrow IPC Stream format for spill files [#14868](https://github.com/apache/datafusion/pull/14868) (davidhewitt) +- refactor(properties): Split properties.rs into smaller modules [#14925](https://github.com/apache/datafusion/pull/14925) (Standing-Man) +- Fix failing extended `sqlite`test on main / update `datafusion-testing` pin [#14940](https://github.com/apache/datafusion/pull/14940) (alamb) +- Revert Datafusion-cli: Redesign the datafusion-cli execution and print, make it totally streaming printing without memory overhead [#14948](https://github.com/apache/datafusion/pull/14948) (alamb) +- Remove invalid bug reproducer. [#14950](https://github.com/apache/datafusion/pull/14950) (wiedld) +- Improve documentation for `DataSourceExec`, `FileScanConfig`, `DataSource` etc [#14941](https://github.com/apache/datafusion/pull/14941) (alamb) +- Do not swap with projection when file is partitioned [#14956](https://github.com/apache/datafusion/pull/14956) (blaginin) +- Minor: Add more projection pushdown tests, clarify comments [#14963](https://github.com/apache/datafusion/pull/14963) (alamb) +- Update labeler components [#14942](https://github.com/apache/datafusion/pull/14942) (alamb) +- Deprecate `Expr::Wildcard` [#14959](https://github.com/apache/datafusion/pull/14959) (linhr) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 38 Andrew Lamb + 35 dependabot[bot] + 14 Piotr Findeisen + 10 Matthijs Brobbel + 10 logan-keede + 9 Bruce Ritchie + 9 xudong.w + 8 Jay Zhan + 8 Qi Zhu + 6 Adam Gutglick + 6 Joseph Koshakow + 5 Ian Lai + 5 Lordworms + 5 Simon Vandel Sillesen + 5 wiedld + 5 zjregee + 4 Dmitrii Blaginin + 4 Kristin Cowalcijk + 4 Peter L + 4 Yongting You + 4 irenjj + 4 oznur-synnada + 3 Andy Yen + 3 Jax Liu + 3 Oleks V + 3 Rohan Krishnaswamy + 3 Tim Saucer + 3 kosiew + 3 niebayes + 3 张林伟 + 2 @clflushopt + 2 Amos Aidoo + 2 Andrey Koshchiy + 2 Berkay Şahin + 2 Carol (Nichols || Goulding) + 2 Christian + 2 Dawei H. + 2 Elia Perantoni + 2 Georgi Krastev + 2 Jonah Gao + 2 Raz Luvaton + 2 Sergey Zhukov + 2 mertak-synnada + 1 Adrian Garcia Badaracco + 1 Alan Tang + 1 Albert Skalt + 1 Alex Huang + 1 Anlin Chen + 1 Artem Osipov + 1 Daniel Hegberg + 1 Daniël Heres + 1 David Hewitt + 1 Duong Cong Toai + 1 Eduard Karacharov + 1 Gabriel + 1 Hussein Awala + 1 Kaifeng Zheng + 1 Landon Gingerich + 1 Leonid Ryzhyk + 1 Li-Lun Lin + 1 Marko Milenković + 1 Matthew Kim + 1 Matthew Turner + 1 Namgung Chan + 1 Nelson Antunes + 1 Owen Leung + 1 Paul J. Davis + 1 Pepijn Van Eeckhoudt + 1 Sasha Syrotenko + 1 Sergei Grebnov + 1 UBarney + 1 Yingwen + 1 Zhen Wang + 1 cht42 + 1 cjw + 1 dentiny + 1 gstvg + 1 m09526 + 1 nuno-faria +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/dev/release/README.md b/dev/release/README.md index e74d84096792..6e4079de8f06 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -280,6 +280,8 @@ Verify that the Cargo.toml in the tarball contains the correct version (cd datafusion/physical-plan && cargo publish) (cd datafusion/physical-optimizer && cargo publish) (cd datafusion/catalog && cargo publish) +(cd datafusion/datasource && cargo publish) +(cd datafusion/catalog-listing && cargo publish) (cd datafusion/functions-table && cargo publish) (cd datafusion/core && cargo publish) (cd datafusion/proto-common && cargo publish) diff --git a/dev/release/verify-release-candidate.sh b/dev/release/verify-release-candidate.sh index 2c0bd216b3ac..9ecbe1bc1713 100755 --- a/dev/release/verify-release-candidate.sh +++ b/dev/release/verify-release-candidate.sh @@ -18,6 +18,37 @@ # under the License. # +# Check that required dependencies are installed +check_dependencies() { + local missing_deps=0 + local required_deps=("curl" "git" "gpg" "cc" "protoc") + + # Either shasum or sha256sum/sha512sum are required + local has_sha_tools=0 + + for dep in "${required_deps[@]}"; do + if ! command -v $dep &> /dev/null; then + echo "Error: $dep is not installed or not in PATH" + missing_deps=1 + fi + done + + # Check for either shasum or sha256sum/sha512sum + if command -v shasum &> /dev/null; then + has_sha_tools=1 + elif command -v sha256sum &> /dev/null && command -v sha512sum &> /dev/null; then + has_sha_tools=1 + else + echo "Error: Neither shasum nor sha256sum/sha512sum are installed or in PATH" + missing_deps=1 + fi + + if [ $missing_deps -ne 0 ]; then + echo "Please install missing dependencies and try again" + exit 1 + fi +} + case $# in 2) VERSION="$1" RC_NUMBER="$2" @@ -31,6 +62,9 @@ set -e set -x set -o pipefail +# Add the dependency check early in the script execution +check_dependencies + SOURCE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)" ARROW_DIR="$(dirname $(dirname ${SOURCE_DIR}))" ARROW_DIST_URL='https://dist.apache.org/repos/dist/dev/datafusion' @@ -117,8 +151,11 @@ test_source_distribution() { # build and test rust + # install the needed version of rust defined in rust-toolchain.toml + rustup toolchain install + # raises on any formatting errors - rustup component add rustfmt --toolchain stable + rustup component add rustfmt cargo fmt --all -- --check # Clone testing repositories into the expected location diff --git a/docs/build.sh b/docs/build.sh index 14464fab40ea..73516e8e9c68 100755 --- a/docs/build.sh +++ b/docs/build.sh @@ -28,4 +28,4 @@ sed -i -e 's/\.\.\/\.\.\/\.\.\//https:\/\/github.com\/apache\/arrow-datafusion\/ python rustdoc_trim.py -make SOURCEDIR=`pwd`/temp html +make SOURCEDIR=`pwd`/temp SPHINXOPTS=-W html diff --git a/docs/source/conf.py b/docs/source/conf.py index 51408f4fa76f..a5a3e66c6b3c 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -106,10 +106,6 @@ "theme_overrides.css" ] -html_js_files = [ - ("https://buttons.github.io/buttons.js", {'async': 'true', 'defer': 'true'}), -] - html_sidebars = { "**": ["docs-sidebar.html"], } diff --git a/docs/source/library-user-guide/api-health.md b/docs/source/contributor-guide/api-health.md similarity index 92% rename from docs/source/library-user-guide/api-health.md rename to docs/source/contributor-guide/api-health.md index 87d3754b21a7..d811bc357445 100644 --- a/docs/source/library-user-guide/api-health.md +++ b/docs/source/contributor-guide/api-health.md @@ -26,6 +26,14 @@ breaking API changes, but they are sometimes necessary. When possible, rather than making breaking API changes, we prefer to deprecate APIs to give users time to adjust to the changes. +## Upgrade Guides + +When making changes that require DataFusion users to make changes to their code +as part of an upgrade please consider adding documentation to the version +specific [Upgrade Guide] + +[upgrade guide]: ../library-user-guide/upgrading.md + ## Breaking Changes In general, a function is part of the public API if it appears on the [docs.rs page] diff --git a/docs/source/contributor-guide/howtos.md b/docs/source/contributor-guide/howtos.md index 556242751ff4..89a1bc7360a1 100644 --- a/docs/source/contributor-guide/howtos.md +++ b/docs/source/contributor-guide/howtos.md @@ -141,9 +141,9 @@ taplo fmt ## How to update protobuf/gen dependencies -The prost/tonic code can be generated by running `./regen.sh`, which in turn invokes the Rust binary located in [gen](./gen) +The prost/tonic code can be generated by running `./regen.sh`, which in turn invokes the Rust binary located in `./gen` -This is necessary after modifying the protobuf definitions or altering the dependencies of [gen](./gen), and requires a +This is necessary after modifying the protobuf definitions or altering the dependencies of `./gen`, and requires a valid installation of [protoc] (see [installation instructions] for details). ```bash diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 65d92e7d7926..e38898db5a92 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -32,7 +32,7 @@ community as well as get more familiar with Rust and the relevant codebases. ## Development Environment -You can find how to setup build and testing environment [here](https://datafusion.apache.org/user-guide/example-usage.html) +You can find how to setup build and testing environment [here](https://datafusion.apache.org/contributor-guide/development_environment.html) ## Finding and Creating Issues to Work On diff --git a/docs/source/contributor-guide/testing.md b/docs/source/contributor-guide/testing.md index 2a9f22d22d66..2868125c7f3d 100644 --- a/docs/source/contributor-guide/testing.md +++ b/docs/source/contributor-guide/testing.md @@ -58,6 +58,18 @@ Like similar systems such as [DuckDB](https://duckdb.org/dev/testing), DataFusio DataFusion has integrated [sqlite's test suite](https://sqlite.org/sqllogictest/doc/trunk/about.wiki) as a supplemental test suite that is run whenever a PR is merged into DataFusion. To run it manually please refer to the [README](https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/README.md#running-tests-sqlite) file for instructions. +## Snapshot testing + +[Insta](https://github.com/mitsuhiko/insta) is used for snapshot testing. Snapshots are generated +and compared on each test run. If the output changes, tests will fail. + +To review the changes, you can use Insta CLI: + +```shell +cargo install cargo-insta +cargo insta review +``` + ## Extended Tests In addition to the standard CI test suite that is run on all PRs prior to merge, diff --git a/docs/source/index.rst b/docs/source/index.rst index 839c896d0b4c..0dc947fdea57 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -22,16 +22,21 @@ Apache DataFusion ================= -.. Code from https://buttons.github.io/ .. raw:: html -

- - Star - - Fork -

- + + + DataFusion is an extensible query engine written in `Rust `_ that uses `Apache Arrow `_ as its in-memory format. @@ -132,8 +137,9 @@ To get started, see library-user-guide/extending-operators library-user-guide/profiling library-user-guide/query-optimizer - library-user-guide/api-health -.. _toc.contributor-guide: + library-user-guide/upgrading + +.. .. _toc.contributor-guide: .. toctree:: :maxdepth: 1 @@ -144,6 +150,7 @@ To get started, see contributor-guide/development_environment contributor-guide/architecture contributor-guide/testing + contributor-guide/api-health contributor-guide/howtos contributor-guide/roadmap contributor-guide/governance diff --git a/docs/source/library-user-guide/query-optimizer.md b/docs/source/library-user-guide/query-optimizer.md index af27bb75053b..03cd7b5bbbbe 100644 --- a/docs/source/library-user-guide/query-optimizer.md +++ b/docs/source/library-user-guide/query-optimizer.md @@ -401,7 +401,7 @@ interval arithmetic to take an expression such as `a > 2500 AND a <= 5000` and build an accurate selectivity estimate that can then be used to find more efficient plans. -#### `AnalysisContext` API +### `AnalysisContext` API The `AnalysisContext` serves as a shared knowledge base during expression evaluation and boundary analysis. Think of it as a dynamic repository that maintains information about: @@ -414,7 +414,7 @@ What makes `AnalysisContext` particularly powerful is its ability to propagate i through the expression tree. As each node in the expression tree is analyzed, it can both read from and write to this shared context, allowing for sophisticated boundary analysis and inference. -#### `ColumnStatistics` for Cardinality Estimation +### `ColumnStatistics` for Cardinality Estimation Column statistics form the foundation of optimization decisions. Rather than just tracking simple metrics, DataFusion's `ColumnStatistics` provides a rich set of information including: diff --git a/docs/source/library-user-guide/upgrading.md b/docs/source/library-user-guide/upgrading.md new file mode 100644 index 000000000000..a6679cbea9ad --- /dev/null +++ b/docs/source/library-user-guide/upgrading.md @@ -0,0 +1,215 @@ + + +# Upgrade Guides + +## DataFusion `46.0.0` + +### Use `invoke_with_args` instead of `invoke()` and `invoke_batch()` + +DataFusion is moving to a consistent API for invoking ScalarUDFs, +[`ScalarUDFImpl::invoke_with_args()`], and deprecating +[`ScalarUDFImpl::invoke()`], [`ScalarUDFImpl::invoke_batch()`], and [`ScalarUDFImpl::invoke_no_args()`] + +If you see errors such as the following it means the older APIs are being used: + +```text +This feature is not implemented: Function concat does not implement invoke but called +``` + +To fix this error, use [`ScalarUDFImpl::invoke_with_args()`] instead, as shown +below. See [PR 14876] for an example. + +Given existing code like this: + +```rust +# /* +impl ScalarUDFImpl for SparkConcat { +... + fn invoke_batch(&self, args: &[ColumnarValue], number_rows: usize) -> Result { + if args + .iter() + .any(|arg| matches!(arg.data_type(), DataType::List(_))) + { + ArrayConcat::new().invoke_batch(args, number_rows) + } else { + ConcatFunc::new().invoke_batch(args, number_rows) + } + } +} +# */ +``` + +To + +```rust +# /* comment out so they don't run +impl ScalarUDFImpl for SparkConcat { + ... + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args + .args + .iter() + .any(|arg| matches!(arg.data_type(), DataType::List(_))) + { + ArrayConcat::new().invoke_with_args(args) + } else { + ConcatFunc::new().invoke_with_args(args) + } + } +} + # */ +``` + +[`scalarudfimpl::invoke()`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.invoke +[`scalarudfimpl::invoke_batch()`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.invoke_batch +[`scalarudfimpl::invoke_no_args()`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.invoke_no_args +[`scalarudfimpl::invoke_with_args()`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html#method.invoke_with_args +[pr 14876]: https://github.com/apache/datafusion/pull/14876 + +### `ParquetExec`, `AvroExec`, `CsvExec`, `JsonExec` deprecated + +DataFusion 46 has a major change to how the built in DataSources are organized. +Instead of individual `ExecutionPlan`s for the different file formats they now +all use `DataSourceExec` and the format specific information is embodied in new +traits `DataSource` and `FileSource`. + +Here is more information about + +- [Design Ticket] +- Change PR [PR #14224] +- Example of an Upgrade [PR in delta-rs] + +[design ticket]: https://github.com/apache/datafusion/issues/13838 +[pr #14224]: https://github.com/apache/datafusion/pull/14224 +[pr in delta-rs]: https://github.com/delta-io/delta-rs/pull/3261 + +### Cookbook: Changes to `ParquetExecBuilder` + +Code that looks for `ParquetExec` like this will no longer work: + +```rust +# /* comment to avoid running + if let Some(parquet_exec) = plan.as_any().downcast_ref::() { + // Do something with ParquetExec here + } +# */ +``` + +Instead, with `DataSourceExec`, the same information is now on `FileScanConfig` and +`ParquetSource`. The equivalent code is + +```rust +# /* comment to avoid running +if let Some(datasource_exec) = plan.as_any().downcast_ref::() { + if let Some(scan_config) = datasource_exec.data_source().as_any().downcast_ref::() { + // FileGroups, and other information is on the FileScanConfig + // parquet + if let Some(parquet_source) = scan_config.file_source.as_any().downcast_ref::() + { + // Information on PruningPredicates and parquet options are here + } +} +# */ +``` + +### Cookbook: Changes to `ParquetExecBuilder` + +Likewise code that builds `ParquetExec` using the `ParquetExecBuilder` such as +the following must be changed: + +```rust +# /* comment to avoid running +let mut exec_plan_builder = ParquetExecBuilder::new( + FileScanConfig::new(self.log_store.object_store_url(), file_schema) + .with_projection(self.projection.cloned()) + .with_limit(self.limit) + .with_table_partition_cols(table_partition_cols), +) +.with_schema_adapter_factory(Arc::new(DeltaSchemaAdapterFactory {})) +.with_table_parquet_options(parquet_options); + +// Add filter +if let Some(predicate) = logical_filter { + if config.enable_parquet_pushdown { + exec_plan_builder = exec_plan_builder.with_predicate(predicate); + } +}; +# */ +``` + +New code should use `FileScanConfig` to build the appropriate `DataSourceExec`: + +```rust +# /* comment to avoid running +let mut file_source = ParquetSource::new(parquet_options) + .with_schema_adapter_factory(Arc::new(DeltaSchemaAdapterFactory {})); + +// Add filter +if let Some(predicate) = logical_filter { + if config.enable_parquet_pushdown { + file_source = file_source.with_predicate(Arc::clone(&file_schema), predicate); + } +}; + +let file_scan_config = FileScanConfig::new( + self.log_store.object_store_url(), + file_schema, + Arc::new(file_source), +) +.with_statistics(stats) +.with_projection(self.projection.cloned()) +.with_limit(self.limit) +.with_table_partition_cols(table_partition_cols); + +// Build the actual scan like this +parquet_scan: file_scan_config.build(), +# */ +``` + +### `datafusion-cli` no longer automatically unescapes strings + +`datafusion-cli` previously would incorrectly unescape string literals (see [ticket] for more details). + +To escape `'` in SQL literals, use `''`: + +```sql +> select 'it''s escaped'; ++----------------------+ +| Utf8("it's escaped") | ++----------------------+ +| it's escaped | ++----------------------+ +1 row(s) fetched. +``` + +To include special characters (such as newlines via `\n`) you can use an `E` literal string. For example + +```sql +> select 'foo\nbar'; ++------------------+ +| Utf8("foo\nbar") | ++------------------+ +| foo\nbar | ++------------------+ +1 row(s) fetched. +Elapsed 0.005 seconds. +``` + +[ticket]: https://github.com/apache/datafusion/issues/13286 diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c93bed963834..b6b53cfe49b3 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -68,7 +68,7 @@ Environment variables are read during `SessionConfig` initialisation so they mus | datafusion.execution.parquet.statistics_enabled | page | (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.max_statistics_size | 4096 | (writing) Sets max statistics size for any column. If NULL, uses default parquet writer setting max_statistics_size is deprecated, currently it is not being used | | datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. | -| datafusion.execution.parquet.created_by | datafusion version 45.0.0 | (writing) Sets "created by" property | +| datafusion.execution.parquet.created_by | datafusion version 46.0.0 | (writing) Sets "created by" property | | datafusion.execution.parquet.column_index_truncate_length | 64 | (writing) Sets column index truncate length | | datafusion.execution.parquet.statistics_truncate_length | NULL | (writing) Sets statictics truncate length. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.data_page_row_count_limit | 20000 | (writing) Sets best effort maximum number of rows in data page | @@ -122,9 +122,12 @@ Environment variables are read during `SessionConfig` initialisation so they mus | datafusion.explain.show_statistics | false | When set to true, the explain statement will print operator statistics for physical plans | | datafusion.explain.show_sizes | true | When set to true, the explain statement will print the partition sizes | | datafusion.explain.show_schema | false | When set to true, the explain statement will print schema information | +| datafusion.explain.format | indent | Display format of explain. Default is "indent". When set to "tree", it will print the plan in a tree-rendered format. | | datafusion.sql_parser.parse_float_as_decimal | false | When set to true, SQL parser will parse float as decimal type | | datafusion.sql_parser.enable_ident_normalization | true | When set to true, SQL parser will normalize ident (convert ident to lowercase when not quoted) | | datafusion.sql_parser.enable_options_value_normalization | false | When set to true, SQL parser will normalize options value (convert value to lowercase). Note that this option is ignored and will be removed in the future. All case-insensitive values are normalized automatically. | | datafusion.sql_parser.dialect | generic | Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB and Databricks. | | datafusion.sql_parser.support_varchar_with_length | true | If true, permit lengths for `VARCHAR` such as `VARCHAR(20)`, but ignore the length. If false, error if a `VARCHAR` with a length is specified. The Arrow type system does not have a notion of maximum string length and thus DataFusion can not enforce such limits. | -| datafusion.sql_parser.collect_spans | false | When set to true, the source locations relative to the original SQL query (i.e. [`Span`](sqlparser::tokenizer::Span)) will be collected and recorded in the logical plan nodes. | +| datafusion.sql_parser.map_varchar_to_utf8view | false | If true, `VARCHAR` is mapped to `Utf8View` during SQL planning. If false, `VARCHAR` is mapped to `Utf8` during SQL planning. Default is false. | +| datafusion.sql_parser.collect_spans | false | When set to true, the source locations relative to the original SQL query (i.e. [`Span`](https://docs.rs/sqlparser/latest/sqlparser/tokenizer/struct.Span.html)) will be collected and recorded in the logical plan nodes. | +| datafusion.sql_parser.recursion_limit | 50 | Specifies the recursion depth limit when parsing complex SQL Queries | diff --git a/docs/source/user-guide/sql/operators.md b/docs/source/user-guide/sql/operators.md index 51f1a26d0b61..b63f55239621 100644 --- a/docs/source/user-guide/sql/operators.md +++ b/docs/source/user-guide/sql/operators.md @@ -17,7 +17,7 @@ under the License. --> -# Operators +# Operators and Literals ## Numerical Operators @@ -552,3 +552,64 @@ Array Is Contained By | true | +-------------------------------------------------------------------------+ ``` + +## Literals + +Use single quotes for literal values. For example, the string `foo bar` is +referred to using `'foo bar'` + +```sql +select 'foo'; +``` + +### Escaping + +Unlike many other languages, SQL literals do not by default support C-style escape +sequences such as `\n` for newline. Instead all characters in a `'` string are treated +literally. + +To escape `'` in SQL literals, use `''`: + +```sql +> select 'it''s escaped'; ++----------------------+ +| Utf8("it's escaped") | ++----------------------+ +| it's escaped | ++----------------------+ +1 row(s) fetched. +``` + +Strings such as `foo\nbar` mean `\` followed by `n` (not newline): + +```sql +> select 'foo\nbar'; ++------------------+ +| Utf8("foo\nbar") | ++------------------+ +| foo\nbar | ++------------------+ +1 row(s) fetched. +Elapsed 0.005 seconds. +``` + +To add escaped characters such as newline or tab, instead of `\n` you use the +`E` style strings. For example, to add the text with a newline + +```text +foo +bar +``` + +You can use `E'foo\nbar'` + +```sql +> select E'foo\nbar'; ++-----------------+ +| Utf8("foo +bar") | ++-----------------+ +| foo +bar | ++-----------------+ +``` diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md index fb4043c33efc..60ecf7bd78d4 100644 --- a/docs/source/user-guide/sql/scalar_functions.md +++ b/docs/source/user-guide/sql/scalar_functions.md @@ -2524,6 +2524,7 @@ _Alias of [current_date](#current_date)._ - [array_intersect](#array_intersect) - [array_join](#array_join) - [array_length](#array_length) +- [array_max](#array_max) - [array_ndims](#array_ndims) - [array_pop_back](#array_pop_back) - [array_pop_front](#array_pop_front) @@ -2569,6 +2570,7 @@ _Alias of [current_date](#current_date)._ - [list_intersect](#list_intersect) - [list_join](#list_join) - [list_length](#list_length) +- [list_max](#list_max) - [list_ndims](#list_ndims) - [list_pop_back](#list_pop_back) - [list_pop_front](#list_pop_front) @@ -3002,6 +3004,33 @@ array_length(array, dimension) - list_length +### `array_max` + +Returns the maximum value in the array. + +```sql +array_max(array) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_max([3,1,4,2]); ++-----------------------------------------+ +| array_max(List([3,1,4,2])) | ++-----------------------------------------+ +| 4 | ++-----------------------------------------+ +``` + +#### Aliases + +- list_max + ### `array_ndims` Returns the number of dimensions of the array. @@ -3759,6 +3788,10 @@ _Alias of [array_to_string](#array_to_string)._ _Alias of [array_length](#array_length)._ +### `list_max` + +_Alias of [array_max](#array_max)._ + ### `list_ndims` _Alias of [array_ndims](#array_ndims)._ From ed790d67a127da9d437584b98dc6d8b3f538405c Mon Sep 17 00:00:00 2001 From: silezhou Date: Thu, 13 Mar 2025 12:44:28 +0000 Subject: [PATCH 03/10] style: cargo fmt --- datafusion/physical-plan/src/coalesce_batches.rs | 7 ------- datafusion/sql/src/planner.rs | 2 +- datafusion/sql/src/select.rs | 3 --- datafusion/sql/tests/cases/plan_to_sql.rs | 7 ++++--- 4 files changed, 5 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index a96b1bf45a47..5244038b9ae2 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -129,13 +129,6 @@ impl DisplayAs for CoalesceBatchesExec { }; Ok(()) } - DisplayFormatType::TreeRender => { - writeln!(f, "target_batch_size={}", self.target_batch_size)?; - if let Some(fetch) = self.fetch { - write!(f, "limit={fetch}")?; - }; - Ok(()) - } } } } diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 7aa222810843..6ccbe319f160 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -575,7 +575,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } } - SQLDataType::UnsignedBigInt(_) | SQLDataType::UnsignedInt8(_) => Ok(DataType::UInt64), + SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), SQLDataType::Float(_) => Ok(DataType::Float32), SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), SQLDataType::Double(ExactNumberInfo::None) | SQLDataType::DoublePrecision | SQLDataType::Float8 => Ok(DataType::Float64), diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 17c62a969db5..7780e4bec8d0 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -29,7 +29,6 @@ use crate::utils::{ use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{not_impl_err, plan_err, Column, Result}; -use datafusion_common::{not_impl_err, plan_err, Column, Result}; use datafusion_common::{RecursionUnnestOption, UnnestOptions}; use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ @@ -38,8 +37,6 @@ use datafusion_expr::expr_rewriter::{ use datafusion_expr::utils::{ expand_qualified_wildcard, expand_wildcard, expr_as_column_expr, expr_to_columns, find_aggregate_exprs, find_window_exprs, - expand_qualified_wildcard, expand_wildcard, expr_as_column_expr, expr_to_columns, - find_aggregate_exprs, find_window_exprs, }; use datafusion_expr::{ Aggregate, Expr, Filter, GroupingSet, LogicalPlan, LogicalPlanBuilder, diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 33ca5c8ace8f..d2230da96be5 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -21,9 +21,10 @@ use datafusion_expr::test::function_stub::{ count_udaf, max_udaf, min_udaf, sum, sum_udaf, }; use datafusion_expr::{ - cast, col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, - cast, col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, - LogicalPlanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, + cast, cast, col, col, lit, lit, table_scan, table_scan, wildcard, wildcard, + EmptyRelation, EmptyRelation, Expr, Expr, Extension, Extension, LogicalPlan, + LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, + UserDefinedLogicalNodeCore, }; use datafusion_functions::unicode; use datafusion_functions_aggregate::grouping::grouping_udaf; From 5c62884ddd79b5fdf603c09af06cb668d83e3011 Mon Sep 17 00:00:00 2001 From: silezhou Date: Thu, 13 Mar 2025 12:45:38 +0000 Subject: [PATCH 04/10] fix: import --- datafusion/sql/tests/cases/plan_to_sql.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index d2230da96be5..90a0c584675c 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -21,9 +21,8 @@ use datafusion_expr::test::function_stub::{ count_udaf, max_udaf, min_udaf, sum, sum_udaf, }; use datafusion_expr::{ - cast, cast, col, col, lit, lit, table_scan, table_scan, wildcard, wildcard, - EmptyRelation, EmptyRelation, Expr, Expr, Extension, Extension, LogicalPlan, - LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, + cast, col, lit, table_scan, wildcard, + EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; use datafusion_functions::unicode; @@ -40,7 +39,6 @@ use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser}; use sqlparser::ast::Statement; use std::hash::Hash; use std::ops::Add; -use std::ops::Add; use std::sync::Arc; use std::{fmt, vec}; From c4a2a78f8183754eeb461b6ca8db3baf362cf7a3 Mon Sep 17 00:00:00 2001 From: silezhou Date: Thu, 13 Mar 2025 12:52:39 +0000 Subject: [PATCH 05/10] style: cargo fmt --- datafusion/sql/tests/cases/plan_to_sql.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 90a0c584675c..3a84a8347512 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -21,9 +21,8 @@ use datafusion_expr::test::function_stub::{ count_udaf, max_udaf, min_udaf, sum, sum_udaf, }; use datafusion_expr::{ - cast, col, lit, table_scan, wildcard, - EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union, UserDefinedLogicalNode, - UserDefinedLogicalNodeCore, + cast, col, lit, table_scan, wildcard, EmptyRelation, Expr, Extension, LogicalPlan, + LogicalPlanBuilder, Union, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; use datafusion_functions::unicode; use datafusion_functions_aggregate::grouping::grouping_udaf; From c8c9f34da54b02e36f0d130a797a53e96e6f27da Mon Sep 17 00:00:00 2001 From: SileZhou Date: Sat, 15 Mar 2025 14:55:11 +0000 Subject: [PATCH 06/10] fix: struct.slt, update.slt and lateral join --- datafusion/sql/src/relation/join.rs | 1 + datafusion/sql/src/statement.rs | 14 ++++++++++---- datafusion/sqllogictest/test_files/struct.slt | 10 +++++----- datafusion/sqllogictest/test_files/update.slt | 4 ++-- 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 5c818fffa2e0..eeb3fba990c2 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -186,6 +186,7 @@ pub(crate) fn is_lateral_join(join: &Join) -> Result { let is_lateral_syntax = is_lateral(&join.relation); let is_apply_syntax = match join.join_operator { JoinOperator::FullOuter(..) + | JoinOperator::Right(..) | JoinOperator::RightOuter(..) | JoinOperator::RightAnti(..) | JoinOperator::RightSemi(..) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 415c16bce3fc..9038c11fe8c3 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -942,18 +942,24 @@ impl SqlToRel<'_, S> { returning, or, } => { - let from = + let froms = from.map(|update_table_from_kind| match update_table_from_kind { - UpdateTableFromKind::BeforeSet(from) => from[0].clone(), - UpdateTableFromKind::AfterSet(from) => from[0].clone(), + UpdateTableFromKind::BeforeSet(froms) => froms.clone(), + UpdateTableFromKind::AfterSet(froms) => froms.clone(), }); + // TODO: support multiple tables in UPDATE SET FROM + if froms.clone().is_some_and(|f| f.len() > 1) { + println!("---------------------------------------------"); + plan_err!("Multiple tables in UPDATE SET FROM not yet supported")?; + } + let update_from = froms.map(|f| f.first().unwrap().clone()); if returning.is_some() { plan_err!("Update-returning clause not yet supported")?; } if or.is_some() { plan_err!("ON conflict not supported")?; } - self.update_to_plan(table, assignments, from, selection) + self.update_to_plan(table, assignments, update_from, selection) } Statement::Delete(Delete { diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index b547271925aa..d1757df87edd 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/datafusion/sqllogictest/test_files/struct.slt @@ -286,7 +286,7 @@ drop table struct_values; statement ok CREATE OR REPLACE VIEW complex_view AS SELECT { - 'user': { + 'user_information': { 'info': { 'personal': { 'name': 'John Doe', @@ -347,22 +347,22 @@ SELECT { } AS complex_data; query T -SELECT complex_data.user.info.personal.name FROM complex_view; +SELECT complex_data.user_information.info.personal.name FROM complex_view; ---- John Doe query I -SELECT complex_data.user.info.personal.age FROM complex_view; +SELECT complex_data.user_information.info.personal.age FROM complex_view; ---- 30 query T -SELECT complex_data.user.info.address.city FROM complex_view; +SELECT complex_data.user_information.info.address.city FROM complex_view; ---- Anytown query T -SELECT complex_data.user.preferences.languages[2] FROM complex_view; +SELECT complex_data.user_information.preferences.languages[2] FROM complex_view; ---- es diff --git a/datafusion/sqllogictest/test_files/update.slt b/datafusion/sqllogictest/test_files/update.slt index 0f9582b04c58..908d2b34aea4 100644 --- a/datafusion/sqllogictest/test_files/update.slt +++ b/datafusion/sqllogictest/test_files/update.slt @@ -78,8 +78,8 @@ physical_plan_error This feature is not implemented: Unsupported logical plan: D statement ok create table t3(a int, b varchar, c double, d int); -# set from multiple tables, sqlparser only supports from one table -query error DataFusion error: SQL error: ParserError\("Expected end of statement, found: ,"\) +# set from multiple tables, DataFusion only supports from one table +query error DataFusion error: Error during planning: Multiple tables in UPDATE SET FROM not yet supported explain update t1 set b = t2.b, c = t3.a, d = 1 from t2, t3 where t1.a = t2.a and t1.a = t3.a; # test table alias From 6c89207fbe79be432ae7ba65e420f1cd28c652d5 Mon Sep 17 00:00:00 2001 From: SileZhou Date: Thu, 20 Mar 2025 12:44:43 +0000 Subject: [PATCH 07/10] chore: cargo fmt & remove unwrap --- datafusion/sql/src/expr/function.rs | 10 +- datafusion/sql/src/expr/mod.rs | 7 +- datafusion/sql/src/planner.rs | 199 ++++++++++++++++------------ datafusion/sql/src/query.rs | 5 +- datafusion/sql/src/relation/join.rs | 8 +- datafusion/sql/src/statement.rs | 4 +- datafusion/sql/src/unparser/plan.rs | 2 +- 7 files changed, 138 insertions(+), 97 deletions(-) diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 64489a857ff4..436f4388d8a3 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -222,7 +222,15 @@ impl SqlToRel<'_, S> { // (e.g. "foo.bar") for function names yet name.to_string() } else { - crate::utils::normalize_ident(name.0[0].as_ident().unwrap().clone()) + match name.0[0].as_ident() { + Some(ident) => crate::utils::normalize_ident(ident.clone()), + None => { + return plan_err!( + "Expected an identifier in function name, but found {:?}", + name.0[0] + ) + } + } }; if name.eq("make_map") { diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index ef489a3d9656..d29ccdc6a7e9 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -1129,9 +1129,10 @@ impl SqlToRel<'_, S> { } } AccessExpr::Dot(expr) => match expr { - SQLExpr::Value( - Value::SingleQuotedString(s) | Value::DoubleQuotedString(s), - ) => Ok(Some(GetFieldAccess::NamedStructField { + SQLExpr::Value(ValueWithSpan { + value: Value::SingleQuotedString(s) | Value::DoubleQuotedString(s), + span : _ + }) => Ok(Some(GetFieldAccess::NamedStructField { name: ScalarValue::from(s), })), _ => { diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 6ccbe319f160..2170bbe970cc 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -177,16 +177,17 @@ impl IdentNormalizer { } } -/// Struct to store the states used by the Planner. The Planner will leverage the states to resolve -/// CTEs, Views, subqueries and PREPARE statements. The states include +/// Struct to store the states used by the Planner. The Planner will leverage the states +/// to resolve CTEs, Views, subqueries and PREPARE statements. The states include /// Common Table Expression (CTE) provided with WITH clause and /// Parameter Data Types provided with PREPARE statement and the query schema of the /// outer query plan. /// /// # Cloning /// -/// Only the `ctes` are truly cloned when the `PlannerContext` is cloned. This helps resolve -/// scoping issues of CTEs. By using cloning, a subquery can inherit CTEs from the outer query +/// Only the `ctes` are truly cloned when the `PlannerContext` is cloned. +/// This helps resolve scoping issues of CTEs. +/// By using cloning, a subquery can inherit CTEs from the outer query /// and can also define its own private CTEs without affecting the outer query. /// #[derive(Debug, Clone)] @@ -329,7 +330,8 @@ impl PlannerContext { /// by subsequent passes. /// /// Key interfaces are: -/// * [`Self::sql_statement_to_plan`]: Convert a statement (e.g. `SELECT ...`) into a [`LogicalPlan`] +/// * [`Self::sql_statement_to_plan`]: Convert a statement +/// (e.g. `SELECT ...`) into a [`LogicalPlan`] /// * [`Self::sql_to_expr`]: Convert an expression (e.g. `1 + 2`) into an [`Expr`] pub struct SqlToRel<'a, S: ContextProvider> { pub(crate) context_provider: &'a S, @@ -442,7 +444,8 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Ok(plan) } else if idents.len() != plan.schema().fields().len() { plan_err!( - "Source table contains {} columns but only {} names given as column alias", + "Source table contains {} columns but only {} \ + names given as column alias", plan.schema().fields().len(), idents.len() ) @@ -556,16 +559,21 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { SQLDataType::Boolean | SQLDataType::Bool => Ok(DataType::Boolean), SQLDataType::TinyInt(_) => Ok(DataType::Int8), SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16), - SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => Ok(DataType::Int32), + SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => { + Ok(DataType::Int32) + } SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64), SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8), - SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => Ok(DataType::UInt16), - SQLDataType::IntUnsigned(_) | SQLDataType::IntegerUnsigned(_) | SQLDataType::Int4Unsigned(_) => { - Ok(DataType::UInt32) - } + SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => { + Ok(DataType::UInt16) + } + SQLDataType::IntUnsigned(_) + | SQLDataType::IntegerUnsigned(_) + | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32), SQLDataType::Varchar(length) => { match (length, self.options.support_varchar_with_length) { - (Some(_), false) => plan_err!("does not support Varchar with length, please set `support_varchar_with_length` to be true"), + (Some(_), false) => plan_err!("does not support Varchar with length, \ + please set `support_varchar_with_length` to be true"), _ => { if self.options.map_varchar_to_utf8view { Ok(DataType::Utf8View) @@ -575,83 +583,90 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { } } } - SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), + SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => { + Ok(DataType::UInt64) + } SQLDataType::Float(_) => Ok(DataType::Float32), SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), - SQLDataType::Double(ExactNumberInfo::None) | SQLDataType::DoublePrecision | SQLDataType::Float8 => Ok(DataType::Float64), - SQLDataType::Double(ExactNumberInfo::Precision(_)|ExactNumberInfo::PrecisionAndScale(_, _)) => { - not_impl_err!("Unsupported SQL type (precision/scale not supported) {sql_type}") - } - SQLDataType::Char(_) - | SQLDataType::Text - | SQLDataType::String(_) => Ok(DataType::Utf8), + SQLDataType::Double(ExactNumberInfo::None) + | SQLDataType::DoublePrecision + | SQLDataType::Float8 => Ok(DataType::Float64), + SQLDataType::Double( + ExactNumberInfo::Precision(_) | ExactNumberInfo::PrecisionAndScale(_, _), + ) => { + not_impl_err!( + "Unsupported SQL type (precision/scale not supported) {sql_type}" + ) + } + SQLDataType::Char(_) | SQLDataType::Text | SQLDataType::String(_) => { + Ok(DataType::Utf8) + } SQLDataType::Timestamp(precision, tz_info) - if precision.is_none() || [0, 3, 6, 9].contains(&precision.unwrap()) => { - let tz = if matches!(tz_info, TimezoneInfo::Tz) - || matches!(tz_info, TimezoneInfo::WithTimeZone) - { - // Timestamp With Time Zone - // INPUT : [SQLDataType] TimestampTz + [Config] Time Zone - // OUTPUT: [ArrowDataType] Timestamp - self.context_provider.options().execution.time_zone.clone() - } else { - // Timestamp Without Time zone - None - }; - let precision = match precision { - Some(0) => TimeUnit::Second, - Some(3) => TimeUnit::Millisecond, - Some(6) => TimeUnit::Microsecond, - None | Some(9) => TimeUnit::Nanosecond, - _ => unreachable!(), - }; - Ok(DataType::Timestamp(precision, tz.map(Into::into))) - } + if precision.is_none() || [0, 3, 6, 9].contains(&precision.unwrap()) => + { + let tz = if matches!(tz_info, TimezoneInfo::Tz) + || matches!(tz_info, TimezoneInfo::WithTimeZone) + { + // Timestamp With Time Zone + // INPUT : [SQLDataType] TimestampTz + [Config] Time Zone + // OUTPUT: [ArrowDataType] Timestamp + self.context_provider.options().execution.time_zone.clone() + } else { + // Timestamp Without Time zone + None + }; + let precision = match precision { + Some(0) => TimeUnit::Second, + Some(3) => TimeUnit::Millisecond, + Some(6) => TimeUnit::Microsecond, + None | Some(9) => TimeUnit::Nanosecond, + _ => unreachable!(), + }; + Ok(DataType::Timestamp(precision, tz.map(Into::into))) + } SQLDataType::Date => Ok(DataType::Date32), SQLDataType::Time(None, tz_info) => { - if matches!(tz_info, TimezoneInfo::None) - || matches!(tz_info, TimezoneInfo::WithoutTimeZone) - { - Ok(DataType::Time64(TimeUnit::Nanosecond)) - } else { - // We don't support TIMETZ and TIME WITH TIME ZONE for now - not_impl_err!( - "Unsupported SQL type {sql_type:?}" - ) - } - } + if matches!(tz_info, TimezoneInfo::None) + || matches!(tz_info, TimezoneInfo::WithoutTimeZone) + { + Ok(DataType::Time64(TimeUnit::Nanosecond)) + } else { + // We don't support TIMETZ and TIME WITH TIME ZONE for now + not_impl_err!("Unsupported SQL type {sql_type:?}") + } + } SQLDataType::Numeric(exact_number_info) - | SQLDataType::Decimal(exact_number_info) => { - let (precision, scale) = match *exact_number_info { - ExactNumberInfo::None => (None, None), - ExactNumberInfo::Precision(precision) => (Some(precision), None), - ExactNumberInfo::PrecisionAndScale(precision, scale) => { - (Some(precision), Some(scale)) - } - }; - make_decimal_type(precision, scale) + | SQLDataType::Decimal(exact_number_info) => { + let (precision, scale) = match *exact_number_info { + ExactNumberInfo::None => (None, None), + ExactNumberInfo::Precision(precision) => (Some(precision), None), + ExactNumberInfo::PrecisionAndScale(precision, scale) => { + (Some(precision), Some(scale)) } + }; + make_decimal_type(precision, scale) + } SQLDataType::Bytea => Ok(DataType::Binary), SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), SQLDataType::Struct(fields, _) => { - let fields = fields - .iter() - .enumerate() - .map(|(idx, field)| { - let data_type = self.convert_data_type(&field.field_type)?; - let field_name = match &field.field_name{ - Some(ident) => ident.clone(), - None => Ident::new(format!("c{idx}")) - }; - Ok(Arc::new(Field::new( - self.ident_normalizer.normalize(field_name), - data_type, - true, - ))) - }) - .collect::>>()?; - Ok(DataType::Struct(Fields::from(fields))) - } + let fields = fields + .iter() + .enumerate() + .map(|(idx, field)| { + let data_type = self.convert_data_type(&field.field_type)?; + let field_name = match &field.field_name { + Some(ident) => ident.clone(), + None => Ident::new(format!("c{idx}")), + }; + Ok(Arc::new(Field::new( + self.ident_normalizer.normalize(field_name), + data_type, + true, + ))) + }) + .collect::>>()?; + Ok(DataType::Struct(Fields::from(fields))) + } SQLDataType::Nvarchar(_) | SQLDataType::JSON | SQLDataType::Uuid @@ -843,7 +858,7 @@ pub(crate) fn idents_to_table_reference( pub fn object_name_to_qualifier( sql_table_name: &ObjectName, enable_normalization: bool, -) -> String { +) -> Result { let columns = vec!["table_name", "table_schema", "table_catalog"].into_iter(); let normalizer = IdentNormalizer::new(enable_normalization); sql_table_name @@ -852,12 +867,22 @@ pub fn object_name_to_qualifier( .rev() .zip(columns) .map(|(object_name_part, column_name)| { - format!( - r#"{} = '{}'"#, - column_name, - normalizer.normalize(object_name_part.as_ident().unwrap().clone()) - ) + object_name_part + .as_ident() + .map(|ident| { + format!( + r#"{} = '{}'"#, + column_name, + normalizer.normalize(ident.clone()) + ) + }) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "Expected identifier, but found: {:?}", + object_name_part + )) + }) }) - .collect::>() - .join(" AND ") + .collect::>>() + .map(|parts| parts.join(" AND ")) } diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index cbbf83ff056e..7a2fd78d7cfd 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -22,8 +22,9 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use crate::stack::StackGuard; use datafusion_common::{not_impl_err, Constraints, DFSchema, Result}; use datafusion_expr::expr::Sort; +use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ - CreateMemoryTable, DdlStatement, Distinct, Expr, LogicalPlan, LogicalPlanBuilder, + CreateMemoryTable, DdlStatement, Distinct, LogicalPlan, LogicalPlanBuilder, }; use sqlparser::ast::{ Expr as SQLExpr, Offset as SQLOffset, OrderBy, OrderByExpr, OrderByKind, Query, @@ -157,7 +158,7 @@ fn to_order_by_exprs(order_by: Option) -> Result> { /// Returns the order by expressions from the query with the select expressions. pub(crate) fn to_order_by_exprs_with_select( order_by: Option, - _select_exprs: Option>, // TODO: ORDER BY ALL + _select_exprs: Option>, // TODO: ORDER BY ALL ) -> Result> { let Some(OrderBy { kind, interpolate }) = order_by else { // If no order by, return an empty array. diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index eeb3fba990c2..bdc30863a1cd 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -136,7 +136,13 @@ impl SqlToRel<'_, S> { ) } else { let id = object_names.swap_remove(0); - Ok(self.ident_normalizer.normalize(id.as_ident().unwrap().clone())) + id.as_ident() + .ok_or_else(|| { + datafusion_common::DataFusionError::Plan( + "Expected identifier in USING clause".to_string(), + ) + }) + .map(|ident| self.ident_normalizer.normalize(ident.clone())) } }) .collect::>>()?; diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 7a0f6f256e1f..0e461165667e 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -2087,7 +2087,7 @@ impl SqlToRel<'_, S> { let where_clause = object_name_to_qualifier( &sql_table_name, self.options.enable_ident_normalization, - ); + )?; if !self.has_table("information_schema", "columns") { return plan_err!( @@ -2212,7 +2212,7 @@ ON p.function_name = r.routine_name let where_clause = object_name_to_qualifier( &sql_table_name, self.options.enable_ident_normalization, - ); + )?; // Do a table lookup to verify the table exists let table_ref = self.object_name_to_table_reference(sql_table_name)?; diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 95609bb920f0..eb99d1e27031 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -668,7 +668,7 @@ impl Unparser<'_> { )); } exists_select.projection(vec![ast::SelectItem::UnnamedExpr( - ast::Expr::Value(ast::Value::Number("1".to_string(), false)), + ast::Expr::value(ast::Value::Number("1".to_string(), false)), )]); query_builder.body(Box::new(SetExpr::Select(Box::new( exists_select.build()?, From 39ac8095adbdc36075916a031962a11a598e7dfc Mon Sep 17 00:00:00 2001 From: SileZhou Date: Sat, 22 Mar 2025 06:34:49 +0000 Subject: [PATCH 08/10] fix: remove meaningless comment for rustfmt --- datafusion/sql/src/planner.rs | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 2170bbe970cc..374e5c17e811 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -572,8 +572,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32), SQLDataType::Varchar(length) => { match (length, self.options.support_varchar_with_length) { - (Some(_), false) => plan_err!("does not support Varchar with length, \ - please set `support_varchar_with_length` to be true"), + (Some(_), false) => plan_err!( + "does not support Varchar with length, \ + please set `support_varchar_with_length` to be true" + ), _ => { if self.options.map_varchar_to_utf8view { Ok(DataType::Utf8View) @@ -686,9 +688,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::CharVarying(_) | SQLDataType::CharacterLargeObject(_) | SQLDataType::CharLargeObject(_) - // Unsupported precision | SQLDataType::Timestamp(_, _) - // Precision is not supported | SQLDataType::Time(Some(_), _) | SQLDataType::Dec(_) | SQLDataType::BigNumeric(_) @@ -699,7 +699,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::Float64 | SQLDataType::JSONB | SQLDataType::Unspecified - // Clickhouse datatypes | SQLDataType::Int16 | SQLDataType::Int32 | SQLDataType::Int128 @@ -721,7 +720,6 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::Nullable(_) | SQLDataType::LowCardinality(_) | SQLDataType::Trigger - // MySQL datatypes | SQLDataType::TinyBlob | SQLDataType::MediumBlob | SQLDataType::LongBlob @@ -734,15 +732,12 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { | SQLDataType::SignedInteger | SQLDataType::Unsigned | SQLDataType::UnsignedInteger - // BigQuery UDFs | SQLDataType::AnyType - // Postgres datatypes | SQLDataType::Table(_) | SQLDataType::VarBit(_) - | SQLDataType::GeometricType(_) - => not_impl_err!( - "Unsupported SQL type {sql_type:?}" - ), + | SQLDataType::GeometricType(_) => { + not_impl_err!("Unsupported SQL type {sql_type:?}") + } } } From 4dca5516a59e039f89640f1ca46d698c281d39b9 Mon Sep 17 00:00:00 2001 From: SileZhou Date: Sat, 22 Mar 2025 07:06:43 +0000 Subject: [PATCH 09/10] refactor: plan_datafusion_err, remove clone --- datafusion/datasource/src/statistics.rs | 6 +++--- datafusion/sql/src/planner.rs | 15 +++++++++++---- datafusion/sql/src/query.rs | 2 +- datafusion/sql/src/relation/join.rs | 6 +++--- datafusion/sql/src/select.rs | 2 +- datafusion/sql/src/statement.rs | 10 ++++++---- 6 files changed, 25 insertions(+), 16 deletions(-) diff --git a/datafusion/datasource/src/statistics.rs b/datafusion/datasource/src/statistics.rs index 9df5aa993d43..cd002a96683a 100644 --- a/datafusion/datasource/src/statistics.rs +++ b/datafusion/datasource/src/statistics.rs @@ -30,7 +30,7 @@ use arrow::{ compute::SortColumn, row::{Row, Rows}, }; -use datafusion_common::{plan_err, DataFusionError, Result}; +use datafusion_common::{plan_datafusion_err, plan_err, DataFusionError, Result}; use datafusion_physical_expr::{expressions::Column, PhysicalSortExpr}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -202,10 +202,10 @@ impl MinMaxStatistics { .zip(max_values.column_by_name(column.name())) } .ok_or_else(|| { - DataFusionError::Plan(format!( + plan_datafusion_err!( "missing column in MinMaxStatistics::new: '{}'", column.name() - )) + ) }) }) .collect::>>()? diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 374e5c17e811..180017ee9c19 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -770,8 +770,15 @@ pub fn object_name_to_table_reference( let ObjectName(object_name_parts) = object_name; let idents = object_name_parts .into_iter() - .map(|object_name_part| object_name_part.as_ident().unwrap().clone()) - .collect(); + .map(|object_name_part| { + object_name_part.as_ident().cloned().ok_or_else(|| { + plan_datafusion_err!( + "Expected identifier, but found: {:?}", + object_name_part + ) + }) + }) + .collect::>>()?; idents_to_table_reference(idents, enable_normalization) } @@ -872,10 +879,10 @@ pub fn object_name_to_qualifier( ) }) .ok_or_else(|| { - DataFusionError::Plan(format!( + plan_datafusion_err!( "Expected identifier, but found: {:?}", object_name_part - )) + ) }) }) .collect::>>() diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 7a2fd78d7cfd..ea641320c01b 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -158,7 +158,7 @@ fn to_order_by_exprs(order_by: Option) -> Result> { /// Returns the order by expressions from the query with the select expressions. pub(crate) fn to_order_by_exprs_with_select( order_by: Option, - _select_exprs: Option>, // TODO: ORDER BY ALL + _select_exprs: Option<&Vec>, // TODO: ORDER BY ALL ) -> Result> { let Some(OrderBy { kind, interpolate }) = order_by else { // If no order by, return an empty array. diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index bdc30863a1cd..8a3c20e3971b 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -16,7 +16,7 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{not_impl_err, Column, Result}; +use datafusion_common::{not_impl_err, plan_datafusion_err, Column, Result}; use datafusion_expr::{JoinType, LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{ Join, JoinConstraint, JoinOperator, ObjectName, TableFactor, TableWithJoins, @@ -138,8 +138,8 @@ impl SqlToRel<'_, S> { let id = object_names.swap_remove(0); id.as_ident() .ok_or_else(|| { - datafusion_common::DataFusionError::Plan( - "Expected identifier in USING clause".to_string(), + plan_datafusion_err!( + "Expected identifier in USING clause" ) }) .map(|ident| self.ident_normalizer.normalize(ident.clone())) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index e252a3d5a048..2a2d0b3b3eb8 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -94,7 +94,7 @@ impl SqlToRel<'_, S> { )?; let order_by = - to_order_by_exprs_with_select(query_order_by, Some(select_exprs.clone()))?; + to_order_by_exprs_with_select(query_order_by, Some(&select_exprs))?; // Having and group by clause may reference aliases defined in select projection let projected_plan = self.project(base_plan.clone(), select_exprs)?; diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 0e461165667e..436fda89a25e 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -79,6 +79,8 @@ fn object_name_to_string(object_name: &ObjectName) -> String { .map(|object_name_part| { object_name_part .as_ident() + // TODO: It might be better to return an error + // than to silently use a default value. .map_or_else(String::new, ident_to_string) }) .collect::>() @@ -948,15 +950,15 @@ impl SqlToRel<'_, S> { } => { let froms = from.map(|update_table_from_kind| match update_table_from_kind { - UpdateTableFromKind::BeforeSet(froms) => froms.clone(), - UpdateTableFromKind::AfterSet(froms) => froms.clone(), + UpdateTableFromKind::BeforeSet(froms) => froms, + UpdateTableFromKind::AfterSet(froms) => froms, }); // TODO: support multiple tables in UPDATE SET FROM - if froms.clone().is_some_and(|f| f.len() > 1) { + if froms.as_ref().is_some_and(|f| f.len() > 1) { println!("---------------------------------------------"); plan_err!("Multiple tables in UPDATE SET FROM not yet supported")?; } - let update_from = froms.map(|f| f.first().unwrap().clone()); + let update_from = froms.and_then(|mut f| f.pop()); if returning.is_some() { plan_err!("Update-returning clause not yet supported")?; } From bc3768fab374c5ba2cc5eebbffed64d57512dd85 Mon Sep 17 00:00:00 2001 From: jonahgao Date: Mon, 24 Mar 2025 09:49:27 +0800 Subject: [PATCH 10/10] remove println --- datafusion/sql/src/statement.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 436fda89a25e..fc6cb0d32fef 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -955,7 +955,6 @@ impl SqlToRel<'_, S> { }); // TODO: support multiple tables in UPDATE SET FROM if froms.as_ref().is_some_and(|f| f.len() > 1) { - println!("---------------------------------------------"); plan_err!("Multiple tables in UPDATE SET FROM not yet supported")?; } let update_from = froms.and_then(|mut f| f.pop());