pixeltable 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl

pixeltable/catalog/table_version.py

@@ -6,7 +6,7 @@ import inspect
 import logging
 import time
 import uuid
-from typing import TYPE_CHECKING, Any, Iterable, Optional
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional
 from uuid import UUID
 
 import sqlalchemy as sql
@@ -453,7 +453,9 @@ class TableVersion:
         self.idxs_by_name[idx_name] = idx_info
 
         # add the columns and update the metadata
-        status = self._add_columns([val_col, undo_col], conn)
+        # TODO support on_error='abort' for indices; it's tricky because of the way metadata changes are entangled
+        # with the database operations
+        status = self._add_columns([val_col, undo_col], conn, print_stats=False, on_error='ignore')
         # now create the index structure
         idx.create_index(self._store_idx_name(idx_id), val_col, conn)
 
@@ -478,7 +480,7 @@ class TableVersion:
             self._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
             _logger.info(f'Dropped index {idx_md.name} on table {self.name}')
 
-    def add_column(self, col: Column, print_stats: bool = False) -> UpdateStatus:
+    def add_column(self, col: Column, print_stats: bool, on_error: Literal['abort', 'ignore']) -> UpdateStatus:
         """Adds a column to the table.
         """
         assert not self.is_snapshot
@@ -498,9 +500,8 @@ class TableVersion:
         preceding_schema_version = self.schema_version
         self.schema_version = self.version
         with Env.get().engine.begin() as conn:
-            status = self._add_columns([col], conn, print_stats=print_stats)
+            status = self._add_columns([col], conn, print_stats=print_stats, on_error=on_error)
             _ = self._add_default_index(col, conn)
-            # TODO: what to do about errors?
             self._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
         _logger.info(f'Added column {col.name} to table {self.name}, new version: {self.version}')
 
@@ -512,7 +513,13 @@ class TableVersion:
         _logger.info(f'Column {col.name}: {msg}')
         return status
 
-    def _add_columns(self, cols: Iterable[Column], conn: sql.engine.Connection, print_stats: bool = False) -> UpdateStatus:
+    def _add_columns(
+        self,
+        cols: Iterable[Column],
+        conn: sql.engine.Connection,
+        print_stats: bool,
+        on_error: Literal['abort', 'ignore']
+    ) -> UpdateStatus:
         """Add and populate columns within the current transaction"""
         cols = list(cols)
         row_count = self.store_tbl.count(conn=conn)
@@ -550,10 +557,14 @@ class TableVersion:
             try:
                 plan.ctx.set_conn(conn)
                 plan.open()
-                num_excs = self.store_tbl.load_column(col, plan, value_expr_slot_idx, conn)
+                try:
+                    num_excs = self.store_tbl.load_column(col, plan, value_expr_slot_idx, conn, on_error)
+                except sql.exc.DBAPIError as exc:
+                    # Wrap the DBAPIError in an excs.Error to unify processing in the subsequent except block
+                    raise excs.Error(f'SQL error during execution of computed column `{col.name}`:\n{exc}') from exc
                 if num_excs > 0:
                     cols_with_excs.append(col)
-            except sql.exc.DBAPIError as e:
+            except excs.Error as exc:
                 self.cols.pop()
                 for col in cols:
                     # remove columns that we already added
@@ -564,7 +575,7 @@ class TableVersion:
                     del self.cols_by_id[col.id]
                 # we need to re-initialize the sqlalchemy schema
                 self.store_tbl.create_sa_tbl()
-                raise excs.Error(f'Error during SQL execution:\n{e}')
+                raise exc
             finally:
                 plan.close()
 
@@ -689,21 +700,30 @@ class TableVersion:
             plan = Planner.create_insert_plan(self, rows, ignore_errors=not fail_on_exception)
         else:
             plan = Planner.create_df_insert_plan(self, df, ignore_errors=not fail_on_exception)
+
+        # this is a base table; we generate rowids during the insert
+        def rowids() -> Iterator[int]:
+            while True:
+                rowid = self.next_rowid
+                self.next_rowid += 1
+                yield rowid
+
         if conn is None:
             with Env.get().engine.begin() as conn:
