pixeltable 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

pixeltable/catalog/table_version.py

@@ -14,7 +14,6 @@ import sqlalchemy as sql
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
 from pixeltable import exprs, index
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
@@ -223,18 +222,23 @@ class TableVersion:
         view_md: Optional[schema.ViewMd] = None,
     ) -> TableVersionMd:
         user = Env.get().user
+        timestamp = time.time()
 
-        # assign ids
+        # assign ids, create metadata
         cols_by_name: dict[str, Column] = {}
+        column_md: dict[int, schema.ColumnMd] = {}
+        schema_col_md: dict[int, schema.SchemaColumn] = {}
         for pos, col in enumerate(cols):
             col.id = pos
             col.schema_version_add = 0
             cols_by_name[col.name] = col
             if col.is_computed:
                 col.check_value_expr()
+            col_md, sch_md = col.to_md(pos)
+            assert sch_md is not None
+            column_md[col.id] = col_md
+            schema_col_md[col.id] = sch_md
 
-        timestamp = time.time()
-        column_md = cls._create_column_md(cols)
         tbl_id = uuid.uuid4()
         tbl_id_str = str(tbl_id)
         tbl_md = schema.TableMd(
@@ -256,18 +260,15 @@ class TableVersion:
         )
 
         table_version_md = schema.TableVersionMd(
-            tbl_id=tbl_id_str, created_at=timestamp, version=0, schema_version=0, additional_md={}
+            tbl_id=tbl_id_str,
+            created_at=timestamp,
+            version=0,
+            schema_version=0,
+            user=user,
+            update_status=None,
+            additional_md={},
         )
 
-        schema_col_md: dict[int, schema.SchemaColumn] = {}
-        for pos, col in enumerate(cols):
-            md = schema.SchemaColumn(
-                pos=pos,
-                name=col.name,
-                media_validation=col._media_validation.name.lower() if col._media_validation is not None else None,
-            )
-            schema_col_md[col.id] = md
-
         schema_version_md = schema.TableSchemaVersionMd(
             tbl_id=tbl_id_str,
             schema_version=0,
@@ -290,76 +291,11 @@ class TableVersion:
         comment: str,
         media_validation: MediaValidation,
     ) -> tuple[UUID, Optional[TableVersion]]:
-        user = Env.get().user
-
-        # assign ids
-        cols_by_name: dict[str, Column] = {}
-        for pos, col in enumerate(cols):
-            col.id = pos
-            col.schema_version_add = 0
-            cols_by_name[col.name] = col
-            if col.is_computed:
-                col.check_value_expr()
-
-        timestamp = time.time()
-        # create schema.Table
-        # Column.dependent_cols for existing cols is wrong at this point, but init() will set it correctly
-        column_md = cls._create_column_md(cols)
-        tbl_id = uuid.uuid4()
-        tbl_id_str = str(tbl_id)
-        table_md = schema.TableMd(
-            tbl_id=tbl_id_str,
-            name=name,
-            user=user,
-            is_replica=False,
-            current_version=0,
-            current_schema_version=0,
-            next_col_id=len(cols),
-            next_idx_id=0,
-            next_row_id=0,
-            view_sn=0,
-            column_md=column_md,
-            index_md={},
-            external_stores=[],
-            view_md=None,
-            additional_md={},
-        )
-
-        # create schema.TableVersion of the initial version
-        table_version_md = schema.TableVersionMd(
-            tbl_id=tbl_id_str,
-            created_at=timestamp,
-            version=0,
-            schema_version=0,
-            user=user,
-            update_status=None,
-            additional_md={},
-        )
-
-        # create schema.TableSchemaVersion
-        schema_col_md: dict[int, schema.SchemaColumn] = {}
-        for pos, col in enumerate(cols):
-            md = schema.SchemaColumn(
-                pos=pos,
-                name=col.name,
-                media_validation=col._media_validation.name.lower() if col._media_validation is not None else None,
-            )
-            schema_col_md[col.id] = md
-
-        schema_version_md = schema.TableSchemaVersionMd(
-            tbl_id=tbl_id_str,
-            schema_version=0,
-            preceding_schema_version=None,
-            columns=schema_col_md,
-            num_retained_versions=num_retained_versions,
-            comment=comment,
-            media_validation=media_validation.name.lower(),
-            additional_md={},
-        )
-
+        inital_md = cls.create_initial_md(name, cols, num_retained_versions, comment, media_validation, view_md=None)
         cat = pxt.catalog.Catalog.get()
 
