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

pixeltable/store.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 import abc
 import logging
 import sys
+import time
 import warnings
-from typing import Any, Iterable, Iterator, Optional
+from typing import Any, Iterable, Iterator
 
 import more_itertools
 import psycopg
@@ -33,11 +34,11 @@ class StoreBase:
 
     tbl_version: catalog.TableVersionHandle
     sa_md: sql.MetaData
-    sa_tbl: Optional[sql.Table]
+    sa_tbl: sql.Table | None
     _pk_cols: list[sql.Column]
     v_min_col: sql.Column
     v_max_col: sql.Column
-    base: Optional[StoreBase]
+    base: StoreBase | None
 
     # In my cursory experiments this was the optimal batch size: it was an improvement over 5_000 and there was no real
     # benefit to going higher.
@@ -79,12 +80,13 @@ class StoreBase:
         self._pk_cols = [*rowid_cols, self.v_min_col]
         return [*rowid_cols, self.v_min_col, self.v_max_col]
 
-    def create_sa_tbl(self, tbl_version: Optional[catalog.TableVersion] = None) -> None:
+    def create_sa_tbl(self, tbl_version: catalog.TableVersion | None = None) -> None:
         """Create self.sa_tbl from self.tbl_version."""
         if tbl_version is None:
             tbl_version = self.tbl_version.get()
         system_cols = self._create_system_columns()
         all_cols = system_cols.copy()
+        # we captured all columns, including dropped ones: they're still part of the physical table
         for col in [c for c in tbl_version.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
@@ -111,7 +113,10 @@ class StoreBase:
         idx_name = f'vmax_idx_{tbl_version.id.hex}'
         idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using=Env.get().dbms.version_index_type))
 
-        # TODO: Include indices to ensure a completely accurate SA table definition?
+        # we only capture indices visible in this version
+        for idx_info in tbl_version.idxs.values():
+            idx = idx_info.idx.sa_index(tbl_version._store_idx_name(idx_info.id), idx_info.val_col)
+            idxs.append(idx)
 
         self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
         # _logger.debug(f'created sa tbl for {tbl_version.id!s} (sa_tbl={id(self.sa_tbl):x}, tv={id(tbl_version):x})')
@@ -137,35 +142,122 @@ class StoreBase:
         assert isinstance(result, int)
         return result
 
+    def _exec_if_not_exists(self, stmt: str, wait_for_table: bool) -> None:
+        """
+        Execute a statement containing 'IF NOT EXISTS' and ignore any duplicate object-related errors.
+
+        The statement needs to run in a separate transaction, because the expected error conditions will abort the
+        enclosing transaction (and the ability to run additional statements in that same transaction).
+        """
+        while True:
+            with Env.get().begin_xact(for_write=True) as conn:
+                try:
+                    if wait_for_table:
+                        # Try to lock the table to make sure that it exists. This needs to run in the same transaction
+                        # as 'stmt' to avoid a race condition.
+                        # TODO: adapt this for CockroachDB
+                        lock_stmt = f'LOCK TABLE {self._storage_name()} IN ACCESS EXCLUSIVE MODE'
+                        conn.execute(sql.text(lock_stmt))
+                    conn.execute(sql.text(stmt))
+                    return
+                except (sql.exc.IntegrityError, sql.exc.ProgrammingError) as e:
+                    Env.get().console_logger.info(f'{stmt} failed with: {e}')
+                    if (
+                        isinstance(e.orig, psycopg.errors.UniqueViolation)
+                        and 'duplicate key value violates unique constraint' in str(e.orig)
+                    ) or (
+                        isinstance(e.orig, (psycopg.errors.DuplicateObject, psycopg.errors.DuplicateTable))
+                        and 'already exists' in str(e.orig)
+                    ):
+                        # table already exists
+                        return
+                    elif isinstance(e.orig, psycopg.errors.UndefinedTable):
+                        # the Lock Table failed because the table doesn't exist yet; try again
+                        time.sleep(1)
+                        continue
+                    else:
+                        raise
+
+    def _store_tbl_exists(self) -> bool:
+        """Returns True if the store table exists, False otherwise."""
+        with Env.get().begin_xact(for_write=False) as conn:
+            q = (
+                'SELECT COUNT(*) FROM pg_catalog.pg_tables '
+                f"WHERE schemaname = 'public' AND tablename = {self._storage_name()!r}"
+            )
+            res = conn.execute(sql.text(q)).scalar_one()
+            return res == 1
+
     def create(self) -> None:
