This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 00bbb42c96 fix: correctly handle Substrait windows with rows bounds
(and validate executability of test plans) (#11278)
00bbb42c96 is described below
commit 00bbb42c96e041fe679631dae9f36bf51d4d998b
Author: Arttu <[email protected]>
AuthorDate: Sat Jul 6 13:27:35 2024 +0200
fix: correctly handle Substrait windows with rows bounds (and validate
executability of test plans) (#11278)
* add a test for window with rows specified that should fail
* make roundtrip tests actually execute the plan
previously we could create plans that seem valid but cannot actually be
executed by DF
E.g. the newly added test fails now with:
Error: Internal("Rows should be Uint")
* fix consuming bounds
---
datafusion/substrait/src/logical_plan/consumer.rs | 22 ++++++---
datafusion/substrait/src/logical_plan/producer.rs | 13 ++++++
.../tests/cases/roundtrip_logical_plan.rs | 52 ++++++++++------------
datafusion/substrait/tests/testdata/data.csv | 4 +-
4 files changed, 55 insertions(+), 36 deletions(-)
diff --git a/datafusion/substrait/src/logical_plan/consumer.rs
b/datafusion/substrait/src/logical_plan/consumer.rs
index c65943643e..03692399e1 100644
--- a/datafusion/substrait/src/logical_plan/consumer.rs
+++ b/datafusion/substrait/src/logical_plan/consumer.rs
@@ -1544,12 +1544,22 @@ fn from_substrait_bound(
BoundKind::CurrentRow(SubstraitBound::CurrentRow {}) => {
Ok(WindowFrameBound::CurrentRow)
}
- BoundKind::Preceding(SubstraitBound::Preceding { offset }) =>
Ok(
-
WindowFrameBound::Preceding(ScalarValue::Int64(Some(*offset))),
- ),
- BoundKind::Following(SubstraitBound::Following { offset }) =>
Ok(
-
WindowFrameBound::Following(ScalarValue::Int64(Some(*offset))),
- ),
+ BoundKind::Preceding(SubstraitBound::Preceding { offset }) => {
+ if *offset <= 0 {
+ return plan_err!("Preceding bound must be positive");
+ }
+ Ok(WindowFrameBound::Preceding(ScalarValue::UInt64(Some(
+ *offset as u64,
+ ))))
+ }
+ BoundKind::Following(SubstraitBound::Following { offset }) => {
+ if *offset <= 0 {
+ return plan_err!("Following bound must be positive");
+ }
+ Ok(WindowFrameBound::Following(ScalarValue::UInt64(Some(
+ *offset as u64,
+ ))))
+ }
BoundKind::Unbounded(SubstraitBound::Unbounded {}) => {
if is_lower {
Ok(WindowFrameBound::Preceding(ScalarValue::Null))
diff --git a/datafusion/substrait/src/logical_plan/producer.rs
b/datafusion/substrait/src/logical_plan/producer.rs
index 899fec21f8..9595420801 100644
--- a/datafusion/substrait/src/logical_plan/producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer.rs
@@ -2230,6 +2230,7 @@ mod test {
use crate::logical_plan::consumer::{
from_substrait_literal_without_names,
from_substrait_type_without_names,
};
+ use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
use datafusion::arrow::array::GenericListArray;
use datafusion::arrow::datatypes::Field;
use datafusion::common::scalar::ScalarStructBuilder;
@@ -2309,6 +2310,14 @@ mod test {
)?;
round_trip_literal(ScalarStructBuilder::new_null(vec![c0, c1, c2]))?;
+ round_trip_literal(ScalarValue::IntervalYearMonth(Some(17)))?;
+ round_trip_literal(ScalarValue::IntervalMonthDayNano(Some(
+ IntervalMonthDayNano::new(17, 25, 1234567890),
+ )))?;
+
round_trip_literal(ScalarValue::IntervalDayTime(Some(IntervalDayTime::new(
+ 57, 123456,
+ ))))?;
+
Ok(())
}
@@ -2376,6 +2385,10 @@ mod test {
.into(),
))?;
+ round_trip_type(DataType::Interval(IntervalUnit::YearMonth))?;
+ round_trip_type(DataType::Interval(IntervalUnit::MonthDayNano))?;
+ round_trip_type(DataType::Interval(IntervalUnit::DayTime))?;
+
Ok(())
}
diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
index dbc2e404bf..52cfa50683 100644
--- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
+++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
@@ -523,24 +523,6 @@ async fn roundtrip_arithmetic_ops() -> Result<()> {
Ok(())
}
-#[tokio::test]
-async fn roundtrip_interval_literal() -> Result<()> {
- roundtrip(
- "SELECT g from data where g = arrow_cast(INTERVAL '1 YEAR',
'Interval(YearMonth)')",
- )
- .await?;
- roundtrip(
- "SELECT g from data where g = arrow_cast(INTERVAL '1 YEAR',
'Interval(DayTime)')",
- )
- .await?;
- roundtrip(
- "SELECT g from data where g = arrow_cast(INTERVAL '1 YEAR',
'Interval(MonthDayNano)')",
- )
- .await?;
-
- Ok(())
-}
-
#[tokio::test]
async fn roundtrip_like() -> Result<()> {
roundtrip("SELECT f FROM data WHERE f LIKE 'a%b'").await
@@ -650,6 +632,17 @@ async fn simple_window_function() -> Result<()> {
roundtrip("SELECT RANK() OVER (PARTITION BY a ORDER BY b), d, sum(b) OVER
(PARTITION BY a) FROM data;").await
}
+#[tokio::test]
+async fn window_with_rows() -> Result<()> {
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 2 PRECEDING AND
CURRENT ROW) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN CURRENT ROW AND
2 FOLLOWING) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 2 PRECEDING AND
2 FOLLOWING) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN UNBOUNDED
PRECEDING AND 2 FOLLOWING) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 2 PRECEDING AND
UNBOUNDED FOLLOWING) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 2 FOLLOWING AND
4 FOLLOWING) FROM data;").await?;
+ roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 4 PRECEDING AND
2 PRECEDING) FROM data;").await
+}
+
#[tokio::test]
async fn qualified_schema_table_reference() -> Result<()> {
roundtrip("SELECT * FROM public.data;").await
@@ -810,23 +803,20 @@ async fn roundtrip_aggregate_udf() -> Result<()> {
struct Dummy {}
impl Accumulator for Dummy {
- fn state(&mut self) -> datafusion::error::Result<Vec<ScalarValue>> {
- Ok(vec![])
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ Ok(vec![ScalarValue::Float64(None), ScalarValue::UInt32(None)])
}
- fn update_batch(
- &mut self,
- _values: &[ArrayRef],
- ) -> datafusion::error::Result<()> {
+ fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {
Ok(())
}
- fn merge_batch(&mut self, _states: &[ArrayRef]) ->
datafusion::error::Result<()> {
+ fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> {
Ok(())
}
- fn evaluate(&mut self) -> datafusion::error::Result<ScalarValue> {
- Ok(ScalarValue::Float64(None))
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ Ok(ScalarValue::Int64(None))
}
fn size(&self) -> usize {
@@ -1060,6 +1050,8 @@ async fn roundtrip_with_ctx(sql: &str, ctx:
SessionContext) -> Result<()> {
assert_eq!(plan1str, plan2str);
assert_eq!(plan.schema(), plan2.schema());
+
+ DataFrame::new(ctx.state(), plan2).show().await?;
Ok(())
}
@@ -1132,7 +1124,6 @@ async fn create_context() -> Result<SessionContext> {
Field::new("d", DataType::Boolean, true),
Field::new("e", DataType::UInt32, true),
Field::new("f", DataType::Utf8, true),
- Field::new("g", DataType::Interval(IntervalUnit::DayTime), true),
];
let schema = Schema::new(fields);
explicit_options.schema = Some(&schema);
@@ -1195,6 +1186,11 @@ async fn create_all_type_context() ->
Result<SessionContext> {
),
Field::new("decimal_128_col", DataType::Decimal128(10, 2), true),
Field::new("decimal_256_col", DataType::Decimal256(10, 2), true),
+ Field::new(
+ "interval_day_time_col",
+ DataType::Interval(IntervalUnit::DayTime),
+ true,
+ ),
]);
explicit_options.schema = Some(&schema);
explicit_options.has_header = false;
diff --git a/datafusion/substrait/tests/testdata/data.csv
b/datafusion/substrait/tests/testdata/data.csv
index 1b85b166b1..ef2766d295 100644
--- a/datafusion/substrait/tests/testdata/data.csv
+++ b/datafusion/substrait/tests/testdata/data.csv
@@ -1,3 +1,3 @@
a,b,c,d,e,f
-1,2.0,2020-01-01,false,4294967296,'a'
-3,4.5,2020-01-01,true,2147483648,'b'
\ No newline at end of file
+1,2.0,2020-01-01,false,4294967295,'a'
+3,4.5,2020-01-01,true,2147483648,'b'
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