-        tbl_version = cls(tbl_id, table_md, table_version_md, None, schema_version_md, [])
+        tbl_id = UUID(hex=inital_md.tbl_md.tbl_id)
+        tbl_version = cls(tbl_id, inital_md.tbl_md, inital_md.version_md, None, inital_md.schema_version_md, [])
         # TODO: break this up, so that Catalog.create_table() registers tbl_version
         cat._tbl_versions[tbl_id, None] = tbl_version
         tbl_version.init()
@@ -373,8 +309,8 @@ class TableVersion:
             tbl_id=tbl_id,
             dir_id=dir_id,
             tbl_md=tbl_version.tbl_md,
-            version_md=table_version_md,
-            schema_version_md=schema_version_md,
+            version_md=inital_md.version_md,
+            schema_version_md=inital_md.schema_version_md,
         )
         return tbl_id, tbl_version
 
@@ -480,25 +416,7 @@ class TableVersion:
         sorted_column_md = sorted(self.tbl_md.column_md.values(), key=lambda item: item.id)
         for col_md in sorted_column_md:
             schema_col_md = self.schema_version_md.columns.get(col_md.id)
-            col_name = schema_col_md.name if schema_col_md is not None else None
-            media_val = (
-                MediaValidation[schema_col_md.media_validation.upper()]
-                if schema_col_md is not None and schema_col_md.media_validation is not None
-                else None
-            )
-            col = Column(
-                col_id=col_md.id,
-                name=col_name,
-                col_type=ts.ColumnType.from_dict(col_md.col_type),
-                is_pk=col_md.is_pk,
-                stored=col_md.stored,
-                media_validation=media_val,
-                schema_version_add=col_md.schema_version_add,
-                schema_version_drop=col_md.schema_version_drop,
-                value_expr_dict=col_md.value_expr,
-                tbl=self,
-            )
-            col.tbl = self
+            col = Column.from_md(col_md, self, schema_col_md)
             self.cols.append(col)
 
             # populate the lookup structures before Expr.from_dict()
@@ -783,31 +701,22 @@ class TableVersion:
         num_excs = 0
         cols_with_excs: list[Column] = []
         for col in cols_to_add:
+            assert col.id is not None, 'Column id must be set before adding the column'
             excs_per_col = 0
             col.schema_version_add = self.schema_version
             # add the column to the lookup structures now, rather than after the store changes executed successfully,
             # because it might be referenced by the next column's value_expr
             self.cols.append(col)
-            if col.name is not None:
-                self.cols_by_name[col.name] = col
             self.cols_by_id[col.id] = col
-
-            # also add to stored md
-            self._tbl_md.column_md[col.id] = schema.ColumnMd(
-                id=col.id,
-                col_type=col.col_type.as_dict(),
-                is_pk=col.is_pk,
-                schema_version_add=col.schema_version_add,
-                schema_version_drop=col.schema_version_drop,
-                value_expr=col.value_expr.as_dict() if col.value_expr is not None else None,
-                stored=col.stored,
-            )
             if col.name is not None:
-                self._schema_version_md.columns[col.id] = schema.SchemaColumn(
-                    name=col.name,
-                    pos=len(self.cols_by_name),
-                    media_validation=col._media_validation.name.lower() if col._media_validation is not None else None,
-                )
+                self.cols_by_name[col.name] = col
+                col_md, sch_md = col.to_md(len(self.cols_by_name))
+                assert sch_md is not None, 'Schema column metadata must be created for user-facing columns'
+                self._tbl_md.column_md[col.id] = col_md
+                self._schema_version_md.columns[col.id] = sch_md
+            else:
+                col_md, _ = col.to_md()
+                self._tbl_md.column_md[col.id] = col_md
 
             if col.is_stored:
                 self.store_tbl.add_column(col)