-                return self._insert(plan, conn, time.time(), print_stats)
+                return self._insert(plan, conn, time.time(), print_stats=print_stats, rowids=rowids())
         else:
-            return self._insert(plan, conn, time.time(), print_stats)
+            return self._insert(plan, conn, time.time(), print_stats=print_stats, rowids=rowids())
 
     def _insert(
-        self, exec_plan: 'exec.ExecNode', conn: sql.engine.Connection, timestamp: float, print_stats: bool = False,
+        self, exec_plan: 'exec.ExecNode', conn: sql.engine.Connection, timestamp: float, *,
+        rowids: Optional[Iterator[int]] = None, print_stats: bool = False,
     ) -> UpdateStatus:
         """Insert rows produced by exec_plan and propagate to views"""
         # we're creating a new version
         self.version += 1
         result = UpdateStatus()
-        num_rows, num_excs, cols_with_excs = self.store_tbl.insert_rows(exec_plan, conn, v_min=self.version)
-        self.next_rowid = num_rows
+        num_rows, num_excs, cols_with_excs = self.store_tbl.insert_rows(
+            exec_plan, conn, v_min=self.version, rowids=rowids)
         result.num_rows = num_rows
         result.num_excs = num_excs
         result.num_computed_values += exec_plan.ctx.num_computed_exprs * num_rows
@@ -714,7 +734,7 @@ class TableVersion:
         for view in self.mutable_views:
             from pixeltable.plan import Planner
             plan, _ = Planner.create_view_load_plan(view.path, propagates_insert=True)
-            status = view._insert(plan, conn, timestamp, print_stats)
+            status = view._insert(plan, conn, timestamp, print_stats=print_stats)
             result.num_rows += status.num_rows
             result.num_excs += status.num_excs
             result.num_computed_values += status.num_computed_values
@@ -751,9 +771,7 @@ class TableVersion:
                 raise excs.Error(f'Filter {analysis_info.filter} not expressible in SQL')
 
         with Env.get().engine.begin() as conn:
-            plan, updated_cols, recomputed_cols = (
-                Planner.create_update_plan(self.path, update_spec, [], where, cascade)
-            )
+            plan, updated_cols, recomputed_cols = Planner.create_update_plan(self.path, update_spec, [], where, cascade)
             from pixeltable.exprs import SqlElementCache
             result = self.propagate_update(
                 plan, where.sql_expr(SqlElementCache()) if where is not None else None, recomputed_cols,

pixeltable/catalog/table_version_path.py

@@ -91,14 +91,6 @@ class TableVersionPath:
         col = self.tbl_version.cols_by_name[col_name]
         return ColumnRef(col)
 
-    def __getitem__(self, index: object) -> Union[exprs.ColumnRef, pxt.DataFrame]:
-        """Return a ColumnRef for the given column name, or a DataFrame for the given slice.
-        """
-        if isinstance(index, str):
-            # basically <tbl>.<colname>
-            return self.__getattr__(index)
-        return pxt.DataFrame(self).__getitem__(index)
-
     def columns(self) -> list[Column]:
         """Return all user columns visible in this tbl version path, including columns from bases"""
         result = list(self.tbl_version.cols_by_name.values())

pixeltable/catalog/view.py

@@ -52,11 +52,11 @@ class View(Table):
 
     @classmethod
     def _create(
-            cls, dir_id: UUID, name: str, base: TableVersionPath, schema: Dict[str, Any],
-            predicate: 'pxt.exprs.Expr', is_snapshot: bool, num_retained_versions: int, comment: str,
+            cls, dir_id: UUID, name: str, base: TableVersionPath, additional_columns: Dict[str, Any],
+            predicate: Optional['pxt.exprs.Expr'], is_snapshot: bool, num_retained_versions: int, comment: str,
             iterator_cls: Optional[Type[ComponentIterator]], iterator_args: Optional[Dict]
     ) -> View:
-        columns = cls._create_columns(schema)
+        columns = cls._create_columns(additional_columns)
         cls._verify_schema(columns)
 
         # verify that filter can be evaluated in the context of the base

pixeltable/dataframe.py

@@ -8,7 +8,7 @@ import logging
 import mimetypes
 import traceback
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, Iterator, List, Optional, Set, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, Iterator, List, Optional, Sequence, Set, Tuple, Union
 
 import pandas as pd
 import pandas.io.formats.style
@@ -97,8 +97,8 @@ class DataFrameResultSet:
             return self._rows[index[0]][col_idx]
         raise excs.Error(f'Bad index: {index}')
 
-    def __iter__(self) -> DataFrameResultSetIterator:
-        return DataFrameResultSetIterator(self)
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        return (self._row_to_dict(i) for i in range(len(self)))
 
     def __eq__(self, other):
         if not isinstance(other, DataFrameResultSet):
@@ -106,19 +106,6 @@ class DataFrameResultSet:
         return self.to_pandas().equals(other.to_pandas())
 
 
-class DataFrameResultSetIterator:
-    def __init__(self, result_set: DataFrameResultSet):
-        self._result_set = result_set
-        self._idx = 0
-
-    def __next__(self) -> Dict[str, Any]:
-        if self._idx >= len(self._result_set):
-            raise StopIteration
-        row = self._result_set._row_to_dict(self._idx)
-        self._idx += 1
-        return row
-
-
 # # TODO: remove this; it's only here as a reminder that we still need to call release() in the current implementation
 # class AnalysisInfo:
 #     def __init__(self, tbl: catalog.TableVersion):
@@ -296,7 +283,7 @@ class DataFrame:
 
     def _create_query_plan(self) -> exec.ExecNode:
         # construct a group-by clause if we're grouping by a table
-        group_by_clause: List[exprs.Expr] = []
+        group_by_clause: Optional[list[exprs.Expr]] = None
         if self.grouping_tbl is not None:
             assert self.group_by_clause is None
             num_rowid_cols = len(self.grouping_tbl.store_tbl.rowid_columns())
@@ -315,8 +302,8 @@ class DataFrame:
             where_clause=self.where_clause,
             group_by_clause=group_by_clause,
             order_by_clause=self.order_by_clause if self.order_by_clause is not None else [],
-            limit=self.limit_val if self.limit_val is not None else 0,
-        )  # limit_val == 0: no limit_val
+            limit=self.limit_val
+        )
 
 
     def show(self, n: int = 20) -> DataFrameResultSet:
@@ -629,17 +616,15 @@ class DataFrame:
         if self.limit_val is not None:
             raise excs.Error(f'Cannot use `{op_name}` after `limit`')
 
-    def __getitem__(self, index: object) -> DataFrame:
+    def __getitem__(self, index: Union[exprs.Expr, Sequence[exprs.Expr]]) -> DataFrame:
         """
         Allowed:
         - [List[Expr]]/[Tuple[Expr]]: setting the select list
         - [Expr]: setting a single-col select list
         """
-        if isinstance(index, tuple):
-            index = list(index)
         if isinstance(index, exprs.Expr):
-            index = [index]
-        if isinstance(index, list):
+            return self.select(index)
+        if isinstance(index, Sequence):
             return self.select(*index)
         raise TypeError(f'Invalid index type: {type(index)}')
 

pixeltable/env.py

@@ -8,6 +8,7 @@ import importlib.util
 import inspect
 import logging
 import os
+import shutil
 import subprocess
 import sys
 import threading
@@ -15,12 +16,12 @@ import uuid
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar
 from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 
 import pixeltable_pgserver
 import sqlalchemy as sql
-import yaml
+import toml
 from tqdm import TqdmWarning
 
 import pixeltable.exceptions as excs
@@ -64,7 +65,7 @@ class Env:
     _log_to_stdout: bool
     _module_log_level: dict[str, int]  # module name -> log level
     _config_file: Optional[Path]
-    _config: Optional[dict[str, Any]]
+    _config: Optional[Config]
     _stdout_handler: logging.StreamHandler
     _initialized: bool
 
@@ -110,6 +111,7 @@ class Env:
         self._log_to_stdout = False
         self._module_log_level = {}  # module name -> log level
 
+        # config
         self._config_file = None
         self._config = None
 
@@ -119,7 +121,8 @@ class Env:
         self._initialized = False
 
     @property
-    def config(self):
+    def config(self) -> Config:
+        assert self._config is not None
         return self._config
 
     @property
@@ -227,30 +230,13 @@ class Env:
         home = Path(os.environ.get('PIXELTABLE_HOME', str(Path.home() / '.pixeltable')))
         assert self._home is None or self._home == home
         self._home = home
-        self._config_file = Path(os.environ.get('PIXELTABLE_CONFIG', str(self._home / 'config.yaml')))
+        self._config_file = Path(os.environ.get('PIXELTABLE_CONFIG', str(self._home / 'config.toml')))
         self._media_dir = self._home / 'media'
         self._file_cache_dir = self._home / 'file_cache'
         self._dataset_cache_dir = self._home / 'dataset_cache'
         self._log_dir = self._home / 'logs'
         self._tmp_dir = self._home / 'tmp'
 
-        # Read in the config
-        if os.path.isfile(self._config_file):
-            with open(self._config_file, 'r') as stream:
-                try:
-                    self._config = yaml.safe_load(stream)
-                except yaml.YAMLError as exc:
-                    self._logger.error(f'Could not read config file: {self._config_file}')
-                    self._config = {}
-        else:
-            self._config = {}
-
-        # Disable spurious warnings
-        warnings.simplefilter('ignore', category=TqdmWarning)
-        if 'hide_warnings' in self._config and self._config['hide_warnings']:
-            # Disable more warnings
-            warnings.simplefilter('ignore', category=UserWarning)
-
         if self._home.exists() and not self._home.is_dir():
             raise RuntimeError(f'{self._home} is not a directory')
 
@@ -274,6 +260,22 @@ class Env:
         if not self._tmp_dir.exists():
             self._tmp_dir.mkdir()
 
+        # Read in the config
+        self._config = Config.from_file(self._config_file)
+        self._file_cache_size_g = self._config.get_float_value('file_cache_size_g')
+        if self._file_cache_size_g is None:
+            raise excs.Error(
+                'pixeltable/file_cache_size_g is missing from configuration\n'
+                f'(either add a `file_cache_size_g` entry to the `pixeltable` section of {self._config_file},\n'
+                'or set the PIXELTABLE_FILE_CACHE_SIZE_G environment variable)'
+            )
+
+        # Disable spurious warnings
+        warnings.simplefilter('ignore', category=TqdmWarning)
+        if self._config.get_bool_value('hide_warnings'):
+            # Disable more warnings
+            warnings.simplefilter('ignore', category=UserWarning)
+
         # configure _logger to log to a file
         self._logfilename = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + '.log'
         fh = logging.FileHandler(self._log_dir / self._logfilename, mode='w')
@@ -313,7 +315,7 @@ class Env:
         self._db_server = pixeltable_pgserver.get_server(self._pgdata_dir, cleanup_mode=None)
         self._db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
 
-        tz_name = os.environ.get('PXT_TIME_ZONE', self._config.get('pxt_time_zone', None))
+        tz_name = self.config.get_string_value('time_zone')
         if tz_name is not None:
             # Validate tzname
             if not isinstance(tz_name, str):
@@ -340,7 +342,7 @@ class Env:
 
         if create_db:
             from pixeltable.metadata import schema
-            schema.Base.metadata.create_all(self._sa_engine)
+            schema.base_metadata.create_all(self._sa_engine)
             metadata.create_system_info(self._sa_engine)
 
         print(f'Connected to Pixeltable database at: {self.db_url}')
@@ -440,21 +442,18 @@ class Env:
         if cl.client_obj is not None:
             return cl.client_obj  # Already initialized
 
-        # Construct a client. For each client parameter, first check if the parameter is in the environment;
-        # if not, look in Pixeltable config from `config.yaml`.
+        # Construct a client, retrieving each parameter from config.
 
         init_kwargs: dict[str, str] = {}
         for param in cl.param_names:
-            environ = f'{name.upper()}_{param.upper()}'
-            if environ in os.environ:
-                init_kwargs[param] = os.environ[environ]
-            elif name.lower() in self._config and param in self._config[name.lower()]:
-                init_kwargs[param] = self._config[name.lower()][param.lower()]
-            if param not in init_kwargs or init_kwargs[param] == '':
+            arg = self._config.get_string_value(param, section=name)
+            if arg is not None and len(arg) > 0:
+                init_kwargs[param] = arg
+            else:
                 raise excs.Error(
                     f'`{name}` client not initialized: parameter `{param}` is not configured.\n'
-                    f'To fix this, specify the `{environ}` environment variable, or put `{param.lower()}` in '
-                    f'the `{name.lower()}` section of $PIXELTABLE_HOME/config.yaml.'
+                    f'To fix this, specify the `{name.upper()}_{param.upper()}` environment variable, or put `{param.lower()}` in '
+                    f'the `{name.lower()}` section of $PIXELTABLE_HOME/config.toml.'
                 )
 
         cl.client_obj = cl.init_fn(**init_kwargs)
@@ -506,7 +505,6 @@ class Env:
         self.__register_package('spacy')
         self.__register_package('tiktoken')
         self.__register_package('together')
-        self.__register_package('toml')
         self.__register_package('torch')
         self.__register_package('torchvision')
         self.__register_package('transformers')
@@ -643,7 +641,7 @@ def register_client(name: str) -> Callable:
     Pixeltable will attempt to load the client parameters from config. For each
     config parameter:
     - If an environment variable named MY_CLIENT_API_KEY (for example) is set, use it;
-    - Otherwise, look for 'api_key' in the 'my_client' section of config.yaml.
+    - Otherwise, look for 'api_key' in the 'my_client' section of config.toml.
 
     If all config parameters are found, Pixeltable calls the initialization function;
     otherwise it throws an exception.
@@ -660,6 +658,79 @@ def register_client(name: str) -> Callable:
     return decorator
 
 
+class Config:
+    """
+    The (global) Pixeltable configuration, as loaded from `config.toml`. Provides methods for retrieving
+    configuration values, which can be set in the config file or as environment variables.
+    """
+    __config: dict[str, Any]
+
+    T = TypeVar('T')
+
+    @classmethod
+    def from_file(cls, path: Path) -> Config:
+        """
+        Loads configuration from the specified TOML file. If the file does not exist, it will be
+        created and populated with the default configuration.
+        """
+        if os.path.isfile(path):
+            with open(path, 'r') as stream:
+                try:
+                    config_dict = toml.load(stream)
+                except Exception as exc:
+                    raise excs.Error(f'Could not read config file: {str(path)}') from exc
+        else:
+            config_dict = cls.__create_default_config(path)
+            with open(path, 'w') as stream:
+                try:
+                    toml.dump(config_dict, stream)
+                except Exception as exc:
+                    raise excs.Error(f'Could not write config file: {str(path)}') from exc
+            logging.getLogger('pixeltable').info(f'Created default config file at: {str(path)}')
+        return cls(config_dict)
+
+    @classmethod
+    def __create_default_config(cls, config_path: Path) -> dict[str, Any]:
+        free_disk_space_bytes = shutil.disk_usage(config_path.parent).free
+        # Default cache size is 1/5 of free disk space
+        file_cache_size_g = free_disk_space_bytes / 5 / (1 << 30)
+        return {
+            'pixeltable': {
+                'file_cache_size_g': round(file_cache_size_g, 1),
+                'hide_warnings': False,
+            }
+        }
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        self.__config = config
+
+    def get_value(self, key: str, expected_type: type[T], section: str = 'pixeltable') -> Optional[T]:
+        env_var = f'{section.upper()}_{key.upper()}'
+        if env_var in os.environ:
+            value = os.environ[env_var]
+        elif section in self.__config and key in self.__config[section]:
+            value = self.__config[section][key]
+        else:
+            return None
+
+        try:
+            return expected_type(value)  # type: ignore[call-arg]
+        except ValueError:
+            raise excs.Error(f'Invalid value for configuration parameter {section}.{key}: {value}')
+
+    def get_string_value(self, key: str, section: str = 'pixeltable') -> Optional[str]:
+        return self.get_value(key, str, section)
+
+    def get_int_value(self, key: str, section: str = 'pixeltable') -> Optional[int]:
+        return self.get_value(key, int, section)
+
+    def get_float_value(self, key: str, section: str = 'pixeltable') -> Optional[float]:
+        return self.get_value(key, float, section)
+
+    def get_bool_value(self, key: str, section: str = 'pixeltable') -> Optional[bool]:
+        return self.get_value(key, bool, section)
+
+
 _registered_clients: dict[str, ApiClient] = {}
 
 

pixeltable/exceptions.py

@@ -1,6 +1,9 @@
-from typing import List, Any
-from types import TracebackType
 from dataclasses import dataclass
+from types import TracebackType
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from pixeltable import exprs
 
 
 class Error(Exception):
@@ -9,11 +12,11 @@ class Error(Exception):
 
 @dataclass
 class ExprEvalError(Exception):
-    expr: Any  # exprs.Expr, but we're not importing pixeltable.exprs to avoid circular imports
+    expr: 'exprs.Expr'
     expr_msg: str
     exc: Exception
     exc_tb: TracebackType
-    input_vals: List[Any]
+    input_vals: list[Any]
     row_num: int
 
 

pixeltable/exec/__init__.py

@@ -8,4 +8,4 @@ from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
 from .media_validation_node import MediaValidationNode
 from .row_update_node import RowUpdateNode
-from .sql_node import SqlLookupNode, SqlScanNode
+from .sql_node import SqlLookupNode, SqlScanNode, SqlAggregationNode, SqlNode

pixeltable/exec/aggregation_node.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import Iterable, List, Optional, Any
+from typing import Iterable, Optional, Any, Iterator
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -13,17 +13,29 @@ from .exec_node import ExecNode
 _logger = logging.getLogger('pixeltable')
 
 class AggregationNode(ExecNode):
+    """
+    In-memory aggregation for UDAs.
+
+    At the moment, this returns all results in a single DataRowBatch.
+    """
+    group_by: Optional[list[exprs.Expr]]
+    input_exprs: list[exprs.Expr]
+    agg_fn_eval_ctx: exprs.RowBuilder.EvalCtx
+    agg_fn_calls: list[exprs.FunctionCall]
+    output_batch: DataRowBatch
+
     def __init__(
-            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: List[exprs.Expr],
-            agg_fn_calls: List[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
+            self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: Optional[list[exprs.Expr]],
+            agg_fn_calls: list[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
     ):
         super().__init__(row_builder, group_by + agg_fn_calls, input_exprs, input)
         self.input = input
         self.group_by = group_by
         self.input_exprs = list(input_exprs)
-        self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=input_exprs)
+        self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=self.input_exprs)
         # we need to make sure to refer to the same exprs that RowBuilder.eval() will use
         self.agg_fn_calls = self.agg_fn_eval_ctx.target_exprs
+        # create output_batch here, rather than in __iter__(), so we don't need to remember tbl and row_builder
         self.output_batch = DataRowBatch(tbl, row_builder, 0)
 
     def _reset_agg_state(self, row_num: int) -> None:
@@ -45,17 +57,14 @@ class AggregationNode(ExecNode):
                 input_vals = [row[d.slot_idx] for d in fn_call.dependencies()]
                 raise excs.ExprEvalError(fn_call, expr_msg, e, exc_tb, input_vals, row_num)
 
-    def __next__(self) -> DataRowBatch:
-        if self.output_batch is None:
-            raise StopIteration
-
+    def __iter__(self) -> Iterator[DataRowBatch]:
         prev_row: Optional[exprs.DataRow] = None
-        current_group: Optional[List[Any]] = None  # the values of the group-by exprs
+        current_group: Optional[list[Any]] = None  # the values of the group-by exprs
         num_input_rows = 0
         for row_batch in self.input:
             num_input_rows += len(row_batch)
             for row in row_batch:
-                group = [row[e.slot_idx] for e in self.group_by]
+                group = [row[e.slot_idx] for e in self.group_by] if self.group_by is not None else None
                 if current_group is None:
                     current_group = group
                     self._reset_agg_state(0)
@@ -71,9 +80,7 @@ class AggregationNode(ExecNode):
         self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
         self.output_batch.add_row(prev_row)
 
-        result = self.output_batch
-        result.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
-        self.output_batch = None
-        _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(result.rows)} rows')
-        return result
+        self.output_batch.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(self.output_batch.rows)} rows')
+        yield self.output_batch