neilconway commented on code in PR #23488:
URL: https://github.com/apache/datafusion/pull/23488#discussion_r3691502458
##########
datafusion/substrait/src/logical_plan/producer/expr/field_reference.rs:
##########
@@ -76,20 +84,35 @@ pub(crate) fn try_to_substrait_field_reference(
}
}
-/// Convert an outer reference column to a Substrait field reference.
-/// Outer reference columns reference columns from an outer query scope in
correlated subqueries.
-/// We convert them the same way as regular columns since the subquery plan
will be
-/// reconstructed with the proper schema context during consumption.
+/// Convert an outer reference column to a Substrait field reference with an
+/// `OuterReference` root type.
+///
+/// Outer reference columns reference columns from an enclosing query scope in
+/// correlated subqueries. The column is resolved against the producer's stack
+/// of outer schemas (pushed at each subquery boundary), innermost first, and
+/// the resulting `steps_out` records how many query boundaries the reference
+/// crosses (`steps_out = 1` is the immediately enclosing query).
pub fn from_outer_reference_column(
+ producer: &mut impl SubstraitProducer,
col: &Column,
- schema: &DFSchemaRef,
) -> datafusion::common::Result<Expression> {
- // OuterReferenceColumn is converted similarly to a regular column
reference.
- // The schema provided should be the schema context in which the outer
reference
- // column appears. During Substrait round-trip, the consumer will
reconstruct
- // the outer reference based on the subquery context.
- let index = schema.index_of_column(col)?;
- substrait_field_ref(index)
+ let mut steps_out = 1;
+ while let Some(outer_schema) = producer.get_outer_schema(steps_out) {
+ if let Some(index) = outer_schema.maybe_index_of_column(col) {
+ return substrait_field_ref_with_root(
+ index,
+ RootType::OuterReference(OuterReference {
+ steps_out: steps_out as u32,
+ }),
+ );
+ }
+ steps_out += 1;
+ }
+ substrait_err!(
+ "Outer reference column '{col}' could not be resolved against any
outer \
+ query schema. If using a custom SubstraitProducer, ensure it
maintains \
+ the outer schema stack
(push_outer_schema/pop_outer_schema/get_outer_schema)"
+ )
Review Comment:
How about adding a test case for this code path?
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -751,6 +751,187 @@ async fn roundtrip_not_exists_substrait() -> Result<()> {
Ok(())
}
+/// Assert that the Substrait plan contains a field reference with an
+/// `OuterReference` root type at the given `steps_out` depth.
+fn assert_contains_outer_reference(proto: &Plan, steps_out: u32) {
+ let proto_str = format!("{proto:?}");
+ let expected = format!("OuterReference {{ steps_out: {steps_out} }}");
+ assert!(
+ proto_str.contains(&expected),
+ "expected Substrait plan to contain an outer reference with
steps_out={steps_out}"
+ );
+}
Review Comment:
Would be nice to avoid unit tests that match on `format!` / `Debug` output,
if possible.
We could perhaps do this by serializing the plan to JSON and walking the
JSON tree. For example, Claude came up with this:
```rust
/// (steps_out, field index) of every OuterReference in the plan.
fn outer_references(plan: &Plan) -> Vec<(u32, u32)> {
fn walk(value: &serde_json::Value, out: &mut Vec<(u32, u32)>) {
match value {
serde_json::Value::Object(map) => {
// A FieldReference serializes its root type and direct
reference
// as sibling keys, e.g.:
// { "directReference": { "structField": { "field": 3 } },
// "outerReference": { "stepsOut": 1 } }
// proto3 JSON omits zero-valued fields, hence the
unwrap_or(0)s.
if let Some(outer) = map.get("outerReference") {
let steps =
outer.get("stepsOut").and_then(|v|
v.as_u64()).unwrap_or(0);
let field = map
.get("directReference")
.and_then(|d| d.get("structField"))
.and_then(|s| s.get("field"))
.and_then(|f| f.as_u64())
.unwrap_or(0);
out.push((steps as u32, field as u32));
}
map.values().for_each(|v| walk(v, out));
}
serde_json::Value::Array(items) => items.iter().for_each(|v|
walk(v, out)),
_ => {}
}
}
let mut refs = vec![];
walk(&serde_json::to_value(plan).expect("Plan serializes to JSON"), &mut
refs);
refs.sort_unstable(); // key order in the JSON walk isn't proto field
order
refs
}
```
Also nice to check that the field index is what we expect, not just
`steps_out`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