@@ -1628,23 +1537,6 @@ class TableVersion:
             return 1 + self.base.get().num_rowid_columns()
         return 1
 
-    @classmethod
-    def _create_column_md(cls, cols: list[Column]) -> dict[int, schema.ColumnMd]:
-        column_md: dict[int, schema.ColumnMd] = {}
-        for col in cols:
-            value_expr_dict = col.value_expr.as_dict() if col.value_expr is not None else None
-            assert col.is_pk is not None
-            column_md[col.id] = schema.ColumnMd(
-                id=col.id,
-                col_type=col.col_type.as_dict(),
-                is_pk=col.is_pk,
-                schema_version_add=col.schema_version_add,
-                schema_version_drop=col.schema_version_drop,
-                value_expr=value_expr_dict,
-                stored=col.stored,
-            )
-        return column_md
-
     @classmethod
     def _create_stores_md(cls, stores: Iterable[pxt.io.ExternalStore]) -> list[dict[str, Any]]:
         return [

pixeltable/catalog/view.py

@@ -25,6 +25,7 @@ from .tbl_ops import CreateStoreTableOp, LoadViewOp, TableOp
 from .update_status import UpdateStatus
 
 if TYPE_CHECKING:
+    from pixeltable.catalog.table import TableMetadata
     from pixeltable.globals import TableDataSource
 
 _logger = logging.getLogger('pixeltable')
@@ -261,7 +262,7 @@ class View(Table):
         """
         return self._snapshot_only and self._id == self._tbl_version_path.tbl_id
 
-    def _get_metadata(self) -> dict[str, Any]:
+    def _get_metadata(self) -> 'TableMetadata':
         md = super()._get_metadata()
         md['is_view'] = True
         md['is_snapshot'] = self._tbl_version_path.is_snapshot()

pixeltable/dataframe.py

@@ -19,7 +19,6 @@ from typing import (
     Optional,
     Sequence,
     TypeVar,
-    Union,
 )
 
 import pandas as pd
@@ -766,7 +765,7 @@ class DataFrame:
         )
 
     def _create_join_predicate(
-        self, other: catalog.TableVersionPath, on: Union[exprs.Expr, Sequence[exprs.ColumnRef]]
+        self, other: catalog.TableVersionPath, on: exprs.Expr | Sequence[exprs.ColumnRef]
     ) -> exprs.Expr:
         """Verifies user-specified 'on' argument and converts it into a join predicate."""
         col_refs: list[exprs.ColumnRef] = []
@@ -829,7 +828,7 @@ class DataFrame:
     def join(
         self,
         other: catalog.Table,
-        on: Optional[Union[exprs.Expr, Sequence[exprs.ColumnRef]]] = None,
+        on: exprs.Expr | Sequence[exprs.ColumnRef] | None = None,
         how: plan.JoinType.LiteralType = 'inner',
     ) -> DataFrame:
         """

pixeltable/env.py

@@ -17,7 +17,6 @@ import types
 import typing
 import uuid
 import warnings
-from abc import abstractmethod
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -890,6 +889,10 @@ class RateLimitsInfo:
     get_request_resources: Callable[..., dict[str, int]]
 
     resource_limits: dict[str, RateLimitInfo] = field(default_factory=dict)
+    has_exc: bool = False
+
+    def debug_str(self) -> str:
+        return ','.join(info.debug_str() for info in self.resource_limits.values())
 
     def is_initialized(self) -> bool:
         return len(self.resource_limits) > 0
@@ -897,7 +900,7 @@ class RateLimitsInfo:
     def reset(self) -> None:
         self.resource_limits.clear()
 
-    def record(self, **kwargs: Any) -> None:
+    def record(self, reset_exc: bool = False, **kwargs: Any) -> None:
         now = datetime.datetime.now(tz=datetime.timezone.utc)
         if len(self.resource_limits) == 0:
             self.resource_limits = {k: RateLimitInfo(k, now, *v) for k, v in kwargs.items() if v is not None}
@@ -908,14 +911,30 @@ class RateLimitsInfo:
                     f'reset={info.reset_at.strftime(TIME_FORMAT)} delta={(info.reset_at - now).total_seconds()}'
                 )
         else:
