Use Expr::qualified_name() and Column::new() to extract partition keys from window and aggregate operators (#355) (apache#17757) v51

Author: masonh22

datafusion/optimizer/src/push_down_filter.rs

@@ -1002,7 +1002,10 @@ impl OptimizerRule for PushDownFilter {
                 let extract_partition_keys = |func: &WindowFunction| {
                     func.partition_by
                         .iter()
-                        .map(|c| Column::from_qualified_name(c.schema_name().to_string()))
+                        .map(|c| {
+                            let (relation, name) = c.qualified_name();
+                            Column::new(relation, name)
+                        })
                         .collect::<HashSet<_>>()
                 };
                 let potential_partition_keys = window
@@ -1545,6 +1548,38 @@ mod tests {
         assert_optimized_plan_eq(plan, expected)
     }
 
+    /// verifies that filters with unusual identifier names are pushed down through window functions
+    #[test]
+    fn filter_window_special_identifier() -> Result<()> {
+        let schema = Schema::new(vec![
+            Field::new("$a", DataType::UInt32, false),
+            Field::new("$b", DataType::UInt32, false),
+            Field::new("$c", DataType::UInt32, false),
+        ]);
+        let table_scan = table_scan(Some("test"), &schema, None)?.build()?;
+
+        let window = Expr::WindowFunction(WindowFunction::new(
+            WindowFunctionDefinition::WindowUDF(
+                datafusion_functions_window::rank::rank_udwf(),
+            ),
+            vec![],
+        ))
+        .partition_by(vec![col("$a"), col("$b")])
+        .order_by(vec![col("$c").sort(true, true)])
+        .build()
+        .unwrap();
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .window(vec![window])?
+            .filter(col("$b").gt(lit(10i64)))?
+            .build()?;
+
+        let expected = "\
+        WindowAggr: windowExpr=[[rank() PARTITION BY [test.$a, test.$b] ORDER BY [test.$c ASC NULLS FIRST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]\
+        \n  TableScan: test, full_filters=[test.$b > Int64(10)]";
+        assert_optimized_plan_eq(plan, expected)
+    }
+
     /// verifies that when partitioning by 'a' and 'b', and filtering by 'a' and 'b', both 'a' and
     /// 'b' are pushed
     #[test]