Merge branch 'main' into add_group_benchmarks_count_distinct

Author: coderfender

datafusion/functions-aggregate/benches/count_distinct.rs

@@ -239,19 +239,28 @@ fn count_distinct_groups_benchmark(c: &mut Criterion) {
                 })
             });
         } else {
+            let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
+            let mut group_rows: Vec<Vec<i64>> = vec![Vec::new(); num_groups];
+            for (idx, &group_idx) in group_indices.iter().enumerate() {
+                if arr.is_valid(idx) {
+                    group_rows[group_idx].push(arr.value(idx));
+                }
+            }
+            let group_arrays: Vec<ArrayRef> = group_rows
+                .iter()
+                .map(|rows| Arc::new(Int64Array::from(rows.clone())) as ArrayRef)
+                .collect();
+
             c.bench_function(&format!("count_distinct_groups {name}"), |b| {
                 b.iter(|| {
                     let mut accumulators: Vec<_> = (0..num_groups)
                         .map(|_| prepare_accumulator(DataType::Int64))
                         .collect();
 
-                    let arr = values.as_any().downcast_ref::<Int64Array>().unwrap();
-                    for (idx, group_idx) in group_indices.iter().enumerate() {
-                        if let Some(val) = arr.value(idx).into() {
-                            let single_val =
-                                Arc::new(Int64Array::from(vec![Some(val)])) as ArrayRef;
-                            accumulators[*group_idx]
-                                .update_batch(std::slice::from_ref(&single_val))
+                    for (group_idx, batch) in group_arrays.iter().enumerate() {
+                        if batch.len() > 0 {
+                            accumulators[group_idx]
+                                .update_batch(std::slice::from_ref(batch))
                                 .unwrap();
                         }
                     }

datafusion/physical-optimizer/src/window_topn.rs

@@ -238,15 +238,15 @@ impl PhysicalOptimizerRule for WindowTopN {
 fn extract_window_limit(
     predicate: &Arc<dyn datafusion_physical_expr::PhysicalExpr>,
 ) -> Option<(usize, usize)> {
-    let binary = predicate.as_any().downcast_ref::<BinaryExpr>()?;
+    let binary = predicate.downcast_ref::<BinaryExpr>()?;
     let op = binary.op();
     let left = binary.left();
     let right = binary.right();
 
     // Try Column op Literal
     if let (Some(col), Some(lit_val)) = (
-        left.as_any().downcast_ref::<Column>(),
-        right.as_any().downcast_ref::<Literal>(),
+        left.downcast_ref::<Column>(),
+        right.downcast_ref::<Literal>(),
     ) {
         let n = scalar_to_usize(lit_val.value())?;
         return match *op {
@@ -258,8 +258,8 @@ fn extract_window_limit(
 
     // Try Literal op Column (flipped)
     if let (Some(lit_val), Some(col)) = (
-        left.as_any().downcast_ref::<Literal>(),
-        right.as_any().downcast_ref::<Column>(),
+        left.downcast_ref::<Literal>(),
+        right.downcast_ref::<Column>(),
     ) {
         let n = scalar_to_usize(lit_val.value())?;
         return match *op {