pixeltable 0.2.13__py3-none-any.whl → 0.2.15__py3-none-any.whl

pixeltable/store.py

@@ -223,35 +223,70 @@ class StoreBase:
         """
         num_excs = 0
         num_rows = 0
-        for row_batch in exec_plan:
-            num_rows += len(row_batch)
-            for result_row in row_batch:
-                values_dict: Dict[sql.Column, Any] = {}
-
-                if col.is_computed:
-                    if result_row.has_exc(value_expr_slot_idx):
-                        num_excs += 1
-                        value_exc = result_row.get_exc(value_expr_slot_idx)
-                        # we store a NULL value and record the exception/exc type
-                        error_type = type(value_exc).__name__
-                        error_msg = str(value_exc)
-                        values_dict = {
-                            col.sa_col: None,
-                            col.sa_errortype_col: error_type,
-                            col.sa_errormsg_col: error_msg
-                        }
-                    else:
-                        val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
-                        if col.col_type.is_media_type():
-                            val = self._move_tmp_media_file(val, col, result_row.pk[-1])
-                        values_dict = {col.sa_col: val}
-
-                update_stmt = sql.update(self.sa_tbl).values(values_dict)
-                for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
-                    update_stmt = update_stmt.where(pk_col == pk_val)
-                log_stmt(_logger, update_stmt)
-                conn.execute(update_stmt)
 
+        # create temp table to store output of exec_plan, with the same primary key as the store table
+        tmp_name = f'temp_{self._storage_name()}'
+        tmp_pk_cols = [sql.Column(col.name, col.type, primary_key=True) for col in self.pk_columns()]
+        tmp_cols = tmp_pk_cols.copy()
+        tmp_val_col = sql.Column(col.sa_col.name, col.sa_col.type)
+        tmp_cols.append(tmp_val_col)
+        # add error columns if the store column records errors
+        if col.records_errors:
+            tmp_errortype_col = sql.Column(col.sa_errortype_col.name, col.sa_errortype_col.type)
+            tmp_cols.append(tmp_errortype_col)
+            tmp_errormsg_col = sql.Column(col.sa_errormsg_col.name, col.sa_errormsg_col.type)
+            tmp_cols.append(tmp_errormsg_col)
+        tmp_tbl = sql.Table(tmp_name, self.sa_md, *tmp_cols, prefixes=['TEMPORARY'])
+        tmp_tbl.create(bind=conn)
+
+        try:
+            # insert rows from exec_plan into temp table
+            for row_batch in exec_plan:
+                num_rows += len(row_batch)
+                tbl_rows: list[dict[str, Any]] = []
+                for result_row in row_batch:
+                    tbl_row: dict[str, Any] = {}
+                    for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
+                        tbl_row[pk_col.name] = pk_val
+
+                    if col.is_computed:
+                        if result_row.has_exc(value_expr_slot_idx):
+                            num_excs += 1
+                            value_exc = result_row.get_exc(value_expr_slot_idx)
+                            # we store a NULL value and record the exception/exc type
+                            error_type = type(value_exc).__name__
+                            error_msg = str(value_exc)
+                            tbl_row[col.sa_col.name] = None
+                            tbl_row[col.sa_errortype_col.name] = error_type
+                            tbl_row[col.sa_errormsg_col.name] = error_msg
+                        else:
+                            val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
+                            if col.col_type.is_media_type():
+                                val = self._move_tmp_media_file(val, col, result_row.pk[-1])
+                            tbl_row[col.sa_col.name] = val
+                            if col.records_errors:
+                                tbl_row[col.sa_errortype_col.name] = None
+                                tbl_row[col.sa_errormsg_col.name] = None
+
+                    tbl_rows.append(tbl_row)
+                conn.execute(sql.insert(tmp_tbl), tbl_rows)
+
+            # update store table with values from temp table
+            update_stmt = sql.update(self.sa_tbl)
+            for pk_col, tmp_pk_col in zip(self.pk_columns(), tmp_pk_cols):
+                update_stmt = update_stmt.where(pk_col == tmp_pk_col)
+            update_stmt = update_stmt.values({col.sa_col: tmp_val_col})
+            if col.records_errors:
+                update_stmt = update_stmt.values({
+                    col.sa_errortype_col: tmp_errortype_col,
+                    col.sa_errormsg_col: tmp_errormsg_col
+                })
+            log_explain(_logger, update_stmt, conn)
+            conn.execute(update_stmt)
+
+        finally:
+            tmp_tbl.drop(bind=conn)
+            self.sa_md.remove(tmp_tbl)
         return num_excs
 
     def insert_rows(
@@ -295,6 +330,8 @@ class StoreBase:
                                     file=sys.stdout
                                 )
                             progress_bar.update(1)
+
+                    # insert batch of rows
                     self._move_tmp_media_files(table_rows, media_cols, v_min)
                     conn.execute(sql.insert(self.sa_tbl), table_rows)
             if progress_bar is not None:

pixeltable/tool/create_test_db_dump.py

@@ -6,7 +6,7 @@ import pathlib
 import subprocess
 from typing import Any
 
-import pgserver
+import pixeltable_pgserver
 import toml
 
 import pixeltable as pxt
@@ -41,7 +41,7 @@ class Dumper:
         md_version = metadata.VERSION
         dump_file = self.output_dir / f'pixeltable-v{md_version:03d}-test.dump.gz'
         _logger.info(f'Creating database dump at: {dump_file}')
-        pg_package_dir = os.path.dirname(pgserver.__file__)
+        pg_package_dir = os.path.dirname(pixeltable_pgserver.__file__)
         pg_dump_binary = f'{pg_package_dir}/pginstall/bin/pg_dump'
         _logger.info(f'Using pg_dump binary at: {pg_dump_binary}')
         with open(dump_file, 'wb') as dump:
@@ -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:
@@ -177,9 +177,9 @@ class Dumper:
         assert t.base_table_image_rot.col in project.stored_proxies
 
     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})
-        
+        def add_column(col_name: str, col_expr: Any, stored: bool = True) -> None:
+            t.add_column(**{f'{col_prefix}_{col_name}': col_expr}, stored=stored)
+
         # arithmetic_expr
         add_column('plus', t.c2 + 6)
         add_column('minus', t.c2 - 5)
@@ -217,7 +217,7 @@ class Dumper:
 
         # image_member_access
         add_column('image_mode', t.c8.mode)
-        add_column('image_rot', t.c8.rotate(180))
+        add_column('image_rot', t.c8.rotate(180), stored=False)
 
         # in_predicate
         add_column('isin_1', t.c1.isin(['test string 1', 'test string 2', 'test string 3']))
@@ -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

@@ -1,5 +1,6 @@
 import ast
 from typing import Optional, Union
+import warnings
 
 import griffe
 import griffe.expressions
@@ -34,6 +35,7 @@ class PxtGriffeExtension(Extension):
         """
         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
@@ -60,6 +62,8 @@ class PxtGriffeExtension(Extension):
             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():

pixeltable/type_system.py

@@ -9,7 +9,7 @@ import urllib.parse
 import urllib.request
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Optional, Tuple, Dict, Callable, List, Union, Sequence, Mapping
+from typing import Any, Iterable, Optional, Tuple, Dict, Callable, List, Union, Sequence, Mapping
 
 import PIL.Image
 import av
@@ -166,24 +166,14 @@ class ColumnType:
         return self._type.name.lower()
 
     def __eq__(self, other: object) -> bool:
-        return self.matches(other) and self.nullable == other.nullable
+        return isinstance(other, ColumnType) and self.matches(other) and self.nullable == other.nullable
 
-    def is_supertype_of(self, other: ColumnType) -> bool:
-        if type(self) != type(other):
-            return False
-        if self.matches(other):
-            return True
-        return self._is_supertype_of(other)
-
-    @abc.abstractmethod
-    def _is_supertype_of(self, other: ColumnType) -> bool:
-        return False
+    def is_supertype_of(self, other: ColumnType, ignore_nullable: bool = False) -> bool:
+        operand = self.copy(nullable=True) if ignore_nullable else self
+        return operand.supertype(other) == operand
 
-    def matches(self, other: object) -> bool:
+    def matches(self, other: ColumnType) -> bool:
         """Two types match if they're equal, aside from nullability"""
-        if not isinstance(other, ColumnType):
-            pass
-        assert isinstance(other, ColumnType)
         if type(self) != type(other):
             return False
         for member_var in vars(self).keys():
@@ -193,56 +183,44 @@ class ColumnType:
                 return False
         return True
 
-    @classmethod
-    def supertype(cls, type1: ColumnType, type2: ColumnType) -> Optional[ColumnType]:
-        if type1 == type2:
-            return type1
+    def supertype(self, other: ColumnType) -> Optional[ColumnType]:
+        if self.copy(nullable=True) == other.copy(nullable=True):
+            return self.copy(nullable=(self.nullable or other.nullable))
 
-        if type1.is_invalid_type():
-            return type2
-        if type2.is_invalid_type():
-            return type1
+        if self.is_invalid_type():
+            return other
+        if other.is_invalid_type():
+            return self
 
-        if type1.is_scalar_type() and type2.is_scalar_type():
-            t = cls.Type.supertype(type1._type, type2._type, cls.common_supertypes)
+        if self.is_scalar_type() and other.is_scalar_type():
+            t = self.Type.supertype(self._type, other._type, self.common_supertypes)
             if t is not None:
-                return cls.make_type(t)
+                return self.make_type(t).copy(nullable=(self.nullable or other.nullable))
             return None
 
-        if type1._type == type2._type:
-            return cls._supertype(type1, type2)
-
         return None
 
-    @classmethod
-    @abc.abstractmethod
-    def _supertype(cls, type1: ColumnType, type2: ColumnType) -> Optional[ColumnType]:
-        """
-        Class-specific implementation of determining the supertype. type1 and type2 are from the same subclass of
-        ColumnType.
-        """
-        pass
-
     @classmethod
     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()
@@ -250,6 +228,26 @@ class ColumnType:
                 return None
         return None
 
+    @classmethod
+    def infer_common_literal_type(cls, vals: Iterable[Any]) -> Optional[ColumnType]:
+        """
+        Returns the most specific type that is a supertype of all literals in `vals`. If no such type
+        exists, returns None.
+
+        Args:
+            vals: A collection of literals.
+        """
+        inferred_type: Optional[ColumnType] = None
+        for val in vals:
+            val_type = cls.infer_literal_type(val)
+            if inferred_type is None:
+                inferred_type = val_type
+            else:
+                inferred_type = inferred_type.supertype(val_type)
+                if inferred_type is None:
+                    return None
+        return inferred_type
+
     @classmethod
     def from_python_type(cls, t: type) -> Optional[ColumnType]:
         if typing.get_origin(t) is typing.Union:
@@ -276,7 +274,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 +423,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 +504,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):
@@ -521,6 +519,28 @@ class JsonType(ColumnType):
         super().__init__(self.Type.JSON, nullable=nullable)
         self.type_spec = type_spec
 
+    def supertype(self, other: ColumnType) -> Optional[JsonType]:
+        if self.type_spec is None:
+            # we don't have a type spec and can accept anything accepted by other
+            return JsonType(nullable=(self.nullable or other.nullable))
+        if other.type_spec is None:
+            # we have a type spec but other doesn't
+            return JsonType(nullable=(self.nullable or other.nullable))
+
+        # we both have type specs; the supertype's type spec is the union of the two
+        type_spec = deepcopy(self.type_spec)
+        for other_field_name, other_field_type in other.type_spec.items():
+            if other_field_name not in type_spec:
+                type_spec[other_field_name] = other_field_type.copy()
+            else:
+                # both type specs have this field
+                field_type = type_spec[other_field_name].supertype(other_field_type)
+                if field_type is None:
+                    # conflicting types
+                    return JsonType(nullable=(self.nullable or other.nullable))
+                type_spec[other_field_name] = field_type
+        return JsonType(type_spec, nullable=(self.nullable or other.nullable))
+
     def _as_dict(self) -> Dict:
         result = super()._as_dict()
         if self.type_spec is not None:
@@ -563,21 +583,22 @@ class JsonType(ColumnType):
 
 
 class ArrayType(ColumnType):
-    def __init__(
-            self, shape: Tuple[Union[int, None], ...], dtype: ColumnType, nullable: bool = False):
+    def __init__(self, shape: Tuple[Union[int, None], ...], dtype: ColumnType, nullable: bool = False):
         super().__init__(self.Type.ARRAY, nullable=nullable)
         self.shape = shape
         assert dtype.is_int_type() or dtype.is_float_type() or dtype.is_bool_type() or dtype.is_string_type()
         self.dtype = dtype._type
 
-    def _supertype(cls, type1: ArrayType, type2: ArrayType) -> Optional[ArrayType]:
-        if len(type1.shape) != len(type2.shape):
+    def supertype(self, other: ColumnType) -> Optional[ArrayType]:
+        if not isinstance(other, ArrayType):
+            return None
+        if len(self.shape) != len(other.shape):
             return None
-        base_type = ColumnType.supertype(type1.dtype, type2.dtype)
+        base_type = self.Type.supertype(self.dtype, other.dtype, self.common_supertypes)
         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)
+        shape = [n1 if n1 == n2 else None for n1, n2 in zip(self.shape, other.shape)]
+        return ArrayType(tuple(shape), self.make_type(base_type), nullable=(self.nullable or other.nullable))
 
     def _as_dict(self) -> Dict:
         result = super()._as_dict()
@@ -609,7 +630,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):
@@ -694,13 +715,13 @@ class ImageType(ColumnType):
             params_str = ''
         return f'{self._type.name.lower()}{params_str}'
 
-    def _is_supertype_of(self, other: ImageType) -> bool:
-        if self.mode != other.mode:
-            return False
-        if self.width is None and self.height is None:
-            return True
-        if self.width != other.width and self.height != other.height:
-            return False
+    def supertype(self, other: ColumnType) -> Optional[ImageType]:
+        if not isinstance(other, ImageType):
+            return None
+        width = self.width if self.width == other.width else None
+        height = self.height if self.height == other.height else None
+        mode = self.mode if self.mode == other.mode else None
+        return ImageType(width=width, height=height, mode=mode, nullable=(self.nullable or other.nullable))
 
     @property
     def size(self) -> Optional[Tuple[int, int]]:

{pixeltable-0.2.13.dist-info → pixeltable-0.2.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pixeltable
-Version: 0.2.13
+Version: 0.2.15
 Summary: Pixeltable: The Multimodal AI Data Plane
 Author: Pixeltable, Inc.
 Author-email: contact@pixeltable.com
@@ -21,9 +21,9 @@ Requires-Dist: more-itertools (>=10.2,<11.0)
 Requires-Dist: numpy (>=1.25)
 Requires-Dist: opencv-python-headless (>=4.7.0.68,<5.0.0.0)
 Requires-Dist: pandas (>=2.0,<3.0)
-Requires-Dist: pgserver (==0.1.4)
 Requires-Dist: pgvector (>=0.2.1,<0.3.0)
 Requires-Dist: pillow (>=9.3.0)
+Requires-Dist: pixeltable-pgserver (==0.2.4)
 Requires-Dist: psutil (>=5.9.5,<6.0.0)
 Requires-Dist: psycopg2-binary (>=2.9.5,<3.0.0)
 Requires-Dist: pymupdf (>=1.24.1,<2.0.0)