Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Support Substrait's IntervalCompound type/literal instead of interval-month-day-nano UDT #12112

Merged
merged 8 commits into from
Oct 29, 2024
Merged

feat: Support Substrait's IntervalCompound type/literal instead of interval-month-day-nano UDT #12112

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e18ec17
feat(substrait): use IntervalCompound instead of interval-month-day-n…
Blizzara Aug 20, 2024
8b628e0
clippy
Blizzara Oct 28, 2024
36402bc
more clippy
Blizzara Oct 28, 2024
94ee836
even more clippy
Blizzara Oct 28, 2024
79f6341
fix precision exponent
Blizzara Oct 29, 2024
ea8bc4a
add a test
Blizzara Oct 29, 2024
d527b4e
update deprecation version
Blizzara Oct 29, 2024
8c215ea
update deprecation comments
Blizzara Oct 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 48 additions & 9 deletions datafusion/substrait/src/logical_plan/consumer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,17 +42,18 @@ use crate::variation_const::{
DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF,
DECIMAL_128_TYPE_VARIATION_REF, DECIMAL_256_TYPE_VARIATION_REF,
DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF,
INTERVAL_MONTH_DAY_NANO_TYPE_NAME, LARGE_CONTAINER_TYPE_VARIATION_REF,
UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
LARGE_CONTAINER_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
VIEW_CONTAINER_TYPE_VARIATION_REF,
};
#[allow(deprecated)]
use crate::variation_const::{
INTERVAL_DAY_TIME_TYPE_REF, INTERVAL_MONTH_DAY_NANO_TYPE_REF,
INTERVAL_YEAR_MONTH_TYPE_REF, TIMESTAMP_MICRO_TYPE_VARIATION_REF,
TIMESTAMP_MILLI_TYPE_VARIATION_REF, TIMESTAMP_NANO_TYPE_VARIATION_REF,
TIMESTAMP_SECOND_TYPE_VARIATION_REF,
INTERVAL_DAY_TIME_TYPE_REF, INTERVAL_MONTH_DAY_NANO_TYPE_NAME,
INTERVAL_MONTH_DAY_NANO_TYPE_REF, INTERVAL_YEAR_MONTH_TYPE_REF,
TIMESTAMP_MICRO_TYPE_VARIATION_REF, TIMESTAMP_MILLI_TYPE_VARIATION_REF,
TIMESTAMP_NANO_TYPE_VARIATION_REF, TIMESTAMP_SECOND_TYPE_VARIATION_REF,
};
use datafusion::arrow::array::{new_empty_array, AsArray};
use datafusion::arrow::temporal_conversions::NANOSECONDS;
use datafusion::common::scalar::ScalarStructBuilder;
use datafusion::dataframe::DataFrame;
use datafusion::logical_expr::expr::InList;
Expand All @@ -71,10 +72,10 @@ use datafusion::{
use std::collections::HashSet;
use std::sync::Arc;
use substrait::proto::exchange_rel::ExchangeKind;
use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode;
use substrait::proto::expression::literal::user_defined::Val;
use substrait::proto::expression::literal::{
IntervalDayToSecond, IntervalYearToMonth, UserDefined,
interval_day_to_second, IntervalCompound, IntervalDayToSecond, IntervalYearToMonth,
UserDefined,
};
use substrait::proto::expression::subquery::SubqueryType;
use substrait::proto::expression::{self, FieldReference, Literal, ScalarFunction};
Expand Down Expand Up @@ -1831,8 +1832,13 @@ fn from_substrait_type(
Ok(DataType::Interval(IntervalUnit::YearMonth))
}
r#type::Kind::IntervalDay(_) => Ok(DataType::Interval(IntervalUnit::DayTime)),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly out of scope of this PR, but do we need to checking the precision on this type now?

Based on https://github.com/substrait-io/substrait/blob/683f4179a058c2c99c04501b920a48ff372356ff/proto/substrait/type.proto#L132-L140 DataFusion can only handle this type with precision loss when precision is set to 0 or 3 (seconds or milliseconds).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup see #12112 (comment). I think we should change it, maybe with some option so that user can decide if they want to throw or accept the precision loss. But will be a separate PR :)

r#type::Kind::IntervalCompound(_) => {
Ok(DataType::Interval(IntervalUnit::MonthDayNano))
}
r#type::Kind::UserDefined(u) => {
// Kept for backwards compatibility, use IntervalCompound instead
if let Some(name) = extensions.types.get(&u.type_reference) {
#[allow(deprecated)]
match name.as_ref() {
INTERVAL_MONTH_DAY_NANO_TYPE_NAME => Ok(DataType::Interval(IntervalUnit::MonthDayNano)),
_ => not_impl_err!(
Expand All @@ -1842,7 +1848,7 @@ fn from_substrait_type(
),
}
} else {
// Kept for backwards compatibility, new plans should include the extension instead
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: I would suggest updating this warning here and further down to something like

// Kept for backwards compatibility
// Producers should use IntervalCompound instead

to make it clearer what/who should be using the new type.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

8c215ea

// Kept for backwards compatibility, use IntervalCompound instead
#[allow(deprecated)]
match u.type_reference {
// Kept for backwards compatibility, use IntervalYear instead
Expand Down Expand Up @@ -2275,6 +2281,7 @@ fn from_substrait_literal(
subseconds,
precision_mode,
})) => {
use interval_day_to_second::PrecisionMode;
// DF only supports millisecond precision, so for any more granular type we lose precision
let milliseconds = match precision_mode {
Some(PrecisionMode::Microseconds(ms)) => ms / 1000,
Expand All @@ -2299,6 +2306,35 @@ fn from_substrait_literal(
Some(LiteralType::IntervalYearToMonth(IntervalYearToMonth { years, months })) => {
ScalarValue::new_interval_ym(*years, *months)
}
Some(LiteralType::IntervalCompound(IntervalCompound {
interval_year_to_month,
interval_day_to_second,
})) => match (interval_year_to_month, interval_day_to_second) {
(
Some(IntervalYearToMonth { years, months }),
Some(IntervalDayToSecond {
days,
seconds,
subseconds,
precision_mode:
Some(interval_day_to_second::PrecisionMode::Precision(p)),
}),
) => {
if *p < 0 || *p > 9 {
return plan_err!(
"Unsupported Substrait interval day to second precision: {}",
p
);
}
let nanos = *subseconds * i64::pow(10, (*p - 9) as u32);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't precision handling here and on L2294 for IntervalDayToSecond be similar? Seems like only 0, 3, 6, 9 are supported over there. This one makes more sense to me, just curious why they are different.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is a very fair question! I think in general the answer is "yes", but there's a catch - the IntervalDayTime in Arrow only supports millisecond precision, so it cannot be done with a single i64::pow. It could be done with two (making it more complicated), or with a fp64::pow (leading to floating point ops).

But there is an additional question of should we allow the silent precision-losing operations (and the answer is probably no, at least not w/o some user config option) so I think that parts needs a better look anyways and I was thinking of doing it separately.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sounds good. I think there are two questions (probably unrelated?) to explore here: 1) precision-losing (for p>3) and 2) whether in-between values are valid (p==1, p==2)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm probably butchering this but should this be 9 - *p instead of *p - 9? If precision is 0 then we want 10^9 right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No I think you’re right 🙈 math is hard.. I’ll fix this tomorrow. Thanks for catching it!

Copy link
Contributor

@alamb alamb Oct 29, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No I think you’re right 🙈 math is hard.. I’ll fix this tomorrow. Thanks for catching it!

Marking this PR as draft so we don't accidentally merge it until it is fixed

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed in 79f6341 and added a test in ea8bc4a, thanks again for catching this @westonpace!

Thanks @alamb, I converted back to PR now and I think it's ready to merge :)

ScalarValue::new_interval_mdn(
*years * 12 + months,
*days,
*seconds as i64 * NANOSECONDS + nanos,
)
}
_ => return plan_err!("Substrait compound interval missing components"),
},
Some(LiteralType::FixedChar(c)) => ScalarValue::Utf8(Some(c.clone())),
Some(LiteralType::UserDefined(user_defined)) => {
// Helper function to prevent duplicating this code - can be inlined once the non-extension path is removed
Expand Down Expand Up @@ -2329,6 +2365,8 @@ fn from_substrait_literal(

if let Some(name) = extensions.types.get(&user_defined.type_reference) {
match name.as_ref() {
// Kept for backwards compatibility - new plans should use IntervalCompound instead
#[allow(deprecated)]
INTERVAL_MONTH_DAY_NANO_TYPE_NAME => {
interval_month_day_nano(user_defined)?
}
Expand Down Expand Up @@ -2379,6 +2417,7 @@ fn from_substrait_literal(
milliseconds,
}))
}
// Kept for backwards compatibility, use IntervalCompound instead
INTERVAL_MONTH_DAY_NANO_TYPE_REF => {
interval_month_day_nano(user_defined)?
}
Expand Down
115 changes: 24 additions & 91 deletions datafusion/substrait/src/logical_plan/producer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ use crate::variation_const::{
UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
};
use datafusion::arrow::array::{Array, GenericListArray, OffsetSizeTrait};
use datafusion::arrow::temporal_conversions::NANOSECONDS;
use datafusion::common::{
exec_err, internal_err, not_impl_err, plan_err, substrait_datafusion_err,
substrait_err, DFSchemaRef, ToDFSchema,
Expand All @@ -58,8 +59,8 @@ use substrait::proto::exchange_rel::{ExchangeKind, RoundRobin, ScatterFields};
use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode;
use substrait::proto::expression::literal::map::KeyValue;
use substrait::proto::expression::literal::{
user_defined, IntervalDayToSecond, IntervalYearToMonth, List, Map,
PrecisionTimestamp, Struct, UserDefined,
IntervalCompound, IntervalDayToSecond, IntervalYearToMonth, List, Map,
PrecisionTimestamp, Struct,
};
use substrait::proto::expression::subquery::InPredicate;
use substrait::proto::expression::window_function::BoundsType;
Expand Down Expand Up @@ -1489,16 +1490,14 @@ fn to_substrait_type(
})),
}),
IntervalUnit::MonthDayNano => {
// Substrait doesn't currently support this type, so we represent it as a UDT
Ok(substrait::proto::Type {
kind: Some(r#type::Kind::UserDefined(r#type::UserDefined {
type_reference: extensions.register_type(
INTERVAL_MONTH_DAY_NANO_TYPE_NAME.to_string(),
),
type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
nullability,
type_parameters: vec![],
})),
kind: Some(r#type::Kind::IntervalCompound(
r#type::IntervalCompound {
type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
nullability,
precision: 9, // nanos
},
)),
})
}
}
Expand Down Expand Up @@ -1892,23 +1891,21 @@ fn to_substrait_literal(
}),
DEFAULT_TYPE_VARIATION_REF,
),
ScalarValue::IntervalMonthDayNano(Some(i)) => {
// IntervalMonthDayNano is internally represented as a 128-bit integer, containing
// months (32bit), days (32bit), and nanoseconds (64bit)
let bytes = i.to_byte_slice();
(
LiteralType::UserDefined(UserDefined {
type_reference: extensions
.register_type(INTERVAL_MONTH_DAY_NANO_TYPE_NAME.to_string()),
type_parameters: vec![],
val: Some(user_defined::Val::Value(ProtoAny {
type_url: INTERVAL_MONTH_DAY_NANO_TYPE_NAME.to_string(),
value: bytes.to_vec().into(),
})),
ScalarValue::IntervalMonthDayNano(Some(i)) => (
LiteralType::IntervalCompound(IntervalCompound {
interval_year_to_month: Some(IntervalYearToMonth {
years: i.months / 12,
months: i.months % 12,
}),
DEFAULT_TYPE_VARIATION_REF,
)
}
interval_day_to_second: Some(IntervalDayToSecond {
days: i.days,
seconds: (i.nanoseconds / NANOSECONDS) as i32,
subseconds: i.nanoseconds % NANOSECONDS,
precision_mode: Some(PrecisionMode::Precision(9)), // nanoseconds
}),
}),
DEFAULT_TYPE_VARIATION_REF,
),
ScalarValue::IntervalDayTime(Some(i)) => (
LiteralType::IntervalDayToSecond(IntervalDayToSecond {
days: i.days,
Expand Down Expand Up @@ -2310,39 +2307,6 @@ mod test {
Ok(())
}

#[test]
fn custom_type_literal_extensions() -> Result<()> {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sadly as we no longer produce custom types, there's no easy way to test for them either :/
maybe if we one day implement some UDT support then that can be used for testing here instead. But for now there's no good replacement.

let mut extensions = Extensions::default();
// IntervalMonthDayNano is represented as a custom type in Substrait
let scalar = ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(
17, 25, 1234567890,
)));
let substrait_literal = to_substrait_literal(&scalar, &mut extensions)?;
let roundtrip_scalar =
from_substrait_literal_without_names(&substrait_literal, &extensions)?;
assert_eq!(scalar, roundtrip_scalar);

assert_eq!(
extensions,
Extensions {
functions: HashMap::new(),
types: HashMap::from([(
0,
INTERVAL_MONTH_DAY_NANO_TYPE_NAME.to_string()
)]),
type_variations: HashMap::new(),
}
);

// Check we fail if we don't propagate extensions
assert!(from_substrait_literal_without_names(
&substrait_literal,
&Extensions::default()
)
.is_err());
Ok(())
}

#[test]
fn round_trip_types() -> Result<()> {
round_trip_type(DataType::Boolean)?;
Expand Down Expand Up @@ -2424,37 +2388,6 @@ mod test {
Ok(())
}

#[test]
fn custom_type_extensions() -> Result<()> {
let mut extensions = Extensions::default();
// IntervalMonthDayNano is represented as a custom type in Substrait
let dt = DataType::Interval(IntervalUnit::MonthDayNano);

let substrait = to_substrait_type(&dt, true, &mut extensions)?;
let roundtrip_dt = from_substrait_type_without_names(&substrait, &extensions)?;
assert_eq!(dt, roundtrip_dt);

assert_eq!(
extensions,
Extensions {
functions: HashMap::new(),
types: HashMap::from([(
0,
INTERVAL_MONTH_DAY_NANO_TYPE_NAME.to_string()
)]),
type_variations: HashMap::new(),
}
);

// Check we fail if we don't propagate extensions
assert!(
from_substrait_type_without_names(&substrait, &Extensions::default())
.is_err()
);

Ok(())
}

#[test]
fn named_struct_names() -> Result<()> {
let mut extensions = Extensions::default();
Expand Down
6 changes: 5 additions & 1 deletion datafusion/substrait/src/variation_const.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,12 +96,16 @@ pub const INTERVAL_DAY_TIME_TYPE_REF: u32 = 2;
/// [`ScalarValue::IntervalMonthDayNano`]: datafusion::common::ScalarValue::IntervalMonthDayNano
#[deprecated(
since = "41.0.0",
note = "Use Substrait `UserDefinedType` with name `INTERVAL_MONTH_DAY_NANO_TYPE_NAME` instead"
note = "Use Substrait `IntervalCompund` type instead"
)]
pub const INTERVAL_MONTH_DAY_NANO_TYPE_REF: u32 = 3;

/// For [`DataType::Interval`] with [`IntervalUnit::MonthDayNano`].
///
/// [`DataType::Interval`]: datafusion::arrow::datatypes::DataType::Interval
/// [`IntervalUnit::MonthDayNano`]: datafusion::arrow::datatypes::IntervalUnit::MonthDayNano
#[deprecated(
since = "42.1.0",
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess that should be next version, what will it be? :)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

43 I believe: #12470

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated in d527b4e

note = "Use Substrait `IntervalCompund` type instead"
)]
pub const INTERVAL_MONTH_DAY_NANO_TYPE_NAME: &str = "interval-month-day-nano";
17 changes: 0 additions & 17 deletions datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,23 +230,6 @@ async fn select_with_reused_functions() -> Result<()> {
Ok(())
}

#[tokio::test]
async fn roundtrip_udt_extensions() -> Result<()> {
let ctx = create_context().await?;
let proto =
roundtrip_with_ctx("SELECT INTERVAL '1 YEAR 1 DAY 1 SECOND' FROM data", ctx)
.await?;
let expected_type = SimpleExtensionDeclaration {
mapping_type: Some(MappingType::ExtensionType(ExtensionType {
extension_uri_reference: u32::MAX,
type_anchor: 0,
name: "interval-month-day-nano".to_string(),
})),
};
assert_eq!(proto.extensions, vec![expected_type]);
Ok(())
}

#[tokio::test]
async fn select_with_filter_date() -> Result<()> {
roundtrip("SELECT * FROM data WHERE c > CAST('2020-01-01' AS DATE)").await
Expand Down
Loading