+            if self.has_exc and not reset_exc:
+                # ignore updates until we're asked to reset
+                _logger.debug(f'rate_limits.record(): ignoring update {kwargs}')
+                return
+            self.has_exc = False
             for k, v in kwargs.items():
                 if v is not None:
                     self.resource_limits[k].update(now, *v)
 
-    @abstractmethod
+    def record_exc(self, exc: Exception) -> None:
+        """Update self.resource_limits based on the exception headers"""
+        self.has_exc = True
+
     def get_retry_delay(self, exc: Exception) -> Optional[float]:
         """Returns number of seconds to wait before retry, or None if not retryable"""
-        pass
+        if len(self.resource_limits) == 0:
+            return 1.0
+        # we're looking for the maximum delay across all depleted resources
+        max_delay = 0.0
+        now = datetime.datetime.now(tz=datetime.timezone.utc)
+        for limit_info in self.resource_limits.values():
+            if limit_info.remaining < 0.05 * limit_info.limit:
+                max_delay = max(max_delay, (limit_info.reset_at - now).total_seconds())
+        return max_delay if max_delay > 0 else None
 
 
 @dataclass
@@ -928,9 +947,15 @@ class RateLimitInfo:
     remaining: int
     reset_at: datetime.datetime
 
+    def debug_str(self) -> str:
+        return (
+            f'{self.resource}@{self.recorded_at.strftime(TIME_FORMAT)}: '
+            f'{self.limit}/{self.remaining}/{self.reset_at.strftime(TIME_FORMAT)}'
+        )
+
     def update(self, recorded_at: datetime.datetime, limit: int, remaining: int, reset_at: datetime.datetime) -> None:
         # we always update everything, even though responses may come back out-of-order: we can't use reset_at to
-        # determine order, because it doesn't increase monotonically (the reeset duration shortens as output_tokens
+        # determine order, because it doesn't increase monotonically (the reset duration shortens as output_tokens
         # are freed up - going from max to actual)
         self.recorded_at = recorded_at
         self.limit = limit
@@ -942,3 +967,16 @@ class RateLimitInfo:
             f'Update {self.resource} rate limit: rem={self.remaining} reset={self.reset_at.strftime(TIME_FORMAT)} '
             f'reset_delta={reset_delta.total_seconds()} recorded_delta={(self.reset_at - recorded_at).total_seconds()}'
         )
+
+
+@dataclass
+class RuntimeCtx:
+    """
+    Container for runtime data provided by the execution system to udfs.
+
+    Udfs that accept the special _runtime_ctx parameter receive an instance of this class.
+    """
+
+    # Indicates a retry attempt following a rate limit error (error code: 429). Requires a 'rate-limits' resource pool.
+    # If True, call RateLimitsInfo.record() with reset_exc=True.
+    is_retry: bool = False

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -4,7 +4,7 @@ import asyncio
 import logging
 import traceback
 from types import TracebackType
-from typing import AsyncIterator, Iterable, Optional, Union
+from typing import AsyncIterator, Iterable, Optional
 
 import numpy as np
 
@@ -49,7 +49,7 @@ class ExprEvalNode(ExecNode):
     # execution state
     tasks: set[asyncio.Task]  # collects all running tasks to prevent them from getting gc'd
     exc_event: asyncio.Event  # set if an exception needs to be propagated
-    error: Optional[Union[Exception]]  # exception that needs to be propagated
+    error: Optional[Exception]  # exception that needs to be propagated
     completed_rows: asyncio.Queue[exprs.DataRow]  # rows that have completed evaluation
     completed_event: asyncio.Event  # set when completed_rows is non-empty
     input_iter: AsyncIterator[DataRowBatch]

pixeltable/exec/expr_eval/schedulers.py

@@ -81,6 +81,8 @@ class RateLimitsScheduler(Scheduler):
         while True:
             if item is None:
                 item = await self.queue.get()
