pixeltable 0.3.6__py3-none-any.whl → 0.3.8__py3-none-any.whl

pixeltable/dataframe.py

@@ -10,13 +10,10 @@ import traceback
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Hashable, Iterator, NoReturn, Optional, Sequence, Union
 
-import numpy as np
 import pandas as pd
 import sqlalchemy as sql
 
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
-from pixeltable import catalog, exec, exprs, plan
+from pixeltable import catalog, exceptions as excs, exec, exprs, plan, type_system as ts
 from pixeltable.catalog import is_valid_identifier
 from pixeltable.catalog.globals import UpdateStatus
 from pixeltable.env import Env
@@ -80,7 +77,7 @@ class DataFrameResultSet:
         if isinstance(index, int):
             return self._row_to_dict(index)
         if isinstance(index, tuple) and len(index) == 2:
-            if not isinstance(index[0], int) or not (isinstance(index[1], str) or isinstance(index[1], int)):
+            if not isinstance(index[0], int) or not isinstance(index[1], (str, int)):
                 raise excs.Error(f'Bad index, expected [<row idx>, <column name | column index>]: {index}')
             if isinstance(index[1], str) and index[1] not in self._col_names:
                 raise excs.Error(f'Invalid column name: {index[1]}')
@@ -96,6 +93,9 @@ class DataFrameResultSet:
             return False
         return self.to_pandas().equals(other.to_pandas())
 
+    def __hash__(self):
+        return hash(self.to_pandas())
+
 
 # # TODO: remove this; it's only here as a reminder that we still need to call release() in the current implementation
 # class AnalysisInfo:
@@ -232,9 +232,8 @@ class DataFrame:
         for var in vars:
             if var.name not in unique_vars:
                 unique_vars[var.name] = var
-            else:
-                if unique_vars[var.name].col_type != var.col_type:
-                    raise excs.Error(f'Multiple definitions of parameter {var.name}')
+            elif unique_vars[var.name].col_type != var.col_type:
+                raise excs.Error(f'Multiple definitions of parameter {var.name}')
         return unique_vars
 
     def parameters(self) -> dict[str, ColumnType]:
@@ -242,17 +241,15 @@ class DataFrame:
 
         Parameters are Variables contained in any component of the DataFrame.
         """
-        vars = self._vars()
-        return {name: var.col_type for name, var in vars.items()}
+        return {name: var.col_type for name, var in self._vars().items()}
 
-    def _exec(self, conn: Optional[sql.engine.Connection] = None) -> Iterator[exprs.DataRow]:
+    def _exec(self) -> Iterator[exprs.DataRow]:
         """Run the query and return rows as a generator.
         This function must not modify the state of the DataFrame, otherwise it breaks dataset caching.
         """
         plan = self._create_query_plan()
 
-        def exec_plan(conn: sql.engine.Connection) -> Iterator[exprs.DataRow]:
-            plan.ctx.set_conn(conn)
+        def exec_plan() -> Iterator[exprs.DataRow]:
             plan.open()
             try:
                 for row_batch in plan:
@@ -260,18 +257,13 @@ class DataFrame:
             finally:
                 plan.close()
 
-        if conn is None:
-            with Env.get().engine.begin() as conn:
-                yield from exec_plan(conn)
-        else:
-            yield from exec_plan(conn)
+        yield from exec_plan()
 
-    async def _aexec(self, conn: sql.engine.Connection) -> AsyncIterator[exprs.DataRow]:
+    async def _aexec(self) -> AsyncIterator[exprs.DataRow]:
         """Run the query and return rows as a generator.
         This function must not modify the state of the DataFrame, otherwise it breaks dataset caching.
         """
         plan = self._create_query_plan()
-        plan.ctx.set_conn(conn)
         plan.open()
         try:
             async for row_batch in plan:
@@ -287,7 +279,7 @@ class DataFrame:
             assert self.group_by_clause is None
             num_rowid_cols = len(self.grouping_tbl.store_tbl.rowid_columns())
             # the grouping table must be a base of self.tbl
-            assert num_rowid_cols <= len(self._first_tbl.tbl_version.store_tbl.rowid_columns())
+            assert num_rowid_cols <= len(self._first_tbl.tbl_version.get().store_tbl.rowid_columns())
             group_by_clause = [exprs.RowidRef(self._first_tbl.tbl_version, idx) for idx in range(num_rowid_cols)]
         elif self.group_by_clause is not None:
             group_by_clause = self.group_by_clause
@@ -327,10 +319,10 @@ class DataFrame:
                 if the DataFrame has an order_by clause.
         """
         if self.order_by_clause is not None:
-            raise excs.Error(f'head() cannot be used with order_by()')
+            raise excs.Error('head() cannot be used with order_by()')
         if self._has_joins():
-            raise excs.Error(f'head() not supported for joins')
-        num_rowid_cols = len(self._first_tbl.tbl_version.store_tbl.rowid_columns())
+            raise excs.Error('head() not supported for joins')
+        num_rowid_cols = len(self._first_tbl.tbl_version.get().store_tbl.rowid_columns())
         order_by_clause = [exprs.RowidRef(self._first_tbl.tbl_version, idx) for idx in range(num_rowid_cols)]
         return self.order_by(*order_by_clause, asc=True).limit(n).collect()
 
@@ -350,10 +342,10 @@ class DataFrame:
                 if the DataFrame has an order_by clause.
         """
         if self.order_by_clause is not None:
-            raise excs.Error(f'tail() cannot be used with order_by()')
+            raise excs.Error('tail() cannot be used with order_by()')
         if self._has_joins():
-            raise excs.Error(f'tail() not supported for joins')
-        num_rowid_cols = len(self._first_tbl.tbl_version.store_tbl.rowid_columns())
+            raise excs.Error('tail() not supported for joins')
+        num_rowid_cols = len(self._first_tbl.tbl_version.get().store_tbl.rowid_columns())
         order_by_clause = [exprs.RowidRef(self._first_tbl.tbl_version, idx) for idx in range(num_rowid_cols)]
         result = self.order_by(*order_by_clause, asc=False).limit(n).collect()
         result._reverse()
@@ -418,7 +410,7 @@ class DataFrame:
         )
 
     def _raise_expr_eval_err(self, e: excs.ExprEvalError) -> NoReturn:
-        msg = f'In row {e.row_num} the {e.expr_msg} encountered exception {type(e.exc).__name__}:\n{str(e.exc)}'
+        msg = f'In row {e.row_num} the {e.expr_msg} encountered exception {type(e.exc).__name__}:\n{e.exc}'
         if len(e.input_vals) > 0:
             input_msgs = [
                 f"'{d}' = {d.col_type.print_value(e.input_vals[i])}" for i, d in enumerate(e.expr.dependencies())
@@ -434,29 +426,27 @@ class DataFrame:
             msg += f'\nStack:\n{nl.join(stack_trace[-1:1:-1])}'
         raise excs.Error(msg) from e
 
-    def _output_row_iterator(self, conn: Optional[sql.engine.Connection] = None) -> Iterator[list]:
-        try:
-            for data_row in self._exec(conn):
-                yield [data_row[e.slot_idx] for e in self._select_list_exprs]
-        except excs.ExprEvalError as e:
-            self._raise_expr_eval_err(e)
-        except sql.exc.DBAPIError as e:
-            raise excs.Error(f'Error during SQL execution:\n{e}')
+    def _output_row_iterator(self) -> Iterator[list]:
+        with Env.get().begin_xact():
+            try:
+                for data_row in self._exec():
+                    yield [data_row[e.slot_idx] for e in self._select_list_exprs]
+            except excs.ExprEvalError as e:
+                self._raise_expr_eval_err(e)
+            except sql.exc.DBAPIError as e:
+                raise excs.Error(f'Error during SQL execution:\n{e}') from e
 
     def collect(self) -> DataFrameResultSet:
-        return self._collect()
+        return DataFrameResultSet(list(self._output_row_iterator()), self.schema)
 
-    def _collect(self, conn: Optional[sql.engine.Connection] = None) -> DataFrameResultSet:
-        return DataFrameResultSet(list(self._output_row_iterator(conn)), self.schema)
-
-    async def _acollect(self, conn: sql.engine.Connection) -> DataFrameResultSet:
+    async def _acollect(self) -> DataFrameResultSet:
         try:
-            result = [[row[e.slot_idx] for e in self._select_list_exprs] async for row in self._aexec(conn)]
+            result = [[row[e.slot_idx] for e in self._select_list_exprs] async for row in self._aexec()]
             return DataFrameResultSet(result, self.schema)
         except excs.ExprEvalError as e:
             self._raise_expr_eval_err(e)
         except sql.exc.DBAPIError as e:
-            raise excs.Error(f'Error during SQL execution:\n{e}')
+            raise excs.Error(f'Error during SQL execution:\n{e}') from e
 
     def count(self) -> int:
         """Return the number of rows in the DataFrame.
@@ -467,7 +457,7 @@ class DataFrame:
         from pixeltable.plan import Planner
 
         stmt = Planner.create_count_stmt(self._first_tbl, self.where_clause)
-        with Env.get().engine.connect() as conn:
+        with Env.get().begin_xact() as conn:
             result: int = conn.execute(stmt).scalar_one()
             assert isinstance(result, int)
             return result
@@ -567,7 +557,7 @@ class DataFrame:
 
         """
         if self.select_list is not None:
-            raise excs.Error(f'Select list already specified')
+            raise excs.Error('Select list already specified')
         for name, _ in named_items.items():
             if not isinstance(name, str) or not is_valid_identifier(name):
                 raise excs.Error(f'Invalid name: {name}')
@@ -653,7 +643,7 @@ class DataFrame:
     ) -> exprs.Expr:
         """Verifies user-specified 'on' argument and converts it into a join predicate."""
         col_refs: list[exprs.ColumnRef] = []
-        joined_tbls = self._from_clause.tbls + [other]
+        joined_tbls = [*self._from_clause.tbls, other]
 
         if isinstance(on, exprs.ColumnRef):
             on = [on]
@@ -663,14 +653,13 @@ class DataFrame:
             if not on.col_type.is_bool_type():
                 raise excs.Error(f"'on': boolean expression expected, but got {on.col_type}: {on}")
             return on
-        else:
-            if not isinstance(on, Sequence) or len(on) == 0:
-                raise excs.Error(f"'on': must be a sequence of column references or a boolean expression")
+        elif not isinstance(on, Sequence) or len(on) == 0:
+            raise excs.Error("'on': must be a sequence of column references or a boolean expression")
 
         assert isinstance(on, Sequence)
         for col_ref in on:
             if not isinstance(col_ref, exprs.ColumnRef):
-                raise excs.Error(f"'on': must be a sequence of column references or a boolean expression")
+                raise excs.Error("'on': must be a sequence of column references or a boolean expression")
             if not col_ref.is_bound_by(joined_tbls):
                 raise excs.Error(f"'on': expression cannot be evaluated in the context of the joined tables: {col_ref}")
             col_refs.append(col_ref)
@@ -765,7 +754,7 @@ class DataFrame:
         join_pred: Optional[exprs.Expr]
         if how == 'cross':
             if on is not None:
-                raise excs.Error(f"'on' not allowed for cross join")
+                raise excs.Error("'on' not allowed for cross join")
             join_pred = None
         else:
             if on is None:
@@ -828,20 +817,20 @@ class DataFrame:
             >>> df = book.group_by(t.genre).select(t.genre, total=sum(t.price)).show()
         """
         if self.group_by_clause is not None:
-            raise excs.Error(f'Group-by already specified')
+            raise excs.Error('Group-by already specified')
         grouping_tbl: Optional[catalog.TableVersion] = None
         group_by_clause: Optional[list[exprs.Expr]] = None
         for item in grouping_items:
             if isinstance(item, catalog.Table):
                 if len(grouping_items) > 1:
-                    raise excs.Error(f'group_by(): only one table can be specified')
+                    raise excs.Error('group_by(): only one table can be specified')
                 if len(self._from_clause.tbls) > 1:
-                    raise excs.Error(f'group_by() with Table not supported for joins')
+                    raise excs.Error('group_by() with Table not supported for joins')
                 # we need to make sure that the grouping table is a base of self.tbl
                 base = self._first_tbl.find_tbl_version(item._tbl_version_path.tbl_id())
                 if base is None or base.id == self._first_tbl.tbl_id():
                     raise excs.Error(f'group_by(): {item._name} is not a base table of {self._first_tbl.tbl_name()}')
-                grouping_tbl = item._tbl_version_path.tbl_version
+                grouping_tbl = item._tbl_version_path.tbl_version.get()
                 break
             if not isinstance(item, exprs.Expr):
                 raise excs.Error(f'Invalid expression in group_by(): {item}')
@@ -943,16 +932,19 @@ class DataFrame:
 
             >>> person = t.select()
 
-            Via the above DataFrame person, update the column 'city' to 'Oakland' and 'state' to 'CA' in the table t:
+            Via the above DataFrame person, update the column 'city' to 'Oakland'
+            and 'state' to 'CA' in the table t:
 
             >>> df = person.update({'city': 'Oakland', 'state': 'CA'})
 
-            Via the above DataFrame person, update the column 'age' to 30 for any rows where 'year' is 2014 in the table t:
+            Via the above DataFrame person, update the column 'age' to 30 for any
+            rows where 'year' is 2014 in the table t:
 
             >>> df = person.where(t.year == 2014).update({'age': 30})
         """
         self._validate_mutable('update', False)
-        return self._first_tbl.tbl_version.update(value_spec, where=self.where_clause, cascade=cascade)
+        with Env.get().begin_xact():
+            return self._first_tbl.tbl_version.get().update(value_spec, where=self.where_clause, cascade=cascade)
 
     def delete(self) -> UpdateStatus:
         """Delete rows form the underlying table of the DataFrame.
@@ -973,8 +965,9 @@ class DataFrame:
         """
         self._validate_mutable('delete', False)
         if not self._first_tbl.is_insertable():
-            raise excs.Error(f'Cannot delete from view')
-        return self._first_tbl.tbl_version.delete(where=self.where_clause)
+            raise excs.Error('Cannot delete from view')
+        with Env.get().begin_xact():
+            return self._first_tbl.tbl_version.get().delete(where=self.where_clause)
 
     def _validate_mutable(self, op_name: str, allow_select: bool) -> None:
         """Tests whether this DataFrame can be mutated (such as by an update operation).
@@ -1020,32 +1013,37 @@ class DataFrame:
 
     @classmethod
     def from_dict(cls, d: dict[str, Any]) -> 'DataFrame':
-        tbls = [catalog.TableVersionPath.from_dict(tbl_dict) for tbl_dict in d['from_clause']['tbls']]
-        join_clauses = [plan.JoinClause(**clause_dict) for clause_dict in d['from_clause']['join_clauses']]
-        from_clause = plan.FromClause(tbls=tbls, join_clauses=join_clauses)
-        select_list = (
-            [(exprs.Expr.from_dict(e), name) for e, name in d['select_list']] if d['select_list'] is not None else None
-        )
-        where_clause = exprs.Expr.from_dict(d['where_clause']) if d['where_clause'] is not None else None
-        group_by_clause = (
-            [exprs.Expr.from_dict(e) for e in d['group_by_clause']] if d['group_by_clause'] is not None else None
-        )
-        grouping_tbl = catalog.TableVersion.from_dict(d['grouping_tbl']) if d['grouping_tbl'] is not None else None
-        order_by_clause = (
-            [(exprs.Expr.from_dict(e), asc) for e, asc in d['order_by_clause']]
-            if d['order_by_clause'] is not None
-            else None
-        )
-        limit_val = exprs.Expr.from_dict(d['limit_val']) if d['limit_val'] is not None else None
-        return DataFrame(
-            from_clause=from_clause,
-            select_list=select_list,
-            where_clause=where_clause,
-            group_by_clause=group_by_clause,
-            grouping_tbl=grouping_tbl,
-            order_by_clause=order_by_clause,
-            limit=limit_val,
-        )
+        # we need to wrap the construction with a transaction, because it might need to load metadata
+        with Env.get().begin_xact():
+            tbls = [catalog.TableVersionPath.from_dict(tbl_dict) for tbl_dict in d['from_clause']['tbls']]
+            join_clauses = [plan.JoinClause(**clause_dict) for clause_dict in d['from_clause']['join_clauses']]
+            from_clause = plan.FromClause(tbls=tbls, join_clauses=join_clauses)
+            select_list = (
+                [(exprs.Expr.from_dict(e), name) for e, name in d['select_list']]
+                if d['select_list'] is not None
+                else None
+            )
+            where_clause = exprs.Expr.from_dict(d['where_clause']) if d['where_clause'] is not None else None
+            group_by_clause = (
+                [exprs.Expr.from_dict(e) for e in d['group_by_clause']] if d['group_by_clause'] is not None else None
+            )
+            grouping_tbl = catalog.TableVersion.from_dict(d['grouping_tbl']) if d['grouping_tbl'] is not None else None
+            order_by_clause = (
+                [(exprs.Expr.from_dict(e), asc) for e, asc in d['order_by_clause']]
+                if d['order_by_clause'] is not None
+                else None
+            )
+            limit_val = exprs.Expr.from_dict(d['limit_val']) if d['limit_val'] is not None else None
+
+            return DataFrame(
+                from_clause=from_clause,
+                select_list=select_list,
+                where_clause=where_clause,
+                group_by_clause=group_by_clause,
+                grouping_tbl=grouping_tbl,
+                order_by_clause=order_by_clause,
+                limit=limit_val,
+            )
 
     def _hash_result_set(self) -> str:
         """Return a hash that changes when the result set changes."""
@@ -1053,7 +1051,7 @@ class DataFrame:
         # add list of referenced table versions (the actual versions, not the effective ones) in order to force cache
         # invalidation when any of the referenced tables changes
         d['tbl_versions'] = [
-            tbl_version.version for tbl in self._from_clause.tbls for tbl_version in tbl.get_tbl_versions()
+            tbl_version.get().version for tbl in self._from_clause.tbls for tbl_version in tbl.get_tbl_versions()
         ]
         summary_string = json.dumps(d)
         return hashlib.sha256(summary_string.encode()).hexdigest()
@@ -1086,7 +1084,8 @@ class DataFrame:
             assert data_file_path.is_file()
             return data_file_path
         else:
-            return write_coco_dataset(self, dest_path)
+            with Env.get().begin_xact():
+                return write_coco_dataset(self, dest_path)
 
     def to_pytorch_dataset(self, image_format: str = 'pt') -> 'torch.utils.data.IterableDataset':
         """
@@ -1130,6 +1129,7 @@ class DataFrame:
         if dest_path.exists():  # fast path: use cache
             assert dest_path.is_dir()
         else:
-            export_parquet(self, dest_path, inline_images=True)
+            with Env.get().begin_xact():
+                export_parquet(self, dest_path, inline_images=True)
 
         return PixeltablePytorchDataset(path=dest_path, image_format=image_format)