From 68ead283bd47de75d8aa79f6df5fe78a957df50c Mon Sep 17 00:00:00 2001
From: Matthew B White <matthew@mh-white.com>
Date: Fri, 13 Dec 2024 20:39:08 +0000
Subject: [PATCH] fix: specify roottype in substrait fieldreference (#13647)

* fix: specify roottype in fieldreference

Signed-off-by: MBWhite <whitemat@uk.ibm.com>

* Fix formatting

Signed-off-by: MBWhite <whitemat@uk.ibm.com>

* review suggestion

Signed-off-by: MBWhite <whitemat@uk.ibm.com>

---------

Signed-off-by: MBWhite <whitemat@uk.ibm.com>
---
 .../substrait/src/logical_plan/producer.rs    | 26 +++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs
index 26d71c7fd3e2..a128b90e6889 100644
--- a/datafusion/substrait/src/logical_plan/producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer.rs
@@ -56,6 +56,7 @@ use datafusion::prelude::Expr;
 use pbjson_types::Any as ProtoAny;
 use substrait::proto::exchange_rel::{ExchangeKind, RoundRobin, ScatterFields};
 use substrait::proto::expression::cast::FailureBehavior;
+use substrait::proto::expression::field_reference::{RootReference, RootType};
 use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode;
 use substrait::proto::expression::literal::map::KeyValue;
 use substrait::proto::expression::literal::{
@@ -2150,7 +2151,7 @@ fn try_to_substrait_field_reference(
                         }),
                     )),
                 })),
-                root_type: None,
+                root_type: Some(RootType::RootReference(RootReference {})),
             })
         }
         _ => substrait_err!("Expect a `Column` expr, but found {expr:?}"),
@@ -2192,13 +2193,14 @@ fn substrait_field_ref(index: usize) -> Result<Expression> {
                     }),
                 )),
             })),
-            root_type: None,
+            root_type: Some(RootType::RootReference(RootReference {})),
         }))),
     })
 }
 
 #[cfg(test)]
 mod test {
+
     use super::*;
     use crate::logical_plan::consumer::{
         from_substrait_extended_expr, from_substrait_literal_without_names,
@@ -2422,6 +2424,26 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn to_field_reference() -> Result<()> {
+        let expression = substrait_field_ref(2)?;
+
+        match &expression.rex_type {
+            Some(RexType::Selection(field_ref)) => {
+                assert_eq!(
+                    field_ref
+                        .root_type
+                        .clone()
+                        .expect("root type should be set"),
+                    RootType::RootReference(RootReference {})
+                );
+            }
+
+            _ => panic!("Should not be anything other than field reference"),
+        }
+        Ok(())
+    }
+
     #[test]
     fn named_struct_names() -> Result<()> {
         let schema = DFSchemaRef::new(DFSchema::try_from(Schema::new(vec![