pixeltable 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

pixeltable/plan.py

@@ -40,7 +40,7 @@ class Analyzer:
 
     def __init__(
             self, tbl: catalog.TableVersionPath, select_list: List[exprs.Expr],
-            where_clause: Optional[exprs.Predicate] = None, group_by_clause: Optional[List[exprs.Expr]] = None,
+            where_clause: Optional[exprs.Expr] = None, group_by_clause: Optional[List[exprs.Expr]] = None,
             order_by_clause: Optional[List[Tuple[exprs.Expr, bool]]] = None):
         if group_by_clause is None:
             group_by_clause = []
@@ -58,7 +58,7 @@ class Analyzer:
         # Where clause of the Select stmt of the SQL scan
         self.sql_where_clause: Optional[exprs.Expr] = None
         # filter predicate applied to output rows of the SQL scan
-        self.filter: Optional[exprs.Predicate] = None
+        self.filter: Optional[exprs.Expr] = None
         # not executable
         #self.similarity_clause: Optional[exprs.ImageSimilarityPredicate] = None
         if where_clause is not None:
@@ -107,7 +107,7 @@ class Analyzer:
         for e in self.group_by_clause:
             if e.sql_expr() is None:
                 raise excs.Error(f'Invalid grouping expression, needs to be expressible in SQL: {e}')
-            if e.contains(filter=lambda e: _is_agg_fn_call(e)):
+            if e._contains(filter=lambda e: _is_agg_fn_call(e)):
                 raise excs.Error(f'Grouping expression contains aggregate function: {e}')
 
         # check that agg fn calls don't have contradicting ordering requirements
@@ -183,7 +183,7 @@ class Planner:
     # TODO: create an exec.CountNode and change this to create_count_plan()
     @classmethod
     def create_count_stmt(
-            cls, tbl: catalog.TableVersionPath, where_clause: Optional[exprs.Predicate] = None
+            cls, tbl: catalog.TableVersionPath, where_clause: Optional[exprs.Expr] = None
     ) -> sql.Select:
         stmt = sql.select(sql.func.count('*'))
         refd_tbl_ids: Set[UUID] = set()
@@ -239,7 +239,7 @@ class Planner:
             cls, tbl: catalog.TableVersionPath,
             update_targets: dict[catalog.Column, exprs.Expr],
             recompute_targets: List[catalog.Column],
