Add unit test

Author: LiaCastaneda

datafusion/core/src/physical_planner.rs

@@ -2205,8 +2205,11 @@ mod tests {
     use datafusion_common::{assert_contains, DFSchemaRef, TableReference};
     use datafusion_execution::runtime_env::RuntimeEnv;
     use datafusion_execution::TaskContext;
-    use datafusion_expr::{col, lit, LogicalPlanBuilder, UserDefinedLogicalNodeCore};
+    use datafusion_expr::{
+        col, lit, LogicalPlanBuilder, Operator, UserDefinedLogicalNodeCore,
+    };
     use datafusion_functions_aggregate::expr_fn::sum;
+    use datafusion_physical_expr::expressions::{BinaryExpr, IsNotNullExpr};
     use datafusion_physical_expr::EquivalenceProperties;
     use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
 
@@ -2723,6 +2726,49 @@ mod tests {
 
         assert_eq!(col.name(), "metric:avg");
     }
+
+    #[tokio::test]
+    async fn test_maybe_fix_nested_column_name_with_colon() {
+        let schema =
+            Schema::new(vec![Field::new("example_column", DataType::Int32, false)]);
+        let schema_ref: SchemaRef = Arc::new(schema);
+
+        // Construct the nested expr
+        let col_expr =
+            Arc::new(Column::new("example_column:1", 0)) as Arc<dyn PhysicalExpr>;
+        let is_not_null_expr = Arc::new(IsNotNullExpr::new(col_expr.clone()));
+
+        // Create a binary expression and put the column inside
+        let binary_expr = Arc::new(BinaryExpr::new(
+            is_not_null_expr.clone(),
+            Operator::Or,
+            is_not_null_expr.clone(),
+        )) as Arc<dyn PhysicalExpr>;
+
+        let fixed_expr =
+            maybe_fix_physical_column_name(Ok(binary_expr), &schema_ref).unwrap();
+
+        let bin = fixed_expr
+            .as_any()
+            .downcast_ref::<BinaryExpr>()
+            .expect("Expected BinaryExpr");
+
+        // Check that both sides where renamed
+        for expr in &[bin.left(), bin.right()] {
+            let is_not_null = expr
+                .as_any()
+                .downcast_ref::<IsNotNullExpr>()
+                .expect("Expected IsNotNull");
+
+            let col = is_not_null
+                .arg()
+                .as_any()
+                .downcast_ref::<Column>()
+                .expect("Expected Column");
+
+            assert_eq!(col.name(), "example_column");
+        }
+    }
     struct ErrorExtensionPlanner {}
 
     #[async_trait]