pixeltable 0.4.0rc1__py3-none-any.whl → 0.4.0rc2__py3-none-any.whl

pixeltable/exec/sql_node.py

@@ -14,6 +14,7 @@ from .exec_node import ExecNode
 
 if TYPE_CHECKING:
     import pixeltable.plan
+    from pixeltable.plan import SampleClause
 
 _logger = logging.getLogger('pixeltable')
 
@@ -64,8 +65,12 @@ def print_order_by_clause(clause: OrderByClause) -> str:
 
 class SqlNode(ExecNode):
     """
-    Materializes data from the store via a Select stmt.
+    Materializes data from the store via an SQL statement.
     This only provides the select list. The subclasses are responsible for the From clause and any additional clauses.
+    The pk columns are not included in the select list.
+    If set_pk is True, they are added to the end of the result set when creating the SQL statement
+    so they can always be referenced as cols[-num_pk_cols:] in the result set.
+    The pk_columns consist of the rowid columns of the target table followed by the version number.
     """
 
     tbl: Optional[catalog.TableVersionPath]
@@ -122,6 +127,7 @@ class SqlNode(ExecNode):
             # we also need to retrieve the pk columns
             assert tbl is not None
             self.num_pk_cols = len(tbl.tbl_version.get().store_tbl.pk_columns())
+            assert self.num_pk_cols > 1
 
         # additional state
         self.result_cursor = None
@@ -139,15 +145,20 @@ class SqlNode(ExecNode):
             if tv is not None:
                 assert tv.is_validated
 
+    def _create_pk_cols(self) -> list[sql.Column]:
+        """Create a list of pk columns"""
+        # we need to retrieve the pk columns
+        if self.set_pk:
+            assert self.tbl is not None
+            assert self.tbl.tbl_version.get().is_validated
+            return self.tbl.tbl_version.get().store_tbl.pk_columns()
+        return []
+
     def _create_stmt(self) -> sql.Select:
         """Create Select from local state"""
 
         assert self.sql_elements.contains_all(self.select_list)
-        sql_select_list = [self.sql_elements.get(e) for e in self.select_list]
-        if self.set_pk:
-            assert self.tbl is not None
-            assert self.tbl.tbl_version.get().is_validated
-            sql_select_list += self.tbl.tbl_version.get().store_tbl.pk_columns()
+        sql_select_list = [self.sql_elements.get(e) for e in self.select_list] + self._create_pk_cols()
         stmt = sql.select(*sql_select_list)
 
         where_clause_element = (
@@ -173,9 +184,10 @@ class SqlNode(ExecNode):
     def _ordering_tbl_ids(self) -> set[UUID]:
         return exprs.Expr.all_tbl_ids(e for e, _ in self.order_by_clause)
 
-    def to_cte(self) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
+    def to_cte(self, keep_pk: bool = False) -> Optional[tuple[sql.CTE, exprs.ExprDict[sql.ColumnElement]]]:
         """
-        Returns a CTE that materializes the output of this node plus a mapping from select list expr to output column
+        Creates a CTE that materializes the output of this node plus a mapping from select list expr to output column.
+        keep_pk: if True, the PK columns are included in the CTE Select statement
 
         Returns:
             (CTE, dict from Expr to output column)
@@ -183,11 +195,13 @@ class SqlNode(ExecNode):
         if self.py_filter is not None:
             # the filter needs to run in Python
             return None
-        self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
         if self.cte is None:
+            if not keep_pk:
+                self.set_pk = False  # we don't need the PK if we use this SqlNode as a CTE
             self.cte = self._create_stmt().cte()
-            assert len(self.cte.c) == len(self.select_list)
-        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))
+        pk_count = self.num_pk_cols if self.set_pk else 0
+        assert len(self.select_list) + pk_count == len(self.cte.c)
+        return self.cte, exprs.ExprDict(zip(self.select_list, self.cte.c))  # skip pk cols
 
     @classmethod
     def retarget_rowid_refs(cls, target: catalog.TableVersionPath, expr_seq: Iterable[exprs.Expr]) -> None:
@@ -293,7 +307,9 @@ class SqlNode(ExecNode):
                 stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
                 _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
             except Exception:
-                pass
+                # log something if we can't log the compiled stmt
+                stmt_str = repr(stmt)
+                _logger.debug(f'SqlLookupNode proto-stmt:\n{stmt_str}')
             self._log_explain(stmt)
 
             conn = Env.get().conn
@@ -510,3 +526,127 @@ class SqlJoinNode(SqlNode):
                 full=join_clause == plan.JoinType.FULL_OUTER,
             )
         return stmt
+
+
+class SqlSampleNode(SqlNode):
+    """
+    Returns rows from a stratified sample with N samples per strata.
+    """
+
+    stratify_exprs: Optional[list[exprs.Expr]]
+    n_samples: Optional[int]
+    fraction_samples: Optional[float]
+    seed: int
+    input_cte: Optional[sql.CTE]
+    pk_count: int
+
+    def __init__(
+        self,
+        row_builder: exprs.RowBuilder,
+        input: SqlNode,
+        select_list: Iterable[exprs.Expr],
+        stratify_exprs: Optional[list[exprs.Expr]] = None,
+        sample_clause: Optional['SampleClause'] = None,
+    ):
+        """
+        Args:
+            select_list: can contain calls to AggregateFunctions
+            stratify_exprs: list of expressions to group by
+            n: number of samples per strata
+        """
+        self.input_cte, input_col_map = input.to_cte(keep_pk=True)
+        self.pk_count = input.num_pk_cols
+        assert self.pk_count > 1
+        sql_elements = exprs.SqlElementCache(input_col_map)
+        super().__init__(input.tbl, row_builder, select_list, sql_elements, set_pk=True)
+        self.stratify_exprs = stratify_exprs
+        self.n_samples = sample_clause.n
+        self.n_per_stratum = sample_clause.n_per_stratum
+        self.fraction_samples = sample_clause.fraction
+        self.seed = sample_clause.seed if sample_clause.seed is not None else 0
+
+    @classmethod
+    def key_sql_expr(cls, seed: sql.ColumnElement, sql_cols: Iterable[sql.ColumnElement]) -> sql.ColumnElement:
+        """Construct expression which is the ordering key for rows to be sampled
+        General SQL form is:
+        - MD5(<seed::text> [ + '___' + <rowid_col_val>::text]+
+        """
+        sql_expr: sql.ColumnElement = sql.cast(seed, sql.Text)
+        for e in sql_cols:
+            sql_expr = sql_expr + sql.literal_column("'___'") + sql.cast(e, sql.Text)
+        sql_expr = sql.func.md5(sql_expr)
+        return sql_expr
+
+    def _create_order_by(self, cte: sql.CTE) -> sql.ColumnElement:
+        """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.seed)), rowid_cols)
+
+    def _create_stmt(self) -> sql.Select:
+        if self.fraction_samples is not None:
+            return self._create_stmt_fraction(self.fraction_samples)
+        return self._create_stmt_n(self.n_samples, self.n_per_stratum)
+
+    def _create_stmt_n(self, n: Optional[int], n_per_stratum: Optional[int]) -> sql.Select:
+        """Create a Select stmt that returns n samples across all strata"""
+        sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
+        order_by = self._create_order_by(self.input_cte)
+
+        # Create a list of all columns plus the rank
+        # Get all columns from the input CTE dynamically
+        select_columns = [*self.input_cte.c]
+        select_columns.append(
+            sql.func.row_number().over(partition_by=sql_strata_exprs, order_by=order_by).label('rank')
+        )
+        row_rank_cte = sql.select(*select_columns).select_from(self.input_cte).cte('row_rank_cte')
+
+        final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
+        if n_per_stratum is not None:
+            return sql.select(*final_columns).filter(row_rank_cte.c.rank <= n_per_stratum)
+        else:
+            secondary_order = self._create_order_by(row_rank_cte)
+            return sql.select(*final_columns).order_by(row_rank_cte.c.rank, secondary_order).limit(n)
+
+    def _create_stmt_fraction(self, fraction_samples: float) -> sql.Select:
+        """Create a Select stmt that returns a fraction of the rows per strata"""
+
+        # Build the strata count CTE
+        # Produces a table of the form:
+        #   ([stratify_exprs], s_s_size)
+        # where s_s_size is the number of samples to take from each stratum
+        sql_strata_exprs = [self.sql_elements.get(e) for e in self.stratify_exprs]
+        per_strata_count_cte = (
+            sql.select(
+                *sql_strata_exprs,
+                sql.func.ceil(fraction_samples * sql.func.count(1).cast(sql.Integer)).label('s_s_size'),
+            )
+            .select_from(self.input_cte)
+            .group_by(*sql_strata_exprs)
+            .cte('per_strata_count_cte')
+        )
+
+        # Build a CTE that ranks the rows within each stratum
+        # Include all columns from the input CTE dynamically
+        order_by = self._create_order_by(self.input_cte)
+        select_columns = [*self.input_cte.c]
+        select_columns.append(
+            sql.func.row_number().over(partition_by=sql_strata_exprs, order_by=order_by).label('rank')
+        )
+        row_rank_cte = sql.select(*select_columns).select_from(self.input_cte).cte('row_rank_cte')
+
+        # Build the join criterion dynamically to accommodate any number of group by columns
+        join_c = sql.true()
+        for col in per_strata_count_cte.c[:-1]:
+            join_c &= row_rank_cte.c[col.name].isnot_distinct_from(col)
+
+        # Join srcp with per_strata_count_cte to limit returns to the requested fraction of rows
+        final_columns = [*row_rank_cte.c[:-1]]  # exclude the rank column
+        stmt = (
+            sql.select(*final_columns)
+            .select_from(row_rank_cte)
+            .join(per_strata_count_cte, join_c)
+            .where(row_rank_cte.c.rank <= per_strata_count_cte.c.s_s_size)
+        )
+
+        return stmt

pixeltable/exprs/data_row.py

@@ -214,6 +214,7 @@ class DataRow:
         """Assign in-memory cell value
         This allows overwriting
         """
+        assert isinstance(idx, int)
         assert self.excs[idx] is None
 
         if (idx in self.img_slot_idxs or idx in self.media_slot_idxs) and isinstance(val, str):
@@ -253,14 +254,15 @@ class DataRow:
         assert self.excs[index] is None
         if self.file_paths[index] is None:
             if filepath is not None:
-                # we want to save this to a file
-                self.file_paths[index] = filepath
-                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
                 image = self.vals[index]
                 assert isinstance(image, PIL.Image.Image)
                 # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
                 # In that case, use WebP instead.
                 format = 'webp' if image.has_transparency_data else 'jpeg'
+                if not filepath.endswith(f'.{format}'):
+                    filepath += f'.{format}'
+                self.file_paths[index] = filepath
+                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
                 image.save(filepath, format=format)
             else:
                 # we discard the content of this cell

pixeltable/exprs/expr.py

@@ -276,6 +276,13 @@ class Expr(abc.ABC):
         tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
         return self._retarget(tbl_versions)
 
+    @classmethod
+    def retarget_list(cls, expr_list: list[Expr], tbl: catalog.TableVersionPath) -> None:
+        """Retarget ColumnRefs in expr_list to the specific TableVersions in tbl."""
+        tbl_versions = {tbl_version.id: tbl_version.get() for tbl_version in tbl.get_tbl_versions()}
+        for i in range(len(expr_list)):
+            expr_list[i] = expr_list[i]._retarget(tbl_versions)
+
     def _retarget(self, tbl_versions: dict[UUID, catalog.TableVersion]) -> Self:
         for i in range(len(self.components)):
             self.components[i] = self.components[i]._retarget(tbl_versions)

pixeltable/exprs/literal.py

@@ -16,6 +16,8 @@ from .sql_element_cache import SqlElementCache
 
 
 class Literal(Expr):
+    val: Any
+
     def __init__(self, val: Any, col_type: Optional[ts.ColumnType] = None):
         if col_type is not None:
             val = col_type.create_literal(val)

pixeltable/func/tools.py

@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
 
 import pydantic
 
-import pixeltable.exceptions as excs
+from pixeltable import exceptions as excs
 
 from .function import Function
 from .signature import Parameter

pixeltable/functions/gemini.py

@@ -103,7 +103,6 @@ def invoke_tools(tools: pxt.func.Tools, response: exprs.Expr) -> exprs.InlineDic
 
 @pxt.udf
 def _gemini_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
-    print(response)
     pxt_tool_calls: dict[str, list[dict]] = {}
     for part in response['candidates'][0]['content']['parts']:
         tool_call = part.get('function_call')

pixeltable/globals.py

@@ -249,13 +249,17 @@ def create_view(
     where: Optional[exprs.Expr] = None
     if isinstance(base, catalog.Table):
         tbl_version_path = base._tbl_version_path
+        sample_clause = None
     elif isinstance(base, DataFrame):
         base._validate_mutable('create_view', allow_select=True)
         if len(base._from_clause.tbls) > 1:
             raise excs.Error('Cannot create a view of a join')
         tbl_version_path = base._from_clause.tbls[0]
         where = base.where_clause
+        sample_clause = base.sample_clause
         select_list = base.select_list
+        if sample_clause is not None and not is_snapshot and not sample_clause.is_repeatable:
+            raise excs.Error('Non-snapshot views cannot be created with non-fractional or stratified sampling')
     else:
         raise excs.Error('`base` must be an instance of `Table` or `DataFrame`')
     assert isinstance(base, (catalog.Table, DataFrame))
@@ -280,6 +284,7 @@ def create_view(
         tbl_version_path,
         select_list=select_list,
         where=where,
+        sample_clause=sample_clause,
         additional_columns=additional_columns,
         is_snapshot=is_snapshot,
         iterator=iterator,

pixeltable/metadata/__init__.py

@@ -8,15 +8,17 @@ from typing import Callable
 import sqlalchemy as sql
 from sqlalchemy import orm
 
+import pixeltable as pxt
+import pixeltable.exceptions as excs
 from pixeltable.utils.console_output import ConsoleLogger
 
 from .schema import SystemInfo, SystemInfoMd
 
 _console_logger = ConsoleLogger(logging.getLogger('pixeltable'))
-
+_logger = logging.getLogger('pixeltable')
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 36
+VERSION = 37
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:
@@ -55,6 +57,13 @@ def upgrade_md(engine: sql.engine.Engine) -> None:
         system_info = session.query(SystemInfo).one().md
         md_version = system_info['schema_version']
         assert isinstance(md_version, int)
+        _logger.info(f'Current database version: {md_version}, installed version: {VERSION}')
+        if md_version > VERSION:
+            raise excs.Error(
+                'This Pixeltable database was created with a newer Pixeltable version '
+                f'than the one currently installed ({pxt.__version__}).\n'
+                'Please update to the latest Pixeltable version by running: pip install --upgrade pixeltable'
+            )
         if md_version == VERSION:
             return
         while md_version < VERSION:

pixeltable/metadata/converters/convert_36.py

@@ -0,0 +1,38 @@
+import logging
+from typing import Any, Optional
+from uuid import UUID
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+_logger = logging.getLogger('pixeltable')
+
+
+@register_converter(version=36)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, table_md_updater=__update_table_md, substitution_fn=__substitute_md)
+
+
+def __update_table_md(table_md: dict, table_id: UUID) -> None:
+    """Update the view metadata to add the sample_clause field if it is missing
+
+    Args:
+        table_md (dict): copy of the original table metadata. this gets updated in place.
+        table_id (UUID): the table id
+
+    """
+    if table_md['view_md'] is None:
+        return
+    if 'sample_clause' not in table_md['view_md']:
+        table_md['view_md']['sample_clause'] = None
+        _logger.info(f'Updating view metadata for table: {table_id}')
+
+
+def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
+    if isinstance(v, dict) and (v.get('_classname') == 'DataFrame'):
+        if 'sample_clause' not in v:
+            v['sample_clause'] = None
+        return k, v
+    return None

pixeltable/metadata/notes.py

@@ -2,6 +2,7 @@
 # rather than as a comment, so that the existence of a description can be enforced by
 # the unit tests when new versions are added.
 VERSION_NOTES = {
+    37: 'Add support for the sample() method on DataFrames',
     36: 'Added Table.lock_dummy',
     35: 'Track reference_tbl in ColumnRef',
     34: 'Set default value for is_pk field in column metadata to False',

pixeltable/metadata/schema.py

@@ -147,6 +147,9 @@ class ViewMd:
     # filter predicate applied to the base table; view-only
     predicate: Optional[dict[str, Any]]
 
+    # sampling predicate applied to the base table; view-only
+    sample_clause: Optional[dict[str, Any]]
+
     # ComponentIterator subclass; only for component views
     iterator_class_fqn: Optional[str]