-        """Create If Not Exists for this table"""
+        """
+        Create or update store table to bring it in sync with self.sa_tbl. Idempotent.
+
+        This runs a sequence of DDL statements (Create Table, Alter Table Add Column, Create Index), each of which
+        is run in its own transaction.
+
+        The exception to that are local replicas, for which TableRestorer creates an enclosing transaction. In theory,
+        this should avoid the potential for race conditions that motivate the error handling present in
+        _exec_if_not_exists() (meaning: we shouldn't see those errors when creating local replicas).
+        TODO: remove the special case for local replicas in order to make the logic easier to reason about.
+        """
+        postgres_dialect = sql.dialects.postgresql.dialect()
+
+        if not self._store_tbl_exists():
+            # run Create Table If Not Exists; we always need If Not Exists to avoid race conditions between concurrent
+            # Pixeltable processes
+            create_stmt = sql.schema.CreateTable(self.sa_tbl, if_not_exists=True).compile(dialect=postgres_dialect)
+            self._exec_if_not_exists(str(create_stmt), wait_for_table=False)
+        else:
+            # ensure that all columns exist by running Alter Table Add Column If Not Exists for all columns
+            for col in self.sa_tbl.columns:
+                stmt = self._add_column_stmt(col)
+                self._exec_if_not_exists(stmt, wait_for_table=True)
+            # TODO: do we also need to ensure that these columns are now visible (ie, is there another potential race
+            # condition here?)
+
+        # ensure that all visible indices exist by running Create Index If Not Exists
+        for index in self.sa_tbl.indexes:
+            create_stmt = sql.schema.CreateIndex(index, if_not_exists=True).compile(dialect=postgres_dialect)
+            self._exec_if_not_exists(str(create_stmt), wait_for_table=True)
+
+    def create_index(self, idx_id: int) -> None:
+        """Create If Not Exists for this index"""
+        idx_info = self.tbl_version.get().idxs[idx_id]
+        sa_idx = idx_info.idx.sa_index(self.tbl_version.get()._store_idx_name(idx_id), idx_info.val_col)
         conn = Env.get().conn
-        stmt = sql.schema.CreateTable(self.sa_tbl).compile(conn)
+        stmt = sql.schema.CreateIndex(sa_idx, if_not_exists=True).compile(conn)
         create_stmt = str(stmt)