+                assert isinstance(item.request.fn_call.fn, func.CallableFunction)
+                assert '_runtime_ctx' in item.request.fn_call.fn.signature.system_parameters
                 if item.num_retries > 0:
                     self.total_retried += 1
 
@@ -97,7 +99,6 @@ class RateLimitsScheduler(Scheduler):
                 continue
 
             # check rate limits
-            _logger.debug(f'checking rate limits for {self.resource_pool}')
             request_resources = self._get_request_resources(item.request)
             limits_info = self._check_resource_limits(request_resources)
             aws: list[Awaitable[None]] = []
@@ -116,21 +117,31 @@ class RateLimitsScheduler(Scheduler):
                 reset_at = limits_info.reset_at
                 if reset_at > now:
                     # we're waiting for the rate limit to reset
-                    wait_for_reset = asyncio.create_task(asyncio.sleep((reset_at - now).total_seconds()))
+                    wait_duration = (reset_at - now).total_seconds()
+                    wait_for_reset = asyncio.create_task(asyncio.sleep(wait_duration))
                     aws.append(wait_for_reset)
-                    _logger.debug(f'waiting for rate limit reset for {self.resource_pool}')
+                    _logger.debug(
+                        f'waiting {wait_duration:.2f}s for rate limit reset of '
+                        f'{self.resource_pool}:{limits_info.resource} (remaining={limits_info.remaining})'
+                    )
 
             if len(aws) > 0:
                 # we have something to wait for
+                report_ts = limits_info.recorded_at
                 done, pending = await asyncio.wait(aws, return_when=asyncio.FIRST_COMPLETED)
                 for task in pending:
                     task.cancel()
                 if completed_aw in done:
                     _logger.debug(f'wait(): completed request for {self.resource_pool}')
                 if wait_for_reset in done:
-                    _logger.debug(f'wait(): rate limit reset for {self.resource_pool}')
-                    # force waiting for another rate limit report before making any scheduling decisions
-                    self.pool_info.reset()
+                    _logger.debug(f'wait(): rate limit reset for {self.resource_pool}:{limits_info.resource}')
+                    last_report_ts = self.pool_info.resource_limits[limits_info.resource].recorded_at
+                    if report_ts == last_report_ts:
+                        # if we haven't seen a new report since we started waiting, force waiting for another rate limit
+                        # report before making any scheduling decisions
+                        # TODO: is it a good idea to discard the information we have?
+                        _logger.debug(f'resetting {self.resource_pool}: currently at {self.pool_info.debug_str()}')
+                        self.pool_info.reset()
                 # re-evaluate current capacity for current item
                 continue
 
@@ -158,16 +169,22 @@ class RateLimitsScheduler(Scheduler):
 
     def _check_resource_limits(self, request_resources: dict[str, int]) -> Optional[env.RateLimitInfo]:
         """Returns the most depleted resource, relative to its limit, or None if all resources are within limits"""
-        candidates: list[tuple[env.RateLimitInfo, float]] = []  # (info, relative usage)
+        candidates: list[tuple[env.RateLimitInfo, float]] = []  # (info, relative remaining)
         for resource, usage in request_resources.items():
-            # 0.05: leave some headroom, we don't have perfect information
             info = self.pool_info.resource_limits[resource]
             est_remaining = info.remaining - self.est_usage[resource] - usage
-            if est_remaining < 0.05 * info.limit:
-                candidates.append((info, est_remaining / info.limit))
-        if len(candidates) == 0:
-            return None
-        return min(candidates, key=lambda x: x[1])[0]
+            candidates.append((info, est_remaining / info.limit))
+        assert len(candidates) > 0
+        candidates.sort(key=lambda x: x[1])  # most depleted first
+        most_depleted = candidates[0]
+        _logger.debug(
+            f'check_resource_limits({request_resources}): '
+            f'most_depleted={most_depleted[0].resource}, rel_remaining={most_depleted[1]}'
+        )
+        # 0.05: leave some headroom, we don't have perfect information
+        if most_depleted[1] < 0.05:
+            return most_depleted[0]
+        return None
 
     async def _exec(self, request: FnCallArgs, exec_ctx: ExecCtx, num_retries: int, is_task: bool) -> None:
         assert all(not row.has_val[request.fn_call.slot_idx] for row in request.rows)
