Skip to content

Commit 68ead28

Browse files
authored
fix: specify roottype in substrait fieldreference (#13647)
* fix: specify roottype in fieldreference Signed-off-by: MBWhite <[email protected]> * Fix formatting Signed-off-by: MBWhite <[email protected]> * review suggestion Signed-off-by: MBWhite <[email protected]> --------- Signed-off-by: MBWhite <[email protected]>
1 parent 8b6daaf commit 68ead28

File tree

1 file changed

+24
-2
lines changed

1 file changed

+24
-2
lines changed

datafusion/substrait/src/logical_plan/producer.rs

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ use datafusion::prelude::Expr;
5656
use pbjson_types::Any as ProtoAny;
5757
use substrait::proto::exchange_rel::{ExchangeKind, RoundRobin, ScatterFields};
5858
use substrait::proto::expression::cast::FailureBehavior;
59+
use substrait::proto::expression::field_reference::{RootReference, RootType};
5960
use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode;
6061
use substrait::proto::expression::literal::map::KeyValue;
6162
use substrait::proto::expression::literal::{
@@ -2150,7 +2151,7 @@ fn try_to_substrait_field_reference(
21502151
}),
21512152
)),
21522153
})),
2153-
root_type: None,
2154+
root_type: Some(RootType::RootReference(RootReference {})),
21542155
})
21552156
}
21562157
_ => substrait_err!("Expect a `Column` expr, but found {expr:?}"),
@@ -2192,13 +2193,14 @@ fn substrait_field_ref(index: usize) -> Result<Expression> {
21922193
}),
21932194
)),
21942195
})),
2195-
root_type: None,
2196+
root_type: Some(RootType::RootReference(RootReference {})),
21962197
}))),
21972198
})
21982199
}
21992200

22002201
#[cfg(test)]
22012202
mod test {
2203+
22022204
use super::*;
22032205
use crate::logical_plan::consumer::{
22042206
from_substrait_extended_expr, from_substrait_literal_without_names,
@@ -2422,6 +2424,26 @@ mod test {
24222424
Ok(())
24232425
}
24242426

2427+
#[test]
2428+
fn to_field_reference() -> Result<()> {
2429+
let expression = substrait_field_ref(2)?;
2430+
2431+
match &expression.rex_type {
2432+
Some(RexType::Selection(field_ref)) => {
2433+
assert_eq!(
2434+
field_ref
2435+
.root_type
2436+
.clone()
2437+
.expect("root type should be set"),
2438+
RootType::RootReference(RootReference {})
2439+
);
2440+
}
2441+
2442+
_ => panic!("Should not be anything other than field reference"),
2443+
}
2444+
Ok(())
2445+
}
2446+
24252447
#[test]
24262448
fn named_struct_names() -> Result<()> {
24272449
let schema = DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![

0 commit comments

Comments
 (0)