aggregates

Author: joocer

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1691
+__build__ = 1695
 __author__ = "@joocer"
-__version__ = "0.26.0-beta.1691"
+__version__ = "0.26.0-beta.1695"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/connectors/cql_connector.py

@@ -109,6 +109,7 @@ def read_dataset(  # type:ignore
         predicates: list = None,
         chunk_size: int = INITIAL_CHUNK_SIZE,  # type:ignore
         limit: int = None,
+        **kwargs,
     ) -> Generator[pyarrow.Table, None, None]:  # type:ignore
         self.chunk_size = chunk_size
 

opteryx/managers/expression/ops.py

@@ -228,7 +228,11 @@ def _inner_filter_operations(arr, operator, value):
         matches = compute.match_like(arr, value).to_numpy(False).astype(dtype=numpy.bool_)
         return numpy.invert(matches)
     if operator == "ILike":
-        return compute.match_like(arr, value, ignore_case=True).to_numpy(False).astype(dtype=numpy.bool_)
+        return (
+            compute.match_like(arr, value, ignore_case=True)
+            .to_numpy(False)
+            .astype(dtype=numpy.bool_)
+        )
     if operator == "NotILike":
         matches = compute.match_like(arr, value, ignore_case=True)
         return numpy.invert(matches)

opteryx/operators/aggregate_and_group_node.py

@@ -68,6 +68,7 @@ def __init__(self, properties: QueryProperties, **parameters):
         self.column_map, self.aggregate_functions = build_aggregations(self.aggregates)
 
         self.buffer = []
+        self.max_buffer_size = 50  # Process in chunks to avoid excessive memory usage
 
     @property
     def config(self):  # pragma: no cover
@@ -85,18 +86,19 @@ def execute(self, morsel: pyarrow.Table, **kwargs):
                 yield EOS
                 return
 
-            # merge all the morsels together into one table, selecting only the columns
-            # we're pretty sure we're going to use - this will fail for datasets
-            # larger than memory
-            table = pyarrow.concat_tables(
-                self.buffer,
-                promote_options="permissive",
-            )
+            # If we have partial results in buffer, do final aggregation
+            if len(self.buffer) > 0:
+                table = pyarrow.concat_tables(
+                    self.buffer,
+                    promote_options="permissive",
+                )
+                table = table.combine_chunks()
+                groups = table.group_by(self.group_by_columns)
+                groups = groups.aggregate(self.aggregate_functions)
+                self.buffer = [groups]  # Replace buffer with final result
 
-            # do the group by and aggregates
-            table = table.combine_chunks()
-            groups = table.group_by(self.group_by_columns)
-            groups = groups.aggregate(self.aggregate_functions)
+            # Now buffer has the final aggregated result
+            groups = self.buffer[0]
 
             # do the secondary activities for ARRAY_AGG
             for node in get_all_nodes_of_type(self.aggregates, select_nodes=(NodeType.AGGREGATOR,)):
@@ -135,4 +137,16 @@ def execute(self, morsel: pyarrow.Table, **kwargs):
         morsel = evaluate_and_append(self.groups, morsel)
 
         self.buffer.append(morsel)
+
+        # If buffer is full, do partial aggregation
+        if len(self.buffer) >= self.max_buffer_size:
+            table = pyarrow.concat_tables(
+                self.buffer,
+                promote_options="permissive",
+            )
+            table = table.combine_chunks()
+            groups = table.group_by(self.group_by_columns)
+            groups = groups.aggregate(self.aggregate_functions)
+            self.buffer = [groups]  # Replace buffer with partial result
+
         yield None

opteryx/operators/read_node.py

@@ -188,7 +188,7 @@ def to_mermaid(self, stats, nid):
             mermaid = f'NODE_{nid}[("**{self.node_type.upper()} (FUNCTION)**<br />'
             mermaid += f"{self.function}<br />"
         else:
-            mermaid = f'NODE_{nid}[(**"{self.node_type.upper()} ({self.connector.__type__})**<br />'
+            mermaid = f'NODE_{nid}[("**{self.node_type.upper()} ({self.connector.__type__})**<br />'
             mermaid += f"{self.connector.dataset}<br />"
         mermaid += BAR
         if self.columns:
@@ -259,7 +259,9 @@ def execute(self, morsel, **kwargs) -> Generator:
         arrow_schema = None
         start_clock = time.monotonic_ns()
         reader = self.connector.read_dataset(
-            columns=self.columns, predicates=self.predicates, limit=self.limit
+            columns=self.columns,
+            predicates=self.predicates,
+            limit=self.limit,
         )
         for morsel in reader:
             # try to make each morsel have the same schema

opteryx/operators/simple_aggregate_and_group_node.py

@@ -59,7 +59,7 @@ def build_finalizer_aggregations(aggregators):
 
 
 class SimpleAggregateAndGroupNode(BasePlanNode):
-    SIMPLE_AGGREGATES = {"SUM", "MIN", "MAX", "COUNT"}
+    SIMPLE_AGGREGATES = {"SUM", "MIN", "MAX", "COUNT", "AVG", "COUNT_DISTINCT"}
 
     def __init__(self, properties: QueryProperties, **parameters):
         BasePlanNode.__init__(self, properties=properties, **parameters)

opteryx/planner/optimizer/__init__.py

@@ -62,7 +62,6 @@
 from opteryx.models import QueryStatistics
 from opteryx.planner.logical_planner import LogicalPlan
 from opteryx.planner.optimizer.strategies import *
-from opteryx.planner.optimizer.strategies.join_groupby_pushdown import JoinGroupByPushdownStrategy
 
 from .strategies.optimization_strategy import OptimizerContext
 
@@ -84,7 +83,6 @@ def __init__(self, statistics: QueryStatistics):
             PredicatePushdownStrategy(statistics),
             ProjectionPushdownStrategy(statistics),
             JoinRewriteStrategy(statistics),
-            JoinGroupByPushdownStrategy(statistics),
             JoinOrderingStrategy(statistics),
             DistinctPushdownStrategy(statistics),
             OperatorFusionStrategy(statistics),