Optimize where exists sub-queries into aggregate and join

datafusion/common/src/scalar.rs

@@ -37,6 +37,8 @@ use std::{convert::TryFrom, fmt, iter::repeat, sync::Arc};
 
 /// Represents a dynamically typed, nullable single value.
 /// This is the single-valued counter-part of arrow’s `Array`.
+/// https://arrow.apache.org/docs/python/api/datatypes.html
+/// https://github.com/apache/arrow/blob/master/format/Schema.fbs#L354-L375
 #[derive(Clone)]
 pub enum ScalarValue {
  /// represents `DataType::Null` (castable to/from any other type)
@@ -75,9 +77,9 @@ pub enum ScalarValue {
  LargeBinary(Option<Vec<u8>>),
  /// list of nested ScalarValue
  List(Option<Vec<ScalarValue>>, Box<DataType>),
- /// Date stored as a signed 32bit int
+ /// Date stored as a signed 32bit int days since UNIX epoch 1970-01-01
  Date32(Option<i32>),
- /// Date stored as a signed 64bit int
+ /// Date stored as a signed 64bit int milliseconds since UNIX epoch 1970-01-01
  Date64(Option<i64>),
  /// Timestamp Second
  TimestampSecond(Option<i64>, Option<String>),
@@ -87,11 +89,14 @@ pub enum ScalarValue {
  TimestampMicrosecond(Option<i64>, Option<String>),
  /// Timestamp Nanoseconds
  TimestampNanosecond(Option<i64>, Option<String>),
- /// Interval with YearMonth unit
+ /// Number of elapsed whole months
  IntervalYearMonth(Option<i32>),
- /// Interval with DayTime unit
+ /// Number of elapsed days and milliseconds (no leap seconds)
+ /// stored as 2 contiguous 32-bit signed integers
  IntervalDayTime(Option<i64>),
- /// Interval with MonthDayNano unit
+ /// A triple of the number of elapsed months, days, and nanoseconds.
+ /// Months and days are encoded as 32-bit signed integers.
+ /// Nanoseconds is encoded as a 64-bit signed integer (no leap seconds).
  IntervalMonthDayNano(Option<i128>),
  /// struct of nested ScalarValue
  Struct(Option<Vec<ScalarValue>>, Box<Vec<Field>>),

datafusion/core/src/execution/context.rs

@@ -99,6 +99,7 @@ use chrono::{DateTime, Utc};
 use datafusion_common::ScalarValue;
 use datafusion_expr::TableSource;
 use datafusion_optimizer::filter_null_join_keys::FilterNullJoinKeys;
+use datafusion_optimizer::subquery_decorrelate::SubqueryDecorrelate;
 use datafusion_sql::{
  parser::DFParser,
  planner::{ContextProvider, SqlToRel},
@@ -1239,6 +1240,7 @@ impl SessionState {
  // of applying other optimizations
  Arc::new(SimplifyExpressions::new()),
  Arc::new(SubqueryFilterToJoin::new()),
+ Arc::new(SubqueryDecorrelate::new()),
  Arc::new(EliminateFilter::new()),
  Arc::new(CommonSubexprEliminate::new()),
  Arc::new(EliminateLimit::new()),

datafusion/core/tests/sql/mod.rs

@@ -108,6 +108,7 @@ mod explain;
 mod idenfifers;
 pub mod information_schema;
 mod partitioned_csv;
+mod subqueries;
 #[cfg(feature = "unicode_expressions")]
 pub mod unicode;
 
@@ -483,7 +484,37 @@ fn get_tpch_table_schema(table: &str) -> Schema {
  Field::new("n_comment", DataType::Utf8, false),
  ]),
 
- _ => unimplemented!(),
+ "supplier" => Schema::new(vec![
+ Field::new("s_suppkey", DataType::Int64, false),
+ Field::new("s_name", DataType::Utf8, false),
+ Field::new("s_address", DataType::Utf8, false),
+ Field::new("s_nationkey", DataType::Int64, false),
+ Field::new("s_phone", DataType::Utf8, false),
+ Field::new("s_acctbal", DataType::Float64, false),
+ Field::new("s_comment", DataType::Utf8, false),
+ ]),
+
+ "partsupp" => Schema::new(vec![
+ Field::new("ps_partkey", DataType::Int64, false),
+ Field::new("ps_suppkey", DataType::Int64, false),
+ Field::new("ps_availqty", DataType::Int32, false),
+ Field::new("ps_supplycost", DataType::Float64, false),
+ Field::new("ps_comment", DataType::Utf8, false),
+ ]),
+
+ "part" => Schema::new(vec![
+ Field::new("p_partkey", DataType::Int64, false),
+ Field::new("p_name", DataType::Utf8, false),
+ Field::new("p_mfgr", DataType::Utf8, false),
+ Field::new("p_brand", DataType::Utf8, false),
+ Field::new("p_type", DataType::Utf8, false),
+ Field::new("p_size", DataType::Int32, false),
+ Field::new("p_container", DataType::Utf8, false),
+ Field::new("p_retailprice", DataType::Float64, false),
+ Field::new("p_comment", DataType::Utf8, false),
+ ]),
+
+ _ => unimplemented!("Table: {}", table),
  }
 }
 

datafusion/core/tests/sql/subqueries.rs

@@ -0,0 +1,67 @@
+use super::*;
+use crate::sql::execute_to_batches;
+use datafusion::assert_batches_eq;
+use datafusion::prelude::SessionContext;
+
+#[tokio::test]
+async fn tpch_q4_correlated() -> Result<()> {
+ let ctx = SessionContext::new();
+ register_tpch_csv(&ctx, "orders").await?;
+ register_tpch_csv(&ctx, "lineitem").await?;
+
+ /*
+ #orders.o_orderpriority ASC NULLS LAST
+ Projection: #orders.o_orderpriority, #COUNT(UInt8(1)) AS order_count
+ Aggregate: groupBy=[[#orders.o_orderpriority]], aggr=[[COUNT(UInt8(1))]]
+ Filter: EXISTS ( -- plan
+ Subquery: Projection: * -- proj
+ Filter: #lineitem.l_orderkey = #orders.o_orderkey -- filter
+ TableScan: lineitem projection=None -- filter.input
+ )
+ TableScan: orders projection=None -- plan.inputs
+ */
+ let sql = r#"
+ select o_orderpriority, count(*) as order_count
+ from orders
+ where exists (
+ select * from lineitem where l_orderkey = o_orderkey and l_commitdate < l_receiptdate)
+ group by o_orderpriority
+ order by o_orderpriority;
+ "#;
+
+ // assert plan
+ let plan = ctx
+ .create_logical_plan(sql)
+ .map_err(|e| format!("{:?} at {}", e, "error"))
+ .unwrap();
+ let plan = ctx
+ .optimize(&plan)
+ .map_err(|e| format!("{:?} at {}", e, "error"))
+ .unwrap();
+ let actual = format!("{}", plan.display_indent());
+ let expected = r#"Sort: #orders.o_orderpriority ASC NULLS LAST
+ Projection: #orders.o_orderpriority, #COUNT(UInt8(1)) AS order_count
+ Aggregate: groupBy=[[#orders.o_orderpriority]], aggr=[[COUNT(UInt8(1))]]
+ Inner Join: #orders.o_orderkey = #lineitem.l_orderkey
+ TableScan: orders projection=[o_orderkey, o_orderpriority]
+ Projection: #lineitem.l_orderkey
+ Aggregate: groupBy=[[#lineitem.l_orderkey]], aggr=[[]]
+ Filter: #lineitem.l_commitdate < #lineitem.l_receiptdate
+ TableScan: lineitem projection=[l_orderkey, l_commitdate, l_receiptdate], partial_filters=[#lineitem.l_commitdate < #lineitem.l_receiptdate]"#
+ .to_string();
+ assert_eq!(actual, expected);
+
+ // assert data
+ let results = execute_to_batches(&ctx, sql).await;
+ let expected = vec![
+ "+-----------------+-------------+",
+ "| o_orderpriority | order_count |",
+ "+-----------------+-------------+",
+ "| 1-URGENT | 1 |",
+ "| 5-LOW | 1 |",
+ "+-----------------+-------------+",
+ ];
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}

datafusion/core/tests/sql/timestamp.rs

@@ -814,3 +814,123 @@ async fn group_by_timestamp_millis() -> Result<()> {
  assert_batches_eq!(expected, &actual);
  Ok(())
 }
+
+#[tokio::test]
+async fn interval_year() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-01-01' + interval '1' year as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1995-01-01 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn add_interval_month() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-01-31' + interval '1' month as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1994-02-28 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn sub_interval_month() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-03-31' - interval '1' month as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1994-02-28 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn sub_month_wrap() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-01-15' - interval '1' month as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1993-12-15 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn add_interval_day() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-01-15' + interval '1' day as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1994-01-16 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn sub_interval_day() -> Result<()> {
+ let ctx = SessionContext::new();
+
+ let sql = "select date '1994-01-01' - interval '1' day as date;";
+ let results = execute_to_batches(&ctx, sql).await;
+
+ let expected = vec![
+ "+------------+",
+ "| date |",
+ "+------------+",
+ "| 1993-12-31 |",
+ "+------------+",
+ ];
+
+ assert_batches_eq!(expected, &results);
+
+ Ok(())
+}

datafusion/core/tests/tpch-csv/part.csv

@@ -0,0 +1,2 @@
+p_partkey,p_name,p_mfgr,p_brand,p_type,p_size,p_container,p_retailprice,p_comment
+1,goldenrod lavender spring chocolate lace,Manufacturer#1,Brand#13,PROMO BURNISHED COPPER,7,JUMBO PKG,901.00,ly. slyly ironi

datafusion/core/tests/tpch-csv/partsupp.csv

@@ -0,0 +1,2 @@
+ps_partkey,ps_suppkey,ps_availqty,ps_supplycost,ps_comment
+67310,7311,100,993.49,ven ideas. quickly even packages print. pending multipliers must have to are fluff

datafusion/core/tests/tpch-csv/supplier.csv

@@ -0,0 +1,2 @@
+s_suppkey,s_name,s_address,s_nationkey,s_phone,s_acctbal,s_comment
+1,Supplier#000000001, N kD4on9OM Ipw3,gf0JBoQDd7tgrzrddZ,17,27-918-335-1736,5755.94,each slyly above the careful

datafusion/optimizer/Cargo.toml

@@ -44,4 +44,5 @@ datafusion-common = { path = "../common", version = "9.0.0" }
 datafusion-expr = { path = "../expr", version = "9.0.0" }
 datafusion-physical-expr = { path = "../physical-expr", version = "9.0.0" }
 hashbrown = { version = "0.12", features = ["raw"] }
+itertools = "0.10"
 log = "^0.4"

datafusion/optimizer/src/lib.rs

@@ -27,6 +27,7 @@ pub mod projection_push_down;
 pub mod reduce_outer_join;
 pub mod simplify_expressions;
 pub mod single_distinct_to_groupby;
+pub mod subquery_decorrelate;
 pub mod subquery_filter_to_join;
 pub mod utils;
 

datafusion/optimizer/src/simplify_expressions.rs

@@ -1951,7 +1951,7 @@ mod tests {
  let date_plus_interval_expr = to_timestamp_expr(ts_string)
  .cast_to(&DataType::Date32, schema)
  .unwrap()
- + Expr::Literal(ScalarValue::IntervalDayTime(Some(123)));
+ + Expr::Literal(ScalarValue::IntervalDayTime(Some(123i64 << 32)));
 
  let plan = LogicalPlanBuilder::from(table_scan.clone())
  .project(vec![date_plus_interval_expr])
@@ -1963,10 +1963,10 @@ mod tests {
 
  // Note that constant folder runs and folds the entire
  // expression down to a single constant (true)
- let expected = "Projection: Date32(\"18636\") AS CAST(totimestamp(Utf8(\"2020-09-08T12:05:00+00:00\")) AS Date32) + IntervalDayTime(\"123\")\
- \n TableScan: test";
+ let expected = r#"Projection: Date32("18636") AS CAST(totimestamp(Utf8("2020-09-08T12:05:00+00:00")) AS Date32) + IntervalDayTime("528280977408")
+ TableScan: test"#;
  let actual = get_optimized_plan_formatted(&plan, &time);
 
- assert_eq!(expected, actual);
+ assert_eq!(actual, expected);
  }
 }