-        if_not_exists_stmt = create_stmt.replace('CREATE TABLE', 'CREATE TABLE IF NOT EXISTS')
-
-        # Postgres seems not to handle concurrent Create Table If Not Exists correctly, we need to ignore the various
-        # errors that can occur when two connections run the same Create Table statement.
-        try:
-            conn.execute(sql.text(if_not_exists_stmt))
-        except (sql.exc.IntegrityError, sql.exc.ProgrammingError) as e:
-            Env.get().console_logger.info(f'StoreBase.create() failed with: {e}')
-            if (
-                isinstance(e.orig, psycopg.errors.UniqueViolation)
-                and 'duplicate key value violates unique constraint "pg_type_typname_nsp_index"' in str(e.orig)
-            ) or (
-                isinstance(e.orig, (psycopg.errors.DuplicateObject, psycopg.errors.DuplicateTable))
-                and 'already exists' in str(e.orig)
-            ):
-                pass
-            else:
-                raise
+        self._exec_if_not_exists(create_stmt, wait_for_table=True)
+
+    def validate(self) -> None:
+        """Validate store table against self.table_version"""
+        with Env.get().begin_xact() as conn:
+            # check that all columns are present
+            q = f'SELECT column_name FROM information_schema.columns WHERE table_name = {self._storage_name()!r}'
+            store_col_info = {row[0] for row in conn.execute(sql.text(q)).fetchall()}
+            tbl_col_info = {col.store_name() for col in self.tbl_version.get().cols if col.is_stored}
+            assert tbl_col_info.issubset(store_col_info)
+
+            # check that all visible indices are present
+            q = f'SELECT indexname FROM pg_indexes WHERE tablename = {self._storage_name()!r}'
+            store_idx_names = {row[0] for row in conn.execute(sql.text(q)).fetchall()}
+            tbl_index_names = {
+                self.tbl_version.get()._store_idx_name(info.id) for info in self.tbl_version.get().idxs.values()
+            }
+            assert tbl_index_names.issubset(store_idx_names)
 
     def drop(self) -> None:
         """Drop store table"""
         conn = Env.get().conn
         self.sa_md.drop_all(bind=conn)
 
+    def _add_column_stmt(self, sa_col: sql.Column) -> str:
+        col_type_str = sa_col.type.compile(dialect=sql.dialects.postgresql.dialect())
+        return (
+            f'ALTER TABLE {self._storage_name()} ADD COLUMN IF NOT EXISTS '
+            f'{sa_col.name} {col_type_str} {"NOT " if not sa_col.nullable else ""} NULL'
+        )
+
     def add_column(self, col: catalog.Column) -> None:
         """Add column(s) to the store-resident table based on a catalog column
 
@@ -174,7 +266,7 @@ class StoreBase:
         """
         assert col.is_stored
         conn = Env.get().conn
-        col_type_str = col.get_sa_col_type().compile(dialect=conn.dialect)
+        col_type_str = col.sa_col_type.compile(dialect=conn.dialect)
         s_txt = f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.store_name()} {col_type_str} NULL'
         added_storage_cols = [col.store_name()]
         if col.stores_cellmd:
@@ -197,34 +289,6 @@ class StoreBase:
         log_stmt(_logger, stmt)
         Env.get().conn.execute(stmt)
 
-    def ensure_updated_schema(self) -> None:
-        from pixeltable.utils.dbms import PostgresqlDbms
-
-        # This should only be called during replica creation where the underlying DBMS is Postgres.
-        assert isinstance(Env.get().dbms, PostgresqlDbms)
-
-        conn = Env.get().conn
-        tv = self.tbl_version.get()
-
-        # Ensure columns exist
-        sql_text = f'SELECT column_name FROM information_schema.columns WHERE table_name = {self._storage_name()!r}'
-        result = conn.execute(sql.text(sql_text))
-        existing_cols = {row[0] for row in result}
-        for col in tv.cols:
-            if col.is_stored and col.store_name() not in existing_cols:
-                _logger.debug(f'Adding missing column {col.store_name()!r} to store table {self._storage_name()!r}')
-                self.add_column(col)
-
-        # Ensure indices exist
-        sql_text = f'SELECT indexname FROM pg_indexes WHERE tablename = {self._storage_name()!r}'
-        result = conn.execute(sql.text(sql_text))
-        existing_idxs = {row[0] for row in result}
-        for idx_name, idx_info in tv.all_idxs.items():
-            store_name = tv._store_idx_name(idx_info.id)
-            if store_name not in existing_idxs:
-                _logger.debug(f'Creating missing index {idx_name!r} on store table {self._storage_name()!r}')
-                idx_info.idx.create_index(store_name, idx_info.val_col)
-
     def load_column(self, col: catalog.Column, exec_plan: ExecNode, abort_on_exc: bool) -> int:
         """Update store column of a computed column with values produced by an execution plan
 
@@ -234,7 +298,7 @@ class StoreBase:
             sql.exc.DBAPIError if there was a SQL error during execution
             excs.Error if on_error='abort' and there was an exception during row evaluation
         """
