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

@@ -1010,9 +1010,8 @@ impl OptimizerRule for PushDownFilter {
                                 .partition_by
                                 .iter()
                                 .map(|c| {
-                                    Column::from_qualified_name(
-                                        c.schema_name().to_string(),
-                                    )
+                                    let (relation, name) = c.qualified_name();
+                                    Column::new(relation, name)
                                 })
                                 .collect::<HashSet<_>>()
                         } else {
@@ -1550,6 +1549,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]