-            where_clause: Optional[exprs.Predicate], cascade: bool
+            where_clause: Optional[exprs.Expr], cascade: bool
     ) -> Tuple[exec.ExecNode, List[str], List[catalog.Column]]:
         """Creates a plan to materialize updated rows.
         The plan:
@@ -288,6 +288,81 @@ class Planner:
         recomputed_user_cols = [c for c in recomputed_cols if c.name is not None]
         return plan, [f'{c.tbl.name}.{c.name}' for c in updated_cols + recomputed_user_cols], recomputed_user_cols
 
+    @classmethod
+    def create_batch_update_plan(
+            cls, tbl: catalog.TableVersionPath,
+            batch: list[dict[catalog.Column, exprs.Expr]], rowids: list[tuple[int, ...]],
+            cascade: bool
+    ) -> Tuple[exec.ExecNode, exec.RowUpdateNode, sql.ClauseElement, List[catalog.Column], List[catalog.Column]]:
+        """
+        Returns:
+        - root node of the plan to produce the updated rows
+        - RowUpdateNode of plan
+        - Where clause for deleting the current versions of updated rows
+        - list of columns that are getting updated
+        - list of user-visible columns that are being recomputed
+        """
+        assert isinstance(tbl, catalog.TableVersionPath)
+        target = tbl.tbl_version  # the one we need to update
+        sa_key_cols: list[sql.Column] = []
+        key_vals: list[tuple] = []
+        if len(rowids) > 0:
+            sa_key_cols = target.store_tbl.rowid_columns()
+            key_vals = rowids
+        else:
+            pk_cols = target.primary_key_columns()
+            sa_key_cols = [c.sa_col for c in pk_cols]
+            key_vals = [tuple(row[col].val for col in pk_cols) for row in batch]
+
+        # retrieve all stored cols and all target exprs
+        updated_cols = batch[0].keys() - target.primary_key_columns()
+        recomputed_cols = target.get_dependent_columns(updated_cols) if cascade else set()
+        # regardless of cascade, we need to update all indices on any updated column
+        idx_val_cols = target.get_idx_val_columns(updated_cols)
+        recomputed_cols.update(idx_val_cols)
+        # we only need to recompute stored columns (unstored ones are substituted away)
+        recomputed_cols = {c for c in recomputed_cols if c.is_stored}
+        recomputed_base_cols = {col for col in recomputed_cols if col.tbl == target}
+        copied_cols = [
+            col for col in target.cols if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
+        ]
+        select_list = [exprs.ColumnRef(col) for col in copied_cols]
+        select_list.extend([exprs.ColumnRef(col) for col in updated_cols])
+
+        recomputed_exprs = \
+            [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols]
+        # the RowUpdateNode updates columns in-place, ie, in the original ColumnRef; no further sustitution is needed
+        select_list.extend(recomputed_exprs)
+
+        # ExecNode tree (from bottom to top):
+        # - SqlLookupNode to retrieve the existing rows
+        # - RowUpdateNode to update the retrieved rows
+        # - ExprEvalNode to evaluate the remaining output exprs
+        analyzer = Analyzer(tbl, select_list)
+        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], analyzer.sql_exprs)
+        analyzer.finalize(row_builder)
+        plan = exec.SqlLookupNode(tbl, row_builder, analyzer.sql_exprs, sa_key_cols, key_vals)
+        delete_where_clause = plan.where_clause
+        col_vals = [{col: row[col].val for col in updated_cols} for row in batch]
+        plan = row_update_node = exec.RowUpdateNode(tbl, key_vals, len(rowids) > 0, col_vals, row_builder, plan)
+        if not cls._is_contained_in(analyzer.select_list, analyzer.sql_exprs):
+            # we need an ExprEvalNode to evaluate the remaining output exprs
+            plan = exec.ExprEvalNode(row_builder, analyzer.select_list, analyzer.sql_exprs, input=plan)
+        # update row builder with column information
+        all_base_cols = copied_cols + list(updated_cols) + list(recomputed_base_cols)  # same order as select_list
+        row_builder.substitute_exprs(select_list, remove_duplicates=False)
+        for i, col in enumerate(all_base_cols):
+            plan.row_builder.add_table_column(col, select_list[i].slot_idx)
+
+        ctx = exec.ExecContext(row_builder)
+        # we're returning everything to the user, so we might as well do it in a single batch
+        ctx.batch_size = 0
+        plan.set_ctx(ctx)
+        recomputed_user_cols = [c for c in recomputed_cols if c.name is not None]
+        return (
+            plan, row_update_node, delete_where_clause, list(updated_cols) + recomputed_user_cols, recomputed_user_cols
+        )
+
     @classmethod
     def create_view_update_plan(
             cls, view: catalog.TableVersionPath, recompute_targets: List[catalog.Column]
@@ -505,7 +580,7 @@ class Planner:
     @classmethod
     def create_query_plan(
             cls, tbl: catalog.TableVersionPath, select_list: Optional[List[exprs.Expr]] = None,
-            where_clause: Optional[exprs.Predicate] = None, group_by_clause: Optional[List[exprs.Expr]] = None,
+            where_clause: Optional[exprs.Expr] = None, group_by_clause: Optional[List[exprs.Expr]] = None,
             order_by_clause: Optional[List[Tuple[exprs.Expr, bool]]] = None, limit: Optional[int] = None,
             with_pk: bool = False, ignore_errors: bool = False, exact_version_only: Optional[List[catalog.TableVersion]] = None
     ) -> exec.ExecNode:
@@ -597,7 +672,7 @@ class Planner:
         return plan
 
     @classmethod
-    def analyze(cls, tbl: catalog.TableVersionPath, where_clause: exprs.Predicate) -> Analyzer:
+    def analyze(cls, tbl: catalog.TableVersionPath, where_clause: exprs.Expr) -> Analyzer:
         return Analyzer(tbl, [], where_clause=where_clause)
 
     @classmethod

pixeltable/tool/create_test_db_dump.py

@@ -61,7 +61,7 @@ class Dumper:
         info_dict = {'pixeltable-dump': {
             'metadata-version': md_version,
             'git-sha': git_sha,
-            'datetime': datetime.datetime.utcnow(),
+            'datetime': datetime.datetime.now(tz=datetime.timezone.utc),
             'user': user
         }}
         with open(info_file, 'w') as info:
@@ -179,7 +179,7 @@ class Dumper:
     def __add_expr_columns(self, t: pxt.Table, col_prefix: str, include_expensive_functions=False) -> None:
         def add_column(col_name: str, col_expr: Any) -> None:
             t.add_column(**{f'{col_prefix}_{col_name}': col_expr})
-        
+
         # arithmetic_expr
         add_column('plus', t.c2 + 6)
         add_column('minus', t.c2 - 5)
@@ -208,7 +208,7 @@ class Dumper:
         add_column('not', ~(t.c2 > 20))
 
         # function_call
-        add_column('function_call', pxt.functions.string.str_format('{0} {key}', t.c1, key=t.c1))  # library function
+        add_column('function_call', pxt.functions.string.format('{0} {key}', t.c1, key=t.c1))  # library function
         add_column('test_udf', test_udf_stored(t.c2))  # stored udf
         add_column('test_udf_batched', test_udf_stored_batched(t.c1, upper=False))  # batched stored udf
         if include_expensive_functions:
@@ -242,8 +242,7 @@ class Dumper:
         add_column('str_const', 'str')
         add_column('int_const', 5)
         add_column('float_const', 5.0)
-        add_column('timestamp_const_1', datetime.datetime.utcnow())
-        add_column('timestamp_const_2', datetime.date.today())
+        add_column('timestamp_const_1', datetime.datetime.now(tz=datetime.timezone.utc))
 
         # type_cast
         add_column('astype', t.c2.astype(FloatType()))

pixeltable/tool/doc_plugins/griffe.py

@@ -0,0 +1,81 @@
+import ast
+from typing import Optional, Union
+import warnings
+
+import griffe
+import griffe.expressions
+from griffe import Extension, Object, ObjectNode
+
+import pixeltable as pxt
+
+logger = griffe.get_logger(__name__)
+
+class PxtGriffeExtension(Extension):
+    """Implementation of a Pixeltable custom griffe extension."""
+
+    def on_instance(self, node: Union[ast.AST, ObjectNode], obj: Object) -> None:
+        if obj.docstring is None:
+            # Skip over entities without a docstring
+            return
+
+        if isinstance(obj, griffe.Function):
+            # See if the (Python) function has a @pxt.udf decorator
+            if any(
+                isinstance(dec.value, griffe.expressions.Expr) and dec.value.canonical_path in ['pixeltable.func.udf', 'pixeltable.udf']
+                for dec in obj.decorators
+            ):
+                # Update the template
+                self.__modify_pxt_udf(obj)
+
+    def __modify_pxt_udf(self, func: griffe.Function) -> None:
+        """
+        Instructs the doc snippet for `func` to use the custom Pixeltable UDF jinja template, and
+        converts all type hints to Pixeltable column type references, in accordance with the @udf
+        decorator behavior.
+        """
+        func.extra['mkdocstrings']['template'] = 'udf.html.jinja'
+        # Dynamically load the UDF reference so we can inspect the Pixeltable signature directly
+        warnings.simplefilter("ignore")
+        udf = griffe.dynamic_import(func.path)
+        assert isinstance(udf, pxt.Function)
+        # Convert the return type to a Pixeltable type reference
+        func.returns = self.__column_type_to_display_str(udf.signature.get_return_type())
+        # Convert the parameter types to Pixeltable type references
+        for griffe_param in func.parameters:
+            assert isinstance(griffe_param.annotation, griffe.expressions.Expr)
+            if griffe_param.name not in udf.signature.parameters:
+                logger.warning(f'Parameter `{griffe_param.name}` not found in signature for UDF: {udf.display_name}')
+                continue
+            pxt_param = udf.signature.parameters[griffe_param.name]
+            griffe_param.annotation = self.__column_type_to_display_str(pxt_param.col_type)
+
+    def __column_type_to_display_str(self, column_type: Optional[pxt.ColumnType]) -> str:
+        # TODO: When we enhance the Pixeltable type system, we may want to refactor some of this logic out.
+        #   I'm putting it here for now though.
+        if column_type is None:
+            return 'None'
+        if column_type.is_string_type():
+            base = 'str'
+        elif column_type.is_int_type():
+            base = 'int'
+        elif column_type.is_float_type():
+            base = 'float'
+        elif column_type.is_bool_type():
+            base = 'bool'
+        elif column_type.is_timestamp_type():
+            base = 'datetime'
+        elif column_type.is_array_type():
+            base = 'ArrayT'
+        elif column_type.is_json_type():
+            base = 'JsonT'
+        elif column_type.is_image_type():
+            base = 'ImageT'
+        elif column_type.is_video_type():
+            base = 'VideoT'
+        elif column_type.is_audio_type():
+            base = 'AudioT'
+        elif column_type.is_document_type():
+            base = 'DocumentT'
+        else:
+            assert False
+        return f'Optional[{base}]' if column_type.nullable else base

pixeltable/tool/doc_plugins/mkdocstrings.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+
+
+def get_templates_path() -> Path:
+    """Implementation of the 'mkdocstrings.python.templates' plugin for custom jinja templates."""
+    return Path(__file__).parent / "templates"

pixeltable/tool/doc_plugins/templates/material/udf.html.jinja

@@ -0,0 +1,135 @@
+{#- Template for Pixeltable UDFs. Cargo-culted (with modification) from _base/function.html.jinja. -#}
+
+{% block logs scoped %}
+  {#- Logging block.
+
+  This block can be used to log debug messages, deprecation messages, warnings, etc.
+  -#}
+  {{ log.debug("Rendering " + function.path) }}
+{% endblock logs %}
+
+{% import "language"|get_template as lang with context %}
+{#- Language module providing the `t` translation method. -#}
+
+<div class="doc doc-object doc-function">
+  {% with obj = function, html_id = function.path %}
+
+    {% if root %}
+      {% set show_full_path = config.show_root_full_path %}
+      {% set root_members = True %}
+    {% elif root_members %}
+      {% set show_full_path = config.show_root_members_full_path or config.show_object_full_path %}
+      {% set root_members = False %}
+    {% else %}
+      {% set show_full_path = config.show_object_full_path %}
+    {% endif %}
+
+    {% set function_name = function.path if show_full_path else function.name %}
+    {#- Brief or full function name depending on configuration. -#}
+    {% set symbol_type = "udf" %}
+    {#- Symbol type: method when parent is a class, function otherwise. -#}
+
+    {% if not root or config.show_root_heading %}
+      {% filter heading(
+          heading_level,
+          role="function",
+          id=html_id,
+          class="doc doc-heading",
+          toc_label=(('<code class="doc-symbol doc-symbol-toc doc-symbol-' + symbol_type + '"></code>&nbsp;')|safe if config.show_symbol_type_toc else '') + function.name,
+        ) %}
+
+        {% block heading scoped %}
+          {#- Heading block.
+
+          This block renders the heading for the function.
+          -#}
+          {% if config.show_symbol_type_heading %}<code class="doc-symbol doc-symbol-heading doc-symbol-{{ symbol_type }}"></code>{% endif %}
+          {% if config.separate_signature %}
+            <span class="doc doc-object-name doc-function-name">{{ function_name }}</span>
+          {% else %}
+            {%+ filter highlight(language="python", inline=True) %}
+              {{ function_name }}{% include "signature"|get_template with context %}
+            {% endfilter %}
+          {% endif %}
+        {% endblock heading %}
+
+        {% block labels scoped %}
+          {#- Labels block.
+
+          This block renders the labels for the function.
+          -#}
+          {% with labels = function.labels %}
+            {% include "labels"|get_template with context %}
+          {% endwith %}
+        {% endblock labels %}
+
+      {% endfilter %}
+
+      {% block signature scoped %}
+        {#- Signature block.
+
+        This block renders the signature for the function.
+        -#}
+        {% if config.separate_signature %}
+          {% filter format_signature(function, config.line_length, crossrefs=config.signature_crossrefs) %}
+            {{ function.name }}
+          {% endfilter %}
+        {% endif %}
+      {% endblock signature %}
+
+    {% else %}
+
+      {% if config.show_root_toc_entry %}
+        {% filter heading(
+            heading_level,
+            role="function",
+            id=html_id,
+            toc_label=(('<code class="doc-symbol doc-symbol-toc doc-symbol-' + symbol_type + '"></code>&nbsp;')|safe if config.show_symbol_type_toc else '') + function.name,
+            hidden=True,
+          ) %}
+        {% endfilter %}
+      {% endif %}
+      {% set heading_level = heading_level - 1 %}
+    {% endif %}
+
+    <div class="doc doc-contents {% if root %}first{% endif %}">
+      {% block contents scoped %}
+        {#- Contents block.
+
+        This block renders the contents of the function.
+        It contains other blocks that users can override.
+        Overriding the contents block allows to rearrange the order of the blocks.
+        -#}
+        {% block docstring scoped %}
+          {#- Docstring block.
+
+          This block renders the docstring for the function.
+          -#}
+          {% with docstring_sections = function.docstring.parsed %}
+            {% include "docstring"|get_template with context %}
+          {% endwith %}
+        {% endblock docstring %}
+
+        {% block source scoped %}
+          {#- Source block.
+
+          This block renders the source code for the function.
+          -#}
+          {% if config.show_source and function.source %}
+            <details class="quote">
+              <summary>{{ lang.t("Source code in") }} <code>
+                {%- if function.relative_filepath.is_absolute() -%}
+                  {{ function.relative_package_filepath }}
+                {%- else -%}
+                  {{ function.relative_filepath }}
+                {%- endif -%}
+              </code></summary>
+              {{ function.source|highlight(language="python", linestart=function.lineno, linenums=True) }}
+            </details>
+          {% endif %}
+        {% endblock source %}
+      {% endblock contents %}
+    </div>
+
+  {% endwith %}
+</div>

pixeltable/type_system.py

@@ -183,7 +183,7 @@ class ColumnType:
         """Two types match if they're equal, aside from nullability"""
         if not isinstance(other, ColumnType):
             pass
-        assert isinstance(other, ColumnType)
+        assert isinstance(other, ColumnType), type(other)
         if type(self) != type(other):
             return False
         for member_var in vars(self).keys():
@@ -206,7 +206,7 @@ class ColumnType:
         if type1.is_scalar_type() and type2.is_scalar_type():
             t = cls.Type.supertype(type1._type, type2._type, cls.common_supertypes)
             if t is not None:
-                return cls.make_type(t)
+                return cls.make_type(t).copy(nullable=(type1.nullable or type2.nullable))
             return None
 
         if type1._type == type2._type:
@@ -227,22 +227,23 @@ class ColumnType:
     def infer_literal_type(cls, val: Any) -> Optional[ColumnType]:
         if isinstance(val, str):
             return StringType()
+        if isinstance(val, bool):
+            # We have to check bool before int, because isinstance(b, int) is True if b is a Python bool
+            return BoolType()
         if isinstance(val, int):
             return IntType()
         if isinstance(val, float):
             return FloatType()
-        if isinstance(val, bool):
-            return BoolType()
-        if isinstance(val, datetime.datetime) or isinstance(val, datetime.date):
+        if isinstance(val, datetime.datetime):
             return TimestampType()
         if isinstance(val, PIL.Image.Image):
-            return ImageType(width=val.width, height=val.height)
+            return ImageType(width=val.width, height=val.height, mode=val.mode)
         if isinstance(val, np.ndarray):
             col_type = ArrayType.from_literal(val)
             if col_type is not None:
                 return col_type
             # this could still be json-serializable
-        if isinstance(val, dict) or isinstance(val, np.ndarray):
+        if isinstance(val, dict) or isinstance(val, list) or isinstance(val, np.ndarray):
             try:
                 JsonType().validate_literal(val)
                 return JsonType()
@@ -276,7 +277,7 @@ class ColumnType:
                 return FloatType()
             if base is bool:
                 return BoolType()
-            if base is datetime.date or base is datetime.datetime:
+            if base is datetime.datetime:
                 return TimestampType()
             if issubclass(base, Sequence) or issubclass(base, Mapping):
                 return JsonType()
@@ -425,7 +426,7 @@ class StringType(ColumnType):
     def conversion_fn(self, target: ColumnType) -> Optional[Callable[[Any], Any]]:
         if not target.is_timestamp_type():
             return None
-        def convert(val: str) -> Optional[datetime]:
+        def convert(val: str) -> Optional[datetime.datetime]:
             try:
                 dt = datetime.datetime.fromisoformat(val)
                 return dt
@@ -506,8 +507,8 @@ class TimestampType(ColumnType):
         return sql.TIMESTAMP()
 
     def _validate_literal(self, val: Any) -> None:
-        if not isinstance(val, datetime.datetime) and not isinstance(val, datetime.date):
-            raise TypeError(f'Expected datetime.datetime or datetime.date, got {val.__class__.__name__}')
+        if not isinstance(val, datetime.datetime):
+            raise TypeError(f'Expected datetime.datetime, got {val.__class__.__name__}')
 
     def _create_literal(self, val: Any) -> Any:
         if isinstance(val, str):
@@ -577,7 +578,7 @@ class ArrayType(ColumnType):
         if base_type is None:
             return None
         shape = [n1 if n1 == n2 else None for n1, n2 in zip(type1.shape, type2.shape)]
-        return ArrayType(tuple(shape), base_type)
+        return ArrayType(tuple(shape), base_type, nullable=(type1.nullable or type2.nullable))
 
     def _as_dict(self) -> Dict:
         result = super()._as_dict()
@@ -609,7 +610,7 @@ class ArrayType(ColumnType):
             dtype = StringType()
         else:
             return None
-        return cls(val.shape, dtype=dtype, nullable=True)
+        return cls(val.shape, dtype=dtype)
 
     def is_valid_literal(self, val: np.ndarray) -> bool:
         if not isinstance(val, np.ndarray):
@@ -695,7 +696,7 @@ class ImageType(ColumnType):
         return f'{self._type.name.lower()}{params_str}'
 
     def _is_supertype_of(self, other: ImageType) -> bool:
-        if self.mode != other.mode:
+        if self.mode is not None and self.mode != other.mode:
             return False
         if self.width is None and self.height is None:
             return True

pixeltable/utils/s3.py

@@ -10,4 +10,4 @@ def get_client() -> Any:
     except AttributeError:
         # No credentials available, use unsigned mode
         config = botocore.config.Config(signature_version=botocore.UNSIGNED)
-        return boto3.client('s3', config=config)
+        return boto3.client('s3', config=config)