Merge pull request #2287 from mabel-dev/#2276-3

#2276
mabel-dev · Jan 21, 2025 · 7be91b2 · 7be91b2
2 parents c0091bc + 49f2b82
commit 7be91b2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/opteryx/__version__.py b/opteryx/__version__.py
@@ -1,4 +1,4 @@
-__build__ = 1006
+__build__ = 1008
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/opteryx/operators/simple_aggregate_node.py b/opteryx/operators/simple_aggregate_node.py
@@ -49,7 +49,10 @@ def collect(self, values):
             elif self.aggregate_type == "MAX":
                 self.current_value = pyarrow.compute.max(values).as_py()
             elif self.aggregate_type == "COUNT_DISTINCT":
-                self.current_value = count_distinct(values.to_numpy().astype(object), FlatHashSet())
+                values = values.to_numpy()
+                if values.dtype != object:
+                    values = values.astype(object)
+                self.current_value = count_distinct(values, FlatHashSet())
             elif self.aggregate_type != "COUNT":
                 raise ValueError(f"Unsupported aggregate type: {self.aggregate_type}")
         else:
@@ -60,7 +63,10 @@ def collect(self, values):
             elif self.aggregate_type == "MAX":
                 self.current_value = max(self.current_value, pyarrow.compute.max(values).as_py())
             elif self.aggregate_type == "COUNT_DISTINCT":
-                self.current_value = count_distinct(values.to_numpy().astype(object), self.current_value)
+                values = values.to_numpy()
+                if values.dtype != object:
+                    values = values.astype(object)
+                self.current_value = count_distinct(values, self.current_value)
             elif self.aggregate_type != "COUNT":
                 raise ValueError(f"Unsupported aggregate type: {self.aggregate_type}")