-        assert col.tbl.id == self.tbl_version.id
+        assert col.get_tbl().id == self.tbl_version.id
         num_excs = 0
         num_rows = 0
         # create temp table to store output of exec_plan, with the same primary key as the store table
@@ -304,7 +368,7 @@ class StoreBase:
         exec_plan: ExecNode,
         v_min: int,
         show_progress: bool = True,
-        rowids: Optional[Iterator[int]] = None,
+        rowids: Iterator[int] | None = None,
         abort_on_exc: bool = False,
     ) -> tuple[set[int], RowCountStats]:
         """Insert rows into the store table and update the catalog table's md
@@ -316,7 +380,7 @@ class StoreBase:
         num_excs = 0
         num_rows = 0
         cols_with_excs: set[int] = set()
-        progress_bar: Optional[tqdm] = None  # create this only after we started executing
+        progress_bar: tqdm | None = None  # create this only after we started executing
         row_builder = exec_plan.row_builder
 
         store_col_names = row_builder.store_column_names()
@@ -389,7 +453,7 @@ class StoreBase:
         # stmt_text = f'INSERT INTO {self.sa_tbl.name} ({col_names_str}) VALUES ({placeholders_str})'
         # conn.exec_driver_sql(stmt_text, table_rows)
 
-    def _versions_clause(self, versions: list[Optional[int]], match_on_vmin: bool) -> sql.ColumnElement[bool]:
+    def _versions_clause(self, versions: list[int | None], match_on_vmin: bool) -> sql.ColumnElement[bool]:
         """Return filter for base versions"""
         v = versions[0]
         if v is None:
@@ -407,9 +471,9 @@ class StoreBase:
     def delete_rows(
         self,
         current_version: int,
-        base_versions: list[Optional[int]],
+        base_versions: list[int | None],
         match_on_vmin: bool,
-        where_clause: Optional[sql.ColumnElement[bool]],
+        where_clause: sql.ColumnElement[bool] | None,
     ) -> int:
         """Mark rows as deleted that are live and were created prior to current_version.
         Also: populate the undo columns
@@ -535,7 +599,7 @@ class StoreComponentView(StoreView):
         self.rowid_cols.append(self.pos_col)
         return self.rowid_cols
 
-    def create_sa_tbl(self, tbl_version: Optional[catalog.TableVersion] = None) -> None:
+    def create_sa_tbl(self, tbl_version: catalog.TableVersion | None = None) -> None:
         if tbl_version is None:
             tbl_version = self.tbl_version.get()
         super().create_sa_tbl(tbl_version)

pixeltable/type_system.py

@@ -10,7 +10,7 @@ import typing
 import urllib.parse
 import urllib.request
 from pathlib import Path
-from typing import Any, ClassVar, Iterable, Literal, Mapping, Optional, Sequence, Union
+from typing import Any, ClassVar, Iterable, Literal, Mapping, Sequence, Union
 
 from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
 
@@ -51,11 +51,11 @@ class ColumnType:
         @classmethod
         def supertype(
             cls,
-            type1: Optional['ColumnType.Type'],
-            type2: Optional['ColumnType.Type'],
+            type1: 'ColumnType.Type' | None,
+            type2: 'ColumnType.Type' | None,
             # we need to pass this in because we can't easily append it as a class member
             common_supertypes: dict[tuple['ColumnType.Type', 'ColumnType.Type'], 'ColumnType.Type'],
-        ) -> Optional['ColumnType.Type']:
+        ) -> 'ColumnType.Type' | None:
             if type1 == type2:
                 return type1
             t = common_supertypes.get((type1, type2))
@@ -188,7 +188,7 @@ class ColumnType:
         if as_schema:
             return base_str if self.nullable else f'Required[{base_str}]'
         else:
