chore: fix emit_to calls

apache · Sep 30, 2024 · 268aa91 · 268aa91
1 parent a60d599
commit 268aa91
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 60 deletions.
diff --git a/datafusion/functions-aggregate-common/src/aggregate/min_max/groups_accumulator_max_view.rs b/datafusion/functions-aggregate-common/src/aggregate/min_max/groups_accumulator_max_view.rs
@@ -81,15 +81,11 @@ impl GroupsAccumulator for GroupsAccumulatorMaxStringView {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let num_groups = match emit_to {
-            EmitTo::All => self.states.len(),
-            EmitTo::First(n) => std::cmp::max(n, self.states.len()),
-        };
+        let states = emit_to.take_needed(&mut self.states);
 
         let mut builder = BinaryViewBuilder::new();
 
-        for i in 0..num_groups {
-            let value = &self.states[i];
+        for value in states {
             if value.is_empty() {
                 builder.append_null();
             } else {
@@ -98,28 +94,15 @@ impl GroupsAccumulator for GroupsAccumulatorMaxStringView {
         }
 
         let array = Arc::new(builder.finish()) as ArrayRef;
-
-        match emit_to {
-            EmitTo::All => {
-                self.states.clear();
-            }
-            EmitTo::First(n) => {
-                self.states.drain(0..n);
-            }
-        }
         Ok(array)
     }
 
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        let num_groups = match emit_to {
-            EmitTo::All => self.states.len(),
-            EmitTo::First(n) => std::cmp::max(n, self.states.len()),
-        };
+        let states = emit_to.take_needed(&mut self.states);
 
         let mut builder = BinaryViewBuilder::new();
 
-        for i in 0..num_groups {
-            let value = &self.states[i];
+        for value in states {
             if value.is_empty() {
                 builder.append_null();
             } else {
@@ -128,15 +111,6 @@ impl GroupsAccumulator for GroupsAccumulatorMaxStringView {
         }
 
         let array = Arc::new(builder.finish()) as ArrayRef;
-
-        match emit_to {
-            EmitTo::All => {
-                self.states.clear();
-            }
-            EmitTo::First(n) => {
-                self.states.drain(0..n);
-            }
-        }
         Ok(vec![array])
     }
 

diff --git a/datafusion/functions-aggregate-common/src/aggregate/min_max/groups_accumulator_min_view.rs b/datafusion/functions-aggregate-common/src/aggregate/min_max/groups_accumulator_min_view.rs
@@ -1,3 +1,4 @@
+use std::os::macos::raw::stat;
 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
@@ -81,15 +82,11 @@ impl GroupsAccumulator for GroupsAccumulatorMinStringView {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let num_groups = match emit_to {
-            EmitTo::All => self.states.len(),
-            EmitTo::First(n) => std::cmp::min(n, self.states.len()),
-        };
+        let states = emit_to.take_needed(&mut self.states);
 
         let mut builder = BinaryViewBuilder::new();
 
-        for i in 0..num_groups {
-            let value = &self.states[i];
+        for value in states {
             if value.is_empty() {
                 builder.append_null();
             } else {
@@ -98,28 +95,15 @@ impl GroupsAccumulator for GroupsAccumulatorMinStringView {
         }
 
         let array = Arc::new(builder.finish()) as ArrayRef;
-
-        match emit_to {
-            EmitTo::All => {
-                self.states.clear();
-            }
-            EmitTo::First(n) => {
-                self.states.drain(0..n);
-            }
-        }
         Ok(array)
     }
 
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        let num_groups = match emit_to {
-            EmitTo::All => self.states.len(),
-            EmitTo::First(n) => std::cmp::min(n, self.states.len()),
-        };
+        let states = emit_to.take_needed(&mut self.states);
 
         let mut builder = BinaryViewBuilder::new();
 
-        for i in 0..num_groups {
-            let value = &self.states[i];
+        for value in states {
             if value.is_empty() {
                 builder.append_null();
             } else {
@@ -128,15 +112,6 @@ impl GroupsAccumulator for GroupsAccumulatorMinStringView {
         }
 
         let array = Arc::new(builder.finish()) as ArrayRef;
-
-        match emit_to {
-            EmitTo::All => {
-                self.states.clear();
-            }
-            EmitTo::First(n) => {
-                self.states.drain(0..n);
-            }
-        }
         Ok(vec![array])
     }