Skip to content
Open
273 changes: 221 additions & 52 deletions Cargo.lock

Large diffs are not rendered by default.

5 changes: 3 additions & 2 deletions datafusion/substrait/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,9 @@ half = { workspace = true }
itertools = { workspace = true }
object_store = { workspace = true }
# We need to match the version in substrait, so we don't use the workspace version here
pbjson-types = { version = "0.8.0" }
pbjson-types = { version = "0.9.0" }
prost = { workspace = true }
substrait = { version = "0.63.0", features = ["serde"] }
substrait = { version = "0.65.0", features = ["serde"] }
url = { workspace = true }
tokio = { workspace = true, features = ["fs"] }

Expand All @@ -59,6 +59,7 @@ insta = { workspace = true }
default = ["physical"]
physical = ["datafusion/parquet"]
protoc = ["substrait/protoc"]
protox = ["substrait/protox"]

[package.metadata.docs.rs]
# Use default features ("physical") for docs, plus "protoc". "protoc" is needed
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,19 @@ fn resolve_outer_reference(
outer_ref: &substrait::proto::expression::field_reference::OuterReference,
field_idx: usize,
) -> datafusion::common::Result<Expr> {
let steps_out = outer_ref.steps_out as usize;
use substrait::proto::expression::field_reference::outer_reference::OuterReferenceType;
// `StepsOut` is deprecated in favour of `RelReference`, but a relation
// reference needs anchors that DataFusion does not assign.
#[expect(deprecated)]
let steps_out = match outer_ref.outer_reference_type {
Some(OuterReferenceType::StepsOut(steps_out)) => steps_out as usize,
Some(OuterReferenceType::RelReference(_)) => {
return not_impl_err!(
"OuterReference by relation reference is not supported"
);
}
None => return substrait_err!("OuterReference without a reference type"),
};
let Some(outer_schema) = consumer.get_outer_schema(steps_out) else {
return substrait_err!(
"OuterReference with steps_out={steps_out} \
Expand Down Expand Up @@ -152,6 +164,34 @@ mod tests {
);
}

#[tokio::test]
async fn unsupported_outer_reference_forms_are_rejected() {
use substrait::proto::expression::field_reference::{
OuterReference, outer_reference::OuterReferenceType,
};

let extensions = Extensions::default();
let session_state = SessionContext::new().state();
let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state);
for (reference_type, message) in [
(
Some(OuterReferenceType::RelReference(1)),
"by relation reference is not supported",
),
(None, "without a reference type"),
] {
let mut reference = lambda_field_ref(0, 0);
reference.root_type = Some(RootType::OuterReference(OuterReference {
outer_reference_type: reference_type,
}));
let error =
from_field_reference(&consumer, &reference, DFSchema::empty_ref())
.await
.unwrap_err();
assert_contains!(error.to_string(), message);
}
}

fn lambda_field_ref(field: i32, steps_out: u32) -> FieldReference {
FieldReference {
reference_type: Some(field_reference::ReferenceType::DirectReference(
Expand Down
62 changes: 11 additions & 51 deletions datafusion/substrait/src/logical_plan/consumer/expr/literal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,7 @@ use crate::variation_const::{
INTERVAL_DAY_TIME_TYPE_REF, INTERVAL_MONTH_DAY_NANO_TYPE_NAME,
INTERVAL_MONTH_DAY_NANO_TYPE_REF, INTERVAL_YEAR_MONTH_TYPE_REF,
LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
TIME_64_TYPE_VARIATION_REF, TIMESTAMP_MICRO_TYPE_VARIATION_REF,
TIMESTAMP_MILLI_TYPE_VARIATION_REF, TIMESTAMP_NANO_TYPE_VARIATION_REF,
TIMESTAMP_SECOND_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
VIEW_CONTAINER_TYPE_VARIATION_REF,
};
use datafusion::arrow::array::{AsArray, MapArray, new_empty_array};
Expand All @@ -48,7 +46,6 @@ use substrait::proto::expression::Literal;
use substrait::proto::expression::literal::user_defined::{TypeAnchorType, Val};
use substrait::proto::expression::literal::{
IntervalCompound, IntervalDayToSecond, IntervalYearToMonth, LiteralType,
interval_day_to_second,
};

pub async fn from_literal(
Expand Down Expand Up @@ -104,28 +101,6 @@ pub(crate) fn from_substrait_literal(
},
Some(LiteralType::Fp32(f)) => ScalarValue::Float32(Some(*f)),
Some(LiteralType::Fp64(f)) => ScalarValue::Float64(Some(*f)),
#[expect(deprecated)]
Some(LiteralType::Timestamp(t)) => {
// Kept for backwards compatibility, new plans should use PrecisionTimestamp(Tz) instead
#[expect(deprecated)]
match lit.type_variation_reference {
TIMESTAMP_SECOND_TYPE_VARIATION_REF => {
ScalarValue::TimestampSecond(Some(*t), None)
}
TIMESTAMP_MILLI_TYPE_VARIATION_REF => {
ScalarValue::TimestampMillisecond(Some(*t), None)
}
TIMESTAMP_MICRO_TYPE_VARIATION_REF => {
ScalarValue::TimestampMicrosecond(Some(*t), None)
}
TIMESTAMP_NANO_TYPE_VARIATION_REF => {
ScalarValue::TimestampNanosecond(Some(*t), None)
}
others => {
return substrait_err!("Unknown type variation reference {others}");
}
}
}
Some(LiteralType::PrecisionTimestamp(pt)) => match pt.precision {
0 => ScalarValue::TimestampSecond(Some(pt.value), None),
3 => ScalarValue::TimestampMillisecond(Some(pt.value), None),
Expand Down Expand Up @@ -381,29 +356,17 @@ pub(crate) fn from_substrait_literal(
days,
seconds,
subseconds,
precision_mode,
precision,
})) => {
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 {
#[expect(deprecated)]
Some(PrecisionMode::Microseconds(ms)) => ms / 1000,
None => {
if *subseconds != 0 {
return substrait_err!(
"Cannot set subseconds field of IntervalDayToSecond without setting precision"
);
} else {
0_i32
}
}
Some(PrecisionMode::Precision(0)) => *subseconds as i32 * 1000,
Some(PrecisionMode::Precision(3)) => *subseconds as i32,
Some(PrecisionMode::Precision(6)) => (subseconds / 1000) as i32,
Some(PrecisionMode::Precision(9)) => (subseconds / 1000 / 1000) as i32,
_ => {
let milliseconds = match precision {
0 => *subseconds as i32 * 1000,
3 => *subseconds as i32,
6 => (subseconds / 1000) as i32,
9 => (subseconds / 1000 / 1000) as i32,
p => {
return not_impl_err!(
"Unsupported Substrait interval day to second precision mode: {precision_mode:?}"
"Unsupported Substrait interval day to second precision: {p}"
);
}
};
Expand All @@ -423,8 +386,7 @@ pub(crate) fn from_substrait_literal(
days,
seconds,
subseconds,
precision_mode:
Some(interval_day_to_second::PrecisionMode::Precision(p)),
precision: p,
}),
) => {
if *p < 0 || *p > 9 {
Expand Down Expand Up @@ -607,9 +569,7 @@ mod tests {
days: 3,
seconds: 4,
subseconds: 5,
precision_mode: Some(
interval_day_to_second::PrecisionMode::Precision(6),
),
precision: 6,
}),
})),
};
Expand Down
17 changes: 9 additions & 8 deletions datafusion/substrait/src/logical_plan/consumer/expr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ pub async fn from_substrait_rex(
) -> datafusion::common::Result<Expr> {
match &expression.rex_type {
Some(t) => match t {
RexType::ExecutionContextVariable(_) => {
not_impl_err!("Execution context variables are not supported")
}
RexType::Literal(expr) => consumer.consume_literal(expr).await,
RexType::Selection(expr) => {
consumer.consume_field_reference(expr, input_schema).await
Expand Down Expand Up @@ -92,8 +95,6 @@ pub async fn from_substrait_rex(
consumer.consume_subquery(expr.as_ref(), input_schema).await
}
RexType::Nested(expr) => consumer.consume_nested(expr, input_schema).await,
#[expect(deprecated)]
RexType::Enum(expr) => consumer.consume_enum(expr, input_schema).await,
RexType::DynamicParameter(expr) => {
consumer.consume_dynamic_parameter(expr, input_schema).await
}
Expand Down Expand Up @@ -216,14 +217,14 @@ mod tests {
async fn window_function_with_range_unit_and_no_order_by()
-> datafusion::common::Result<()> {
let substrait = Expression {
rex_type: Some(RexType::WindowFunction(
rex_type: Some(RexType::WindowFunction(Box::new(
substrait::proto::expression::WindowFunction {
function_reference: 0,
bounds_type: BoundsType::Range as i32,
sorts: vec![],
..Default::default()
},
)),
))),
};

let mut consumer = test_consumer();
Expand All @@ -247,12 +248,12 @@ mod tests {
#[tokio::test]
async fn window_function_with_count() -> datafusion::common::Result<()> {
let substrait = Expression {
rex_type: Some(RexType::WindowFunction(
rex_type: Some(RexType::WindowFunction(Box::new(
substrait::proto::expression::WindowFunction {
function_reference: 0,
..Default::default()
},
)),
))),
};

let mut consumer = test_consumer();
Expand All @@ -274,13 +275,13 @@ mod tests {
#[tokio::test]
async fn window_function_with_invalid_invocation() {
let substrait = Expression {
rex_type: Some(RexType::WindowFunction(
rex_type: Some(RexType::WindowFunction(Box::new(
substrait::proto::expression::WindowFunction {
function_reference: 0,
invocation: 3,
..Default::default()
},
)),
))),
};

let mut consumer = test_consumer();
Expand Down
Loading
Loading