-            return f'Optional[{base_str}]' if self.nullable else base_str
+            return f'{base_str} | None' if self.nullable else base_str
 
     def _to_base_str(self) -> str:
         """
@@ -217,7 +217,7 @@ class ColumnType:
         # Default: just compare base types (this works for all types whose only parameter is nullable)
         return self._type == other._type
 
-    def supertype(self, other: ColumnType) -> Optional[ColumnType]:
+    def supertype(self, other: ColumnType) -> ColumnType | None:
         if self == other:
             return self
         if self.matches(other):
@@ -237,7 +237,7 @@ class ColumnType:
         return None
 
     @classmethod
-    def infer_literal_type(cls, val: Any, nullable: bool = False) -> Optional[ColumnType]:
+    def infer_literal_type(cls, val: Any, nullable: bool = False) -> ColumnType | None:
         if val is None:
             return InvalidType(nullable=True)
         if isinstance(val, str):
@@ -271,7 +271,7 @@ class ColumnType:
         return None
 
     @classmethod
-    def infer_common_literal_type(cls, vals: Iterable[Any]) -> Optional[ColumnType]:
+    def infer_common_literal_type(cls, vals: Iterable[Any]) -> ColumnType | None:
         """
         Returns the most specific type that is a supertype of all literals in `vals`. If no such type
         exists, returns None.
@@ -279,7 +279,7 @@ class ColumnType:
         Args:
             vals: A collection of literals.
         """
-        inferred_type: Optional[ColumnType] = None
+        inferred_type: ColumnType | None = None
         for val in vals:
             val_type = cls.infer_literal_type(val)
             if inferred_type is None:
@@ -299,7 +299,7 @@ class ColumnType:
         nullable_default: bool = False,
         allow_builtin_types: bool = True,
         infer_pydantic_json: bool = False,
-    ) -> Optional[ColumnType]:
+    ) -> ColumnType | None:
         """
         Convert a Python type into a Pixeltable `ColumnType` instance.
 
@@ -317,9 +317,9 @@ class ColumnType:
         origin = typing.get_origin(t)
         type_args = typing.get_args(t)
         if origin in (typing.Union, types.UnionType):
-            # Check if `t` has the form Optional[T].
+            # Check if `t` has the form T | None.
             if len(type_args) == 2 and type(None) in type_args:
-                # `t` is a type of the form Optional[T] (equivalently, T | None or None | T).
+                # `t` is a type of the form T | None (equivalently, T | None or None | T).
                 # We treat it as the underlying type but with nullable=True.
                 underlying_py_type = type_args[0] if type_args[1] is type(None) else type_args[1]
                 underlying = cls.from_python_type(
@@ -338,7 +338,7 @@ class ColumnType:
             if isinstance(parameters, ColumnType):
                 return parameters.copy(nullable=nullable_default)
         else:
-            # It's something other than Optional[T], Required[T], or an explicitly annotated type.
+            # It's something other than T | None, Required[T], or an explicitly annotated type.
             if origin is not None:
                 # Discard type parameters to ensure that parameterized types such as `list[T]`
                 # are correctly mapped to Pixeltable types.
@@ -411,7 +411,7 @@ class ColumnType:
         raise excs.Error(f'Unknown type: {t}')
 
     @classmethod
-    def from_json_schema(cls, schema: dict[str, Any]) -> Optional[ColumnType]:
+    def from_json_schema(cls, schema: dict[str, Any]) -> ColumnType | None:
         # We first express the JSON schema as a Python type, and then convert it to a Pixeltable type.
         # TODO: Is there a meaningful fallback if one of these operations fails? (Maybe another use case for a pxt Any
         #     type?)
@@ -704,10 +704,10 @@ class DateType(ColumnType):
 
 
 class JsonType(ColumnType):
-    json_schema: Optional[dict[str, Any]]
-    __validator: Optional[jsonschema.protocols.Validator]
+    json_schema: dict[str, Any] | None
+    __validator: jsonschema.protocols.Validator | None
 
-    def __init__(self, json_schema: Optional[dict[str, Any]] = None, nullable: bool = False):
+    def __init__(self, json_schema: dict[str, Any] | None = None, nullable: bool = False):
         super().__init__(self.Type.JSON, nullable=nullable)
         self.json_schema = json_schema
         if json_schema is None:
@@ -777,7 +777,7 @@ class JsonType(ColumnType):
             return val.model_dump()
         return val
 
-    def supertype(self, other: ColumnType) -> Optional[JsonType]:
+    def supertype(self, other: ColumnType) -> JsonType | None:
         # Try using the (much faster) supertype logic in ColumnType first. That will work if, for example, the types
         # are identical except for nullability. If that doesn't work and both types are JsonType, then we will need to
         # merge their schemas.
@@ -799,7 +799,7 @@ class JsonType(ColumnType):
         )
 
     @classmethod