@@ -188,7 +205,8 @@ class RateLimitsScheduler(Scheduler):
                 for row, result in zip(request.rows, batch_result):
                     row[request.fn_call.slot_idx] = result
             else:
-                result = await pxt_fn.aexec(*request.args, **request.kwargs)
+                request_kwargs = {**request.kwargs, '_runtime_ctx': env.RuntimeCtx(is_retry=num_retries > 0)}
+                result = await pxt_fn.aexec(*request.args, **request_kwargs)
                 request.row[request.fn_call.slot_idx] = result
             end_ts = datetime.datetime.now(tz=datetime.timezone.utc)
             _logger.debug(
@@ -202,10 +220,14 @@ class RateLimitsScheduler(Scheduler):
             self.dispatcher.dispatch(request.rows, exec_ctx)
         except Exception as exc:
             _logger.debug(f'scheduler {self.resource_pool}: exception in slot {request.fn_call.slot_idx}: {exc}')
+            if hasattr(exc, 'response') and hasattr(exc.response, 'headers'):
+                _logger.debug(f'scheduler {self.resource_pool}: exception headers: {exc.response.headers}')
             if self.pool_info is None:
                 # our pool info should be available at this point
                 self._set_pool_info()
             assert self.pool_info is not None
+            self.pool_info.record_exc(exc)
+
             if num_retries < self.MAX_RETRIES:
                 retry_delay = self.pool_info.get_retry_delay(exc)
                 if retry_delay is not None:
@@ -214,7 +236,6 @@ class RateLimitsScheduler(Scheduler):
                     await asyncio.sleep(retry_delay)
                     self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
                     return
-            # TODO: update resource limits reported in exc.response.headers, if present
 
             # record the exception
             _, _, exc_tb = sys.exc_info()

pixeltable/exprs/array_slice.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
@@ -16,7 +16,7 @@ class ArraySlice(Expr):
     Slice operation on an array, eg, t.array_col[:, 1:2].
     """
 
-    def __init__(self, arr: Expr, index: tuple[Union[int, slice], ...]):
+    def __init__(self, arr: Expr, index: tuple[int | slice, ...]):
         assert arr.col_type.is_array_type()
         # determine result type
         super().__init__(arr.col_type)

pixeltable/exprs/data_row.py

@@ -281,6 +281,19 @@ class DataRow:
             pass
         self.vals[index] = None
 
+    def move_tmp_media_file(self, index: int, col: catalog.Column) -> None:
+        """If a media url refers to data in a temporary file, move the data to the MediaStore"""
+        if self.file_urls[index] is None:
+            return
+        assert self.excs[index] is None
+        assert col.col_type.is_media_type()
+        src_path = MediaStore.resolve_tmp_url(self.file_urls[index])
+        if src_path is None:
+            # The media url does not point to a temporary file, leave it as is
+            return
+        new_file_url = MediaStore.relocate_local_media_file(src_path, col)
+        self.file_urls[index] = new_file_url
+
     @property
     def rowid(self) -> tuple[int, ...]:
         return self.pk[:-1]

pixeltable/exprs/expr.py

@@ -7,7 +7,7 @@ import inspect
 import json
 import sys
 import typing
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, Optional, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, Optional, TypeVar, overload
 from uuid import UUID
 
 import numpy as np
@@ -550,7 +550,7 @@ class Expr(abc.ABC):
         else:
             return InPredicate(self, value_set_literal=value_set)
 
-    def astype(self, new_type: Union[ts.ColumnType, type, _AnnotatedAlias]) -> 'exprs.TypeCast':
+    def astype(self, new_type: ts.ColumnType | type | _AnnotatedAlias) -> 'exprs.TypeCast':
         from pixeltable.exprs import TypeCast
 
         # Interpret the type argument the same way we would if given in a schema
@@ -562,7 +562,7 @@ class Expr(abc.ABC):
         return TypeCast(self, col_type)
 
     def apply(
-        self, fn: Callable, *, col_type: Union[ts.ColumnType, type, _AnnotatedAlias, None] = None
+        self, fn: Callable, *, col_type: ts.ColumnType | type | _AnnotatedAlias | None = None
     ) -> 'exprs.FunctionCall':
         if col_type is not None:
             col_type = ts.ColumnType.normalize_type(col_type)
@@ -646,7 +646,7 @@ class Expr(abc.ABC):
 
     def _make_comparison(self, op: ComparisonOperator, other: object) -> 'exprs.Comparison':
         """
-        other: Union[Expr, LiteralPythonTypes]
+        other: Expr | LiteralPythonTypes
         """
         # TODO: check for compatibility
         from .comparison import Comparison
@@ -661,7 +661,7 @@ class Expr(abc.ABC):
     def __neg__(self) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, -1)
 
-    def __add__(self, other: object) -> Union[exprs.ArithmeticExpr, exprs.StringOp]:
+    def __add__(self, other: object) -> exprs.ArithmeticExpr | exprs.StringOp:
         if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
             return self._make_string_expr(StringOperator.CONCAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.ADD, other)
@@ -669,7 +669,7 @@ class Expr(abc.ABC):
     def __sub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __mul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __mul__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(self, str) or (isinstance(self, Expr) and self.col_type.is_string_type()):
             return self._make_string_expr(StringOperator.REPEAT, other)
         return self._make_arithmetic_expr(ArithmeticOperator.MUL, other)
@@ -683,7 +683,7 @@ class Expr(abc.ABC):
     def __floordiv__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._make_arithmetic_expr(ArithmeticOperator.FLOORDIV, other)
 
-    def __radd__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __radd__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
             return self._rmake_string_expr(StringOperator.CONCAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.ADD, other)
@@ -691,7 +691,7 @@ class Expr(abc.ABC):
     def __rsub__(self, other: object) -> 'exprs.ArithmeticExpr':
         return self._rmake_arithmetic_expr(ArithmeticOperator.SUB, other)
 
-    def __rmul__(self, other: object) -> Union['exprs.ArithmeticExpr', 'exprs.StringOp']:
+    def __rmul__(self, other: object) -> 'exprs.ArithmeticExpr' | 'exprs.StringOp':
         if isinstance(other, str) or (isinstance(other, Expr) and other.col_type.is_string_type()):
             return self._rmake_string_expr(StringOperator.REPEAT, other)
         return self._rmake_arithmetic_expr(ArithmeticOperator.MUL, other)
@@ -733,7 +733,7 @@ class Expr(abc.ABC):
 
     def _make_arithmetic_expr(self, op: ArithmeticOperator, other: object) -> 'exprs.ArithmeticExpr':
         """
-        other: Union[Expr, LiteralPythonTypes]
+        other: Expr | LiteralPythonTypes
         """
         # TODO: check for compatibility
         from .arithmetic_expr import ArithmeticExpr

pixeltable/exprs/function_call.py

@@ -4,7 +4,7 @@ import inspect
 import logging
 import sys
 from textwrap import dedent
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Sequence
 
 import sqlalchemy as sql
 
@@ -36,7 +36,7 @@ class FunctionCall(Expr):
     # - a component index, if the parameter is a non-variadic parameter
     # - a list of component indices, if the parameter is a variadic positional parameter
     # - a dict mapping keyword names to component indices, if the parameter is a variadic keyword parameter
-    bound_idxs: dict[str, Union[int, list[int], dict[str, int]]]
+    bound_idxs: dict[str, int | list[int] | dict[str, int]]
 
     return_type: ts.ColumnType
     group_by_start_idx: int

pixeltable/exprs/globals.py

@@ -2,10 +2,9 @@ from __future__ import annotations
 
 import datetime
 import enum
-from typing import Union
 
 # Python types corresponding to our literal types
-LiteralPythonTypes = Union[str, int, float, bool, datetime.datetime, datetime.date]
+LiteralPythonTypes = str | int | float | bool | datetime.datetime | datetime.date
 
 
 def print_slice(s: slice) -> str: