pixeltable 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

pixeltable/store.py

@@ -12,10 +12,8 @@ from typing import Any, Iterator, Literal, Optional, Union
 import sqlalchemy as sql
 from tqdm import TqdmWarning, tqdm
 
-import pixeltable.catalog as catalog
-import pixeltable.env as env
-import pixeltable.exceptions as excs
-from pixeltable import exprs
+from pixeltable import catalog, exceptions as excs, exprs
+from pixeltable.env import Env
 from pixeltable.exec import ExecNode
 from pixeltable.metadata import schema
 from pixeltable.utils.media_store import MediaStore
@@ -33,7 +31,7 @@ class StoreBase:
     - v_max: version at which the row was deleted (or MAX_VERSION if it's still live)
     """
 
-    tbl_version: catalog.TableVersion
+    tbl_version: catalog.TableVersionHandle
     sa_md: sql.MetaData
     sa_tbl: Optional[sql.Table]
     _pk_cols: list[sql.Column]
@@ -44,12 +42,14 @@ class StoreBase:
     __INSERT_BATCH_SIZE = 1000
 
     def __init__(self, tbl_version: catalog.TableVersion):
-        self.tbl_version = tbl_version
+        self.tbl_version = catalog.TableVersionHandle(
+            tbl_version.id, tbl_version.effective_version, tbl_version=tbl_version
+        )
         self.sa_md = sql.MetaData()
         self.sa_tbl = None
         # We need to declare a `base` variable here, even though it's only defined for instances of `StoreView`,
         # since it's referenced by various methods of `StoreBase`
-        self.base = None if tbl_version.base is None else tbl_version.base.store_tbl
+        self.base = tbl_version.base.get().store_tbl if tbl_version.base is not None else None
         self.create_sa_tbl()
 
     def pk_columns(self) -> list[sql.Column]:
@@ -76,7 +76,7 @@ class StoreBase:
         """Create self.sa_tbl from self.tbl_version."""
         system_cols = self._create_system_columns()
         all_cols = system_cols.copy()
-        for col in [c for c in self.tbl_version.cols if c.is_stored]:
+        for col in [c for c in self.tbl_version.get().cols if c.is_stored]:
             # re-create sql.Column for each column, regardless of whether it already has sa_col set: it was bound
             # to the last sql.Table version we created and cannot be reused
             col.create_sa_cols()
@@ -115,7 +115,7 @@ class StoreBase:
 
     def _move_tmp_media_file(self, file_url: Optional[str], col: catalog.Column, v_min: int) -> str:
         """Move tmp media file with given url to Env.media_dir and return new url, or given url if not a tmp_dir file"""
-        pxt_tmp_dir = str(env.Env.get().tmp_dir)
+        pxt_tmp_dir = str(Env.get().tmp_dir)
         if file_url is None:
             return None
         parsed = urllib.parse.urlparse(file_url)
@@ -158,36 +158,36 @@ class StoreBase:
             table_row[pk_col.name] = pk_val
         return table_row, num_excs
 
-    def count(self, conn: Optional[sql.engine.Connection] = None) -> int:
+    def count(self) -> int:
         """Return the number of rows visible in self.tbl_version"""
         stmt = (
             sql.select(sql.func.count('*'))
             .select_from(self.sa_tbl)
-            .where(self.v_min_col <= self.tbl_version.version)
-            .where(self.v_max_col > self.tbl_version.version)
+            .where(self.v_min_col <= self.tbl_version.get().version)
+            .where(self.v_max_col > self.tbl_version.get().version)
         )
-        if conn is None:
-            with env.Env.get().engine.connect() as conn:
-                result = conn.execute(stmt).scalar_one()
-        else:
-            result = conn.execute(stmt).scalar_one()
+        conn = Env.get().conn
+        result = conn.execute(stmt).scalar_one()
         assert isinstance(result, int)
         return result
 
-    def create(self, conn: sql.engine.Connection) -> None:
+    def create(self) -> None:
+        conn = Env.get().conn
         self.sa_md.create_all(bind=conn)
 
-    def drop(self, conn: sql.engine.Connection) -> None:
+    def drop(self) -> None:
         """Drop store table"""
+        conn = Env.get().conn
         self.sa_md.drop_all(bind=conn)
 
-    def add_column(self, col: catalog.Column, conn: sql.engine.Connection) -> None:
+    def add_column(self, col: catalog.Column) -> None:
         """Add column(s) to the store-resident table based on a catalog column
 
         Note that a computed catalog column will require two extra columns (for the computed value and for the error
         message).
         """
         assert col.is_stored
+        conn = Env.get().conn
         col_type_str = col.get_sa_col_type().compile(dialect=conn.dialect)
         stmt = sql.text(f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.store_name()} {col_type_str} NULL')
         log_stmt(_logger, stmt)
@@ -207,8 +207,9 @@ class StoreBase:
         self.create_sa_tbl()
         _logger.info(f'Added columns {added_storage_cols} to storage table {self._storage_name()}')
 
-    def drop_column(self, col: catalog.Column, conn: sql.engine.Connection) -> None:
+    def drop_column(self, col: catalog.Column) -> None:
         """Execute Alter Table Drop Column statement"""
+        conn = Env.get().conn
         stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.store_name()}'
         conn.execute(sql.text(stmt))
         if col.records_errors:
@@ -218,12 +219,7 @@ class StoreBase:
             conn.execute(sql.text(stmt))
 
     def load_column(
-        self,
-        col: catalog.Column,
-        exec_plan: ExecNode,
-        value_expr_slot_idx: int,
-        conn: sql.engine.Connection,
-        on_error: Literal['abort', 'ignore'],
+        self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, on_error: Literal['abort', 'ignore']
     ) -> int:
         """Update store column of a computed column with values produced by an execution plan
 
@@ -250,6 +246,7 @@ class StoreBase:
             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'])
+        conn = Env.get().conn
         tmp_tbl.create(bind=conn)
 
         try:
@@ -280,7 +277,7 @@ class StoreBase:
                         else:
                             if col.col_type.is_image_type() and result_row.file_urls[value_expr_slot_idx] is None:
                                 # we have yet to store this image
-                                filepath = str(MediaStore.prepare_media_path(col.tbl.id, col.id, col.tbl.version))
+                                filepath = str(MediaStore.prepare_media_path(col.tbl.id, col.id, col.tbl.get().version))
                                 result_row.flush_img(value_expr_slot_idx, filepath)
                             val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
                             if col.col_type.is_media_type():
@@ -313,7 +310,6 @@ class StoreBase:
     def insert_rows(
         self,
         exec_plan: ExecNode,
-        conn: sql.engine.Connection,
         v_min: Optional[int] = None,
         show_progress: bool = True,
         rowids: Optional[Iterator[int]] = None,
@@ -324,7 +320,6 @@ class StoreBase:
             number of inserted rows, number of exceptions, set of column ids that have exceptions
         """
         assert v_min is not None
-        exec_plan.ctx.set_conn(conn)
         # TODO: total?
         num_excs = 0
         num_rows = 0
@@ -332,6 +327,8 @@ class StoreBase:
         progress_bar: Optional[tqdm] = None  # create this only after we started executing
         row_builder = exec_plan.row_builder
         media_cols = [info.col for info in row_builder.table_columns if info.col.col_type.is_media_type()]
+        conn = Env.get().conn
+
         try:
             exec_plan.open()
             for row_batch in exec_plan:
@@ -348,7 +345,7 @@ class StoreBase:
                             raise exc
 
                         rowid = (next(rowids),) if rowids is not None else row.pk[:-1]
-                        pk = rowid + (v_min,)
+                        pk = (*rowid, v_min)
                         table_row, num_row_exc = self._create_table_row(row, row_builder, cols_with_excs, pk=pk)
                         num_excs += num_row_exc
                         table_rows.append(table_row)
@@ -357,7 +354,7 @@ class StoreBase:
                             if progress_bar is None:
                                 warnings.simplefilter('ignore', category=TqdmWarning)
                                 progress_bar = tqdm(
-                                    desc=f'Inserting rows into `{self.tbl_version.name}`',
+                                    desc=f'Inserting rows into `{self.tbl_version.get().name}`',
                                     unit=' rows',
                                     ncols=100,
                                     file=sys.stdout,
@@ -378,7 +375,9 @@ class StoreBase:
         v = versions[0]
         if v is None:
             # we're looking at live rows
-            clause = sql.and_(self.v_min_col <= self.tbl_version.version, self.v_max_col == schema.Table.MAX_VERSION)
+            clause = sql.and_(
+                self.v_min_col <= self.tbl_version.get().version, self.v_max_col == schema.Table.MAX_VERSION
+            )
         else:
             # we're looking at a specific version
             clause = self.v_min_col == v if match_on_vmin else self.v_max_col == v
@@ -392,7 +391,6 @@ class StoreBase:
         base_versions: list[Optional[int]],
         match_on_vmin: bool,
         where_clause: Optional[sql.ColumnElement[bool]],
-        conn: sql.engine.Connection,
     ) -> int:
         """Mark rows as deleted that are live and were created prior to current_version.
         Also: populate the undo columns
@@ -414,7 +412,7 @@ class StoreBase:
             sql.true() if len(base_versions) == 0 else self.base._versions_clause(base_versions, match_on_vmin)
         )
         set_clause: dict[sql.Column, Union[int, sql.Column]] = {self.v_max_col: current_version}
-        for index_info in self.tbl_version.idxs_by_name.values():
+        for index_info in self.tbl_version.get().idxs_by_name.values():
             # copy value column to undo column
             set_clause[index_info.undo_col.sa_col] = index_info.val_col.sa_col
             # set value column to NULL
@@ -426,6 +424,7 @@ class StoreBase:
             .where(rowid_join_clause)
             .where(base_versions_clause)
         )
+        conn = Env.get().conn
         log_explain(_logger, stmt, conn)
         status = conn.execute(stmt)
         return status.rowcount
@@ -433,7 +432,7 @@ class StoreBase:
 
 class StoreTable(StoreBase):
     def __init__(self, tbl_version: catalog.TableVersion):
-        assert not tbl_version.is_view()
+        assert not tbl_version.is_view
         super().__init__(tbl_version)
 
     def _create_rowid_columns(self) -> list[sql.Column]:
@@ -449,7 +448,7 @@ class StoreTable(StoreBase):
 
 class StoreView(StoreBase):
     def __init__(self, catalog_view: catalog.TableVersion):
-        assert catalog_view.is_view()
+        assert catalog_view.is_view
         super().__init__(catalog_view)
 
     def _create_rowid_columns(self) -> list[sql.Column]:
@@ -492,7 +491,7 @@ class StoreComponentView(StoreView):
     def create_sa_tbl(self) -> None:
         super().create_sa_tbl()
         # we need to fix up the 'pos' column in TableVersion
-        self.tbl_version.cols_by_name['pos'].sa_col = self.pos_col
+        self.tbl_version.get().cols_by_name['pos'].sa_col = self.pos_col
 
     def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.and_(

pixeltable/type_system.py

@@ -8,7 +8,7 @@ import json
 import typing
 import urllib.parse
 import urllib.request
-from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, Union
+from typing import Any, ClassVar, Iterable, Literal, Mapping, Optional, Sequence, Union
 
 import av
 import jsonschema
@@ -81,9 +81,9 @@ class ColumnType:
         FLOAT32 = (10,)
         FLOAT64 = 11
 
-    scalar_types = {Type.STRING, Type.INT, Type.FLOAT, Type.BOOL, Type.TIMESTAMP}
-    numeric_types = {Type.INT, Type.FLOAT}
-    common_supertypes: dict[tuple[Type, Type], Type] = {
+    scalar_types: ClassVar[set[Type]] = {Type.STRING, Type.INT, Type.FLOAT, Type.BOOL, Type.TIMESTAMP}
+    numeric_types: ClassVar[set[Type]] = {Type.INT, Type.FLOAT}
+    common_supertypes: ClassVar[dict[tuple[Type, Type], Type]] = {
         (Type.BOOL, Type.INT): Type.INT,
         (Type.BOOL, Type.FLOAT): Type.FLOAT,
         (Type.INT, Type.FLOAT): Type.FLOAT,
@@ -150,7 +150,7 @@ class ColumnType:
 
     @classmethod
     def make_type(cls, t: Type) -> ColumnType:
-        assert t != cls.Type.INVALID and t != cls.Type.ARRAY
+        assert t != cls.Type.INVALID
         if t == cls.Type.STRING:
             return StringType()
         if t == cls.Type.INT:
@@ -161,6 +161,8 @@ class ColumnType:
             return BoolType()
         if t == cls.Type.TIMESTAMP:
             return TimestampType()
+        if t == cls.Type.ARRAY:
+            return ArrayType()
         if t == cls.Type.JSON:
             return JsonType()
         if t == cls.Type.IMAGE:
@@ -364,7 +366,7 @@ class ColumnType:
             cls.__raise_exc_for_invalid_type(t)
         return col_type
 
-    __TYPE_SUGGESTIONS: list[tuple[type, str]] = [
+    __TYPE_SUGGESTIONS: ClassVar[list[tuple[type, str]]] = [
         (str, 'pxt.String'),
         (bool, 'pxt.Bool'),
         (int, 'pxt.Int'),
@@ -405,9 +407,8 @@ class ColumnType:
             path = parse_local_file_path(val)
             if path is not None and not path.is_file():
                 raise TypeError(f'File not found: {path}')
-        else:
-            if not isinstance(val, bytes):
-                raise TypeError(f'expected file path or bytes, got {type(val)}')
+        elif not isinstance(val, bytes):
+            raise TypeError(f'expected file path or bytes, got {type(val)}')
 
     @abc.abstractmethod
     def _validate_literal(self, val: Any) -> None:
@@ -475,12 +476,12 @@ class ColumnType:
         # types that refer to external media files
         return self.is_image_type() or self.is_video_type() or self.is_audio_type() or self.is_document_type()
 
+    @classmethod
     @abc.abstractmethod
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         """
         Return corresponding SQLAlchemy type.
         """
-        pass
 
     def to_json_schema(self) -> dict[str, Any]:
         if self.nullable:
@@ -496,14 +497,15 @@ class InvalidType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.INVALID, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.types.NullType()
 
     def print_value(self, val: Any) -> str:
         return str(val)
 
     def _validate_literal(self, val: Any) -> None:
-        assert False
+        raise AssertionError()
 
 
 class StringType(ColumnType):
@@ -513,7 +515,8 @@ class StringType(ColumnType):
     def has_supertype(self):
         return not self.nullable
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.String()
 
     def _to_json_schema(self) -> dict[str, Any]:
@@ -539,7 +542,8 @@ class IntType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.INT, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.BigInteger()
 
     def _to_json_schema(self) -> dict[str, Any]:
@@ -556,7 +560,8 @@ class FloatType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.FLOAT, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.Float()
 
     def _to_json_schema(self) -> dict[str, Any]:
@@ -576,7 +581,8 @@ class BoolType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.BOOL, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.Boolean()
 
     def _to_json_schema(self) -> dict[str, Any]:
@@ -599,7 +605,8 @@ class TimestampType(ColumnType):
     def has_supertype(self):
         return not self.nullable
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.TIMESTAMP(timezone=True)
 
     def _validate_literal(self, val: Any) -> None:
@@ -644,7 +651,8 @@ class JsonType(ColumnType):
     def _from_dict(cls, d: dict) -> ColumnType:
         return cls(json_schema=d.get('json_schema'), nullable=d['nullable'])
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.dialects.postgresql.JSONB()
 
     def _to_json_schema(self) -> dict[str, Any]:
@@ -760,7 +768,7 @@ class JsonType(ColumnType):
         a_type = a.get('type')
         b_type = b.get('type')
 
-        if a_type in ('string', 'integer', 'number', 'boolean', 'object', 'array') and a_type == b_type:
+        if a_type in {'string', 'integer', 'number', 'boolean', 'object', 'array'} and a_type == b_type:
             # a and b both have the same type designation, but are not identical. This can happen if
             # (for example) they have validators or other attributes that differ. In this case, we
             # generalize to {'type': t}, where t is their shared type, with no other qualifications.
@@ -904,7 +912,7 @@ class ArrayType(ColumnType):
         # determine our dtype
         assert isinstance(val, np.ndarray)
         pxttype: Optional[ColumnType] = cls.from_np_dtype(val.dtype, nullable)
-        if pxttype == None:
+        if pxttype is None:
             return None
         return cls(val.shape, dtype=pxttype, nullable=nullable)
 
@@ -962,7 +970,8 @@ class ArrayType(ColumnType):
             return np.array(val, dtype=self.numpy_dtype())
         return val
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.LargeBinary()
 
     def numpy_dtype(self) -> Optional[np.dtype]:
@@ -976,7 +985,7 @@ class ArrayType(ColumnType):
             return np.dtype(np.bool_)
         if self.dtype == self.Type.STRING:
             return np.dtype(np.str_)
-        assert False, self.dtype
+        raise AssertionError(self.dtype)
 
 
 class ImageType(ColumnType):
@@ -1060,7 +1069,8 @@ class ImageType(ColumnType):
         assert 'mode' in d
         return cls(width=d['width'], height=d['height'], mode=d['mode'], nullable=d['nullable'])
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.String()
 
     def _create_literal(self, val: Any) -> Any:
@@ -1094,7 +1104,8 @@ class VideoType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.VIDEO, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         # stored as a file path
         return sql.String()
 
@@ -1126,7 +1137,8 @@ class AudioType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.AUDIO, nullable=nullable)
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         # stored as a file path
         return sql.String()
 
@@ -1168,7 +1180,7 @@ class DocumentType(ColumnType):
                     raise ValueError(f'Invalid document type: {type_str}')
             self._doc_formats = [self.DocumentFormat[type_str.upper()] for type_str in type_strs]
         else:
-            self._doc_formats = [t for t in self.DocumentFormat]
+            self._doc_formats = list(self.DocumentFormat)
 
     def copy(self, nullable: bool) -> ColumnType:
         return DocumentType(doc_formats=self.doc_formats, nullable=nullable)
@@ -1179,7 +1191,8 @@ class DocumentType(ColumnType):
     def __hash__(self) -> int:
         return hash((self._type, self.nullable, self._doc_formats))
 
-    def to_sa_type(self) -> sql.types.TypeEngine:
+    @classmethod
+    def to_sa_type(cls) -> sql.types.TypeEngine:
         # stored as a file path
         return sql.String()
 

pixeltable/utils/filecache.py

@@ -13,6 +13,7 @@ from typing import Optional
 from uuid import UUID
 
 import pixeltable.exceptions as excs
+from pixeltable.config import Config
 from pixeltable.env import Env
 
 _logger = logging.getLogger('pixeltable')
@@ -153,7 +154,7 @@ class FileCache:
                 f'of the evicted file(s) is {round(extra_capacity_needed / (1 << 30), 1)} GiB.\n'
                 f'Consider increasing the cache size to at least {round(suggested_cache_size / (1 << 30), 1)} GiB '
                 f'(it is currently {round(self.capacity_bytes / (1 << 30), 1)} GiB).\n'
-                f'You can do this by setting the value of `file_cache_size_g` in: {str(Env.get()._config_file)}',
+                f'You can do this by setting the value of `file_cache_size_g` in: {Config.get().config_file}',
                 excs.PixeltableWarning,
             )
             self.new_redownload_witnessed = False

{pixeltable-0.3.5.dist-info → pixeltable-0.3.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pixeltable
-Version: 0.3.5
+Version: 0.3.7
 Summary: AI Data Infrastructure: Declarative, Multimodal, and Incremental
 License: Apache-2.0
 Keywords: data-science,machine-learning,database,ai,computer-vision,chatbot,ml,artificial-intelligence,feature-engineering,multimodal,mlops,feature-store,vector-database,llm,genai