-    def __superschema(cls, a: dict[str, Any], b: dict[str, Any]) -> Optional[dict[str, Any]]:
+    def __superschema(cls, a: dict[str, Any], b: dict[str, Any]) -> dict[str, Any] | None:
         # Defining a general type hierarchy over all JSON schemas would be a challenging problem. In order to keep
         # things manageable, we only define a hierarchy among "conforming" schemas, which provides enough generality
         # for the most important use cases (unions for type inference, validation of inline exprs). A schema is
@@ -859,7 +859,7 @@ class JsonType(ColumnType):
         return {}  # Unresolvable type conflict; the supertype is an unrestricted JsonType.
 
     @classmethod
-    def __superschema_with_nulls(cls, a: dict[str, Any], b: dict[str, Any]) -> Optional[dict[str, Any]]:
+    def __superschema_with_nulls(cls, a: dict[str, Any], b: dict[str, Any]) -> dict[str, Any] | None:
         a, a_nullable = cls.__unpack_null_from_schema(a)
         b, b_nullable = cls.__unpack_null_from_schema(b)
 
@@ -888,15 +888,12 @@ class JsonType(ColumnType):
 
 
 class ArrayType(ColumnType):
-    shape: Optional[tuple[Optional[int], ...]]
-    pxt_dtype: Optional[ColumnType]
-    dtype: Optional[ColumnType.Type]
+    shape: tuple[int | None, ...] | None
+    pxt_dtype: ColumnType | None
+    dtype: ColumnType.Type | None
 
     def __init__(
-        self,
-        shape: Optional[tuple[Optional[int], ...]] = None,
-        dtype: Optional[ColumnType] = None,
-        nullable: bool = False,
+        self, shape: tuple[int | None, ...] | None = None, dtype: ColumnType | None = None, nullable: bool = False
     ):
         super().__init__(self.Type.ARRAY, nullable=nullable)
         assert shape is None or dtype is not None, (shape, dtype)  # cannot specify a shape without a dtype
@@ -921,7 +918,7 @@ class ArrayType(ColumnType):
     def __hash__(self) -> int:
         return hash((self._type, self.nullable, self.shape, self.dtype))
 
-    def supertype(self, other: ColumnType) -> Optional[ArrayType]:
+    def supertype(self, other: ColumnType) -> ArrayType | None:
         basic_supertype = super().supertype(other)
         if basic_supertype is not None:
             assert isinstance(basic_supertype, ArrayType)
@@ -934,7 +931,7 @@ class ArrayType(ColumnType):
         if super_dtype is None:
             # if the dtypes are incompatible, then the supertype is a fully general array
             return ArrayType(nullable=(self.nullable or other.nullable))
-        super_shape: Optional[tuple[Optional[int], ...]]
+        super_shape: tuple[int | None, ...] | None
         if self.shape is None or other.shape is None or len(self.shape) != len(other.shape):
             super_shape = None
         else:
@@ -965,7 +962,7 @@ class ArrayType(ColumnType):
         return cls(shape, dtype, nullable=d['nullable'])
 
     @classmethod
