pixeltable 0.4.16__py3-none-any.whl → 0.4.18__py3-none-any.whl

pixeltable/exec/exec_context.py

@@ -1,3 +1,4 @@
+import random
 from typing import Optional
 
 import sqlalchemy as sql
@@ -8,6 +9,17 @@ from pixeltable import exprs
 class ExecContext:
     """Class for execution runtime constants"""
 
+    row_builder: exprs.RowBuilder
+    profile: exprs.ExecProfile
+    show_pbar: bool
+    batch_size: int
+    num_rows: Optional[int]
+    conn: Optional[sql.engine.Connection]
+    pk_clause: Optional[list[sql.ClauseElement]]
+    num_computed_exprs: int
+    ignore_errors: bool
+    random_seed: int  # general-purpose source of randomness with execution scope
+
     def __init__(
         self,
         row_builder: exprs.RowBuilder,
@@ -23,8 +35,9 @@ class ExecContext:
         self.row_builder = row_builder
         self.profile = exprs.ExecProfile(row_builder)
         # num_rows is used to compute the total number of computed cells used for the progress bar
-        self.num_rows: Optional[int] = None
-        self.conn: Optional[sql.engine.Connection] = None  # if present, use this to execute SQL queries
+        self.num_rows = None
+        self.conn = None  # if present, use this to execute SQL queries
         self.pk_clause = pk_clause
         self.num_computed_exprs = num_computed_exprs
         self.ignore_errors = ignore_errors
+        self.random_seed = random.randint(0, 1 << 63)

pixeltable/exec/sql_node.py

@@ -648,7 +648,6 @@ class SqlSampleNode(SqlNode):
         )
         self.stratify_exprs = stratify_exprs
         self.sample_clause = sample_clause
-        assert isinstance(self.sample_clause.seed, int)
 
     @classmethod
     def key_sql_expr(cls, seed: sql.ColumnElement, sql_cols: Iterable[sql.ColumnElement]) -> sql.ColumnElement:
@@ -667,7 +666,9 @@ class SqlSampleNode(SqlNode):
         """Create an expression for randomly ordering rows with a given seed"""
         rowid_cols = [*cte.c[-self.pk_count : -1]]  # exclude the version column
         assert len(rowid_cols) > 0
-        return self.key_sql_expr(sql.literal_column(str(self.sample_clause.seed)), rowid_cols)
+        # If seed is not set in the sample clause, use the random seed given by the execution context
+        seed = self.sample_clause.seed if self.sample_clause.seed is not None else self.ctx.random_seed
+        return self.key_sql_expr(sql.literal_column(str(seed)), rowid_cols)
 
     def _create_stmt(self) -> sql.Select:
         from pixeltable.plan import SampleClause

pixeltable/exprs/arithmetic_expr.py

@@ -4,7 +4,7 @@ from typing import Any, Optional
 
 import sqlalchemy as sql
 
-from pixeltable import exceptions as excs, type_system as ts
+from pixeltable import env, exceptions as excs, type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr
@@ -64,12 +64,18 @@ class ArithmeticExpr(Expr):
         right = sql_elements.get(self._op2)
         if left is None or right is None:
             return None
-        if self.operator == ArithmeticOperator.ADD:
-            return left + right
-        if self.operator == ArithmeticOperator.SUB:
-            return left - right
-        if self.operator == ArithmeticOperator.MUL:
-            return left * right
+        if self.operator in (ArithmeticOperator.ADD, ArithmeticOperator.SUB, ArithmeticOperator.MUL):
+            if env.Env.get().is_using_cockroachdb and self._op1.col_type != self._op2.col_type:
+                if self._op1.col_type != self.col_type:
+                    left = sql.cast(left, self.col_type.to_sa_type())
+                if self._op2.col_type != self.col_type:
+                    right = sql.cast(right, self.col_type.to_sa_type())
+            if self.operator == ArithmeticOperator.ADD:
+                return left + right
+            if self.operator == ArithmeticOperator.SUB:
+                return left - right
+            if self.operator == ArithmeticOperator.MUL:
+                return left * right
         if self.operator == ArithmeticOperator.DIV:
             assert self.col_type.is_float_type()
             # Avoid division by zero errors by converting any zero divisor to NULL.