-    def from_np_dtype(cls, dtype: np.dtype, nullable: bool) -> Optional[ColumnType]:
+    def from_np_dtype(cls, dtype: np.dtype, nullable: bool) -> ColumnType | None:
         """
         Return pixeltable type corresponding to a given simple numpy dtype
         """
@@ -994,10 +991,10 @@ class ArrayType(ColumnType):
         return None
 
     @classmethod
-    def from_literal(cls, val: np.ndarray, nullable: bool = False) -> Optional[ArrayType]:
+    def from_literal(cls, val: np.ndarray, nullable: bool = False) -> ArrayType | None:
         # determine our dtype
         assert isinstance(val, np.ndarray)
-        pxttype: Optional[ColumnType] = cls.from_np_dtype(val.dtype, nullable)
+        pxttype: ColumnType | None = cls.from_np_dtype(val.dtype, nullable)
         if pxttype is None:
             return None
         return cls(val.shape, dtype=pxttype, nullable=nullable)
@@ -1060,7 +1057,7 @@ class ArrayType(ColumnType):
     def to_sa_type(cls) -> sql.types.TypeEngine:
         return sql.LargeBinary()
 
-    def numpy_dtype(self) -> Optional[np.dtype]:
+    def numpy_dtype(self) -> np.dtype | None:
         if self.dtype is None:
             return None
         if self.dtype == self.Type.INT:
@@ -1077,10 +1074,10 @@ class ArrayType(ColumnType):
 class ImageType(ColumnType):
     def __init__(
         self,
-        width: Optional[int] = None,
-        height: Optional[int] = None,
-        size: Optional[tuple[int, int]] = None,
-        mode: Optional[str] = None,
+        width: int | None = None,
+        height: int | None = None,
+        size: tuple[int, int] | None = None,
+        mode: str | None = None,
         nullable: bool = False,
     ):
         # TODO: does it make sense to specify only width or height?
@@ -1121,7 +1118,7 @@ class ImageType(ColumnType):
     def __hash__(self) -> int:
         return hash((self._type, self.nullable, self.size, self.mode))
 
-    def supertype(self, other: ColumnType) -> Optional[ImageType]:
+    def supertype(self, other: ColumnType) -> ImageType | None:
         basic_supertype = super().supertype(other)
         if basic_supertype is not None:
             assert isinstance(basic_supertype, ImageType)
@@ -1136,7 +1133,7 @@ class ImageType(ColumnType):
         return ImageType(width=width, height=height, mode=mode, nullable=(self.nullable or other.nullable))
 
     @property
-    def size(self) -> Optional[tuple[int, int]]:
+    def size(self) -> tuple[int, int] | None:
         if self.width is None or self.height is None:
             return None
         return (self.width, self.height)
@@ -1255,7 +1252,7 @@ class DocumentType(ColumnType):
         TXT = 4
 
         @classmethod
-        def from_extension(cls, ext: str) -> Optional['DocumentType.DocumentFormat']:
+        def from_extension(cls, ext: str) -> 'DocumentType.DocumentFormat' | None:
             if ext in ('.htm', '.html'):
                 return cls.HTML
             if ext == '.md':
@@ -1268,7 +1265,7 @@ class DocumentType(ColumnType):
                 return cls.TXT
             return None
 
-    def __init__(self, nullable: bool = False, doc_formats: Optional[str] = None):
+    def __init__(self, nullable: bool = False, doc_formats: str | None = None):
         super().__init__(self.Type.DOCUMENT, nullable=nullable)
         self.doc_formats = doc_formats
         if doc_formats is not None:
@@ -1365,13 +1362,13 @@ class Array(np.ndarray, _PxtType):
     def __class_getitem__(cls, item: Any) -> _AnnotatedAlias:
         """
         `item` (the type subscript) must be a tuple with exactly two elements (in any order):
-        - A tuple of `Optional[int]`s, specifying the shape of the array
+        - A tuple of `int | None`s, specifying the shape of the array
         - A type, specifying the dtype of the array
         Example: Array[(3, None, 2), pxt.Float]
         """
         params = item if isinstance(item, tuple) else (item,)
-        shape: Optional[tuple] = None
-        dtype: Optional[ColumnType] = None
+        shape: tuple | None = None
+        dtype: ColumnType | None = None
         if not any(isinstance(param, (type, _AnnotatedAlias)) for param in params):
             raise TypeError('Array type parameter must include a dtype.')
         for param in params:
@@ -1411,8 +1408,8 @@ class Image(PIL.Image.Image, _PxtType):
         else:
             # Not a tuple (single arg)
             params = (item,)
-        size: Optional[tuple] = None
-        mode: Optional[str] = None
+        size: tuple | None = None
+        mode: str | None = None
         for param in params:
             if isinstance(param, tuple):
                 if (

pixeltable/utils/__init__.py

@@ -2,7 +2,6 @@ import hashlib
 import urllib.parse
 import urllib.request
 from pathlib import Path
-from typing import Optional
 
 
 def print_perf_counter_delta(delta: float) -> str:
@@ -39,7 +38,7 @@ def sha256sum(path: Path | str) -> str:
     return h.hexdigest()
 
 
-def parse_local_file_path(file_or_url: str) -> Optional[Path]:
+def parse_local_file_path(file_or_url: str) -> Path | None:
     """
     Parses a string that may be either a URL or a local file path.
 

pixeltable/utils/arrow.py

@@ -1,7 +1,7 @@
 import datetime
 import io
 import json
-from typing import TYPE_CHECKING, Any, Iterator, Optional, cast
+from typing import TYPE_CHECKING, Any, Iterator, cast
 
 import numpy as np
 import PIL.Image
@@ -48,7 +48,7 @@ PXT_TO_PA_TYPES: dict[type[ts.ColumnType], pa.DataType] = {
 }
 
 
-def to_pixeltable_type(arrow_type: pa.DataType, nullable: bool) -> Optional[ts.ColumnType]:
+def to_pixeltable_type(arrow_type: pa.DataType, nullable: bool) -> ts.ColumnType | None:
     """Convert a pyarrow DataType to a pixeltable ColumnType if one is defined.
     Returns None if no conversion is currently implemented.
     """
@@ -66,7 +66,7 @@ def to_pixeltable_type(arrow_type: pa.DataType, nullable: bool) -> Optional[ts.C
         return None
 
 
-def to_arrow_type(pixeltable_type: ts.ColumnType) -> Optional[pa.DataType]:
+def to_arrow_type(pixeltable_type: ts.ColumnType) -> pa.DataType | None:
     """Convert a pixeltable DataType to a pyarrow datatype if one is defined.
     Returns None if no conversion is currently implemented.
     """
@@ -240,7 +240,7 @@ def _ar_val_to_pxt_val(val: Any, pxt_type: ts.ColumnType) -> Any:
 
 
 def iter_tuples2(
-    batch: pa.Table | pa.RecordBatch, col_mapping: Optional[dict[str, str]], schema: dict[str, ts.ColumnType]
+    batch: pa.Table | pa.RecordBatch, col_mapping: dict[str, str] | None, schema: dict[str, ts.ColumnType]
 ) -> Iterator[dict[str, Any]]:
     """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
     pydict = to_pydict(batch)

pixeltable/utils/av.py

@@ -5,6 +5,14 @@ import av.stream
 
 from pixeltable.env import Env
 
+# format -> (codec, extension)
+AUDIO_FORMATS: dict[str, tuple[str, str]] = {
+    'wav': ('pcm_s16le', 'wav'),
+    'mp3': ('libmp3lame', 'mp3'),
+    'flac': ('flac', 'flac'),
+    'mp4': ('aac', 'm4a'),
+}
+
 
 def get_metadata(path: str) -> dict:
     with av.open(path) as container: