pixeltable 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

pixeltable/env.py

@@ -15,9 +15,7 @@ import sys
 import threading
 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
@@ -102,6 +100,8 @@ class Env:
     def _init_env(cls, reinit_db: bool = False) -> None:
         assert not cls.__initializing, 'Circular env initialization detected.'
         cls.__initializing = True
+        if cls._instance is not None:
+            cls._instance._clean_up()
         cls._instance = None
         env = Env()
         env._set_up(reinit_db=reinit_db)
@@ -247,7 +247,7 @@ class Env:
         if self._current_conn is None:
             assert self._current_session is None
             try:
-                self._current_isolation_level = 'SERIALIZABLE' if for_write else 'REPEATABLE_READ'
+                self._current_isolation_level = 'SERIALIZABLE'
                 with (
                     self.engine.connect().execution_options(isolation_level=self._current_isolation_level) as conn,
                     sql.orm.Session(conn) as session,
@@ -486,7 +486,7 @@ class Env:
                 raise excs.Error(error)
             self._logger.info(f'Using database at: {self.db_url}')
         else:
-            self._db_name = os.environ.get('PIXELTABLE_DB', 'pixeltable')
+            self._db_name = config.get_string_value('db') or 'pixeltable'
             self._pgdata_dir = Path(os.environ.get('PIXELTABLE_PGDATA', str(Config.get().home / 'pgdata')))
             # cleanup_mode=None will leave the postgres process running after Python exits
             # cleanup_mode='stop' will terminate the postgres process when Python exits
@@ -558,6 +558,14 @@ class Env:
         finally:
             engine.dispose()
 
+    def _pgserver_terminate_connections_stmt(self) -> str:
+        return f"""
+                SELECT pg_terminate_backend(pg_stat_activity.pid)
+                FROM pg_stat_activity
+                WHERE pg_stat_activity.datname = '{self._db_name}'
+                AND pid <> pg_backend_pid()
+            """
+
     def _drop_store_db(self) -> None:
         assert self._db_name is not None
         engine = sql.create_engine(self._dbms.default_system_db_url(), future=True, isolation_level='AUTOCOMMIT')
@@ -566,13 +574,7 @@ class Env:
             with engine.begin() as conn:
                 # terminate active connections
                 if self._db_server is not None:
-                    stmt = f"""
-                        SELECT pg_terminate_backend(pg_stat_activity.pid)
-                        FROM pg_stat_activity
-                        WHERE pg_stat_activity.datname = '{self._db_name}'
-                        AND pid <> pg_backend_pid()
-                    """
-                    conn.execute(sql.text(stmt))
+                    conn.execute(sql.text(self._pgserver_terminate_connections_stmt()))
                 # drop db
                 stmt = self._dbms.drop_db_stmt(preparer.quote(self._db_name))
                 conn.execute(sql.text(stmt))
@@ -750,12 +752,6 @@ class Env:
             else:
                 os.remove(path)
 
-    def num_tmp_files(self) -> int:
-        return len(glob.glob(f'{self._tmp_dir}/*'))
-
-    def create_tmp_path(self, extension: str = '') -> Path:
-        return self._tmp_dir / f'{uuid.uuid4()}{extension}'
-
     # def get_resource_pool_info(self, pool_id: str, pool_info_cls: Optional[Type[T]]) -> T:
     def get_resource_pool_info(self, pool_id: str, make_pool_info: Optional[Callable[[], T]] = None) -> T:
         """Returns the info object for the given id, creating it if necessary."""
@@ -816,6 +812,63 @@ class Env:
         except Exception as exc:
             raise excs.Error(f'Failed to load spaCy model: {spacy_model}') from exc
 
+    def _clean_up(self) -> None:
+        """
+        Internal cleanup method that properly closes all resources and resets state.
+        This is called before destroying the singleton instance.
+        """
+        assert self._current_session is None
+        assert self._current_conn is None
+
+        # Stop HTTP server
+        if self._httpd is not None:
+            try:
+                self._httpd.shutdown()
+                self._httpd.server_close()
+            except Exception as e:
+                _logger.warning(f'Error stopping HTTP server: {e}')
+
+        # First terminate all connections to the database
+        if self._db_server is not None:
+            assert self._dbms is not None
+            assert self._db_name is not None
+            try:
+                temp_engine = sql.create_engine(self._dbms.default_system_db_url(), isolation_level='AUTOCOMMIT')
+                try:
+                    with temp_engine.begin() as conn:
+                        conn.execute(sql.text(self._pgserver_terminate_connections_stmt()))
+                        _logger.info(f"Terminated all connections to database '{self._db_name}'")
+                except Exception as e:
+                    _logger.warning(f'Error terminating database connections: {e}')
+                finally:
+                    temp_engine.dispose()
+            except Exception as e:
+                _logger.warning(f'Error stopping database server: {e}')
+
+        # Dispose of SQLAlchemy engine (after stopping db server)
+        if self._sa_engine is not None:
+            try:
+                self._sa_engine.dispose()
+            except Exception as e:
+                _logger.warning(f'Error disposing engine: {e}')
+
+        # Close event loop
+        if self._event_loop is not None:
+            try:
+                if self._event_loop.is_running():
+                    self._event_loop.stop()
+                self._event_loop.close()
+            except Exception as e:
+                _logger.warning(f'Error closing event loop: {e}')
+
+        # Remove logging handlers
+        for handler in self._logger.handlers[:]:
+            try:
+                handler.close()
+                self._logger.removeHandler(handler)
+            except Exception as e:
+                _logger.warning(f'Error removing handler: {e}')
+
 
 def register_client(name: str) -> Callable:
     """Decorator that registers a third-party API client for use by Pixeltable.
@@ -890,6 +943,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 +954,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 +965,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 +1001,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 +1021,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/aggregation_node.py

@@ -103,6 +103,6 @@ class AggregationNode(ExecNode):
             self.row_builder.eval(prev_row, self.agg_fn_eval_ctx, profile=self.ctx.profile)
             self.output_batch.add_row(prev_row)
 
-        self.output_batch.flush_imgs(None, self.stored_img_cols, self.flushed_img_slots)
+        self.output_batch.flush_imgs(None, self.row_builder.stored_img_cols, self.flushed_img_slots)
         _logger.debug(f'AggregateNode: consumed {num_input_rows} rows, returning {len(self.output_batch.rows)} rows')
         yield self.output_batch

pixeltable/exec/cache_prefetch_node.py

@@ -12,8 +12,9 @@ from pathlib import Path
 from typing import Any, AsyncIterator, Iterator, Optional
 from uuid import UUID
 
-from pixeltable import env, exceptions as excs, exprs
+from pixeltable import exceptions as excs, exprs
 from pixeltable.utils.filecache import FileCache
+from pixeltable.utils.media_store import TempStore
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -219,7 +220,7 @@ class CachePrefetchNode(ExecNode):
             self.in_flight_requests[f] = url
 
     def __fetch_url(self, url: str) -> tuple[Optional[Path], Optional[Exception]]:
-        """Fetches a remote URL into Env.tmp_dir and returns its path"""
+        """Fetches a remote URL into the TempStore and returns its path"""
         _logger.debug(f'fetching url={url} thread_name={threading.current_thread().name}')
         parsed = urllib.parse.urlparse(url)
         # Use len(parsed.scheme) > 1 here to ensure we're not being passed
@@ -230,7 +231,7 @@ class CachePrefetchNode(ExecNode):
         if parsed.path:
             p = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
             extension = p.suffix
-        tmp_path = env.Env.get().create_tmp_path(extension=extension)
+        tmp_path = TempStore.create_path(extension=extension)
         try:
             _logger.debug(f'Downloading {url} to {tmp_path}')
             if parsed.scheme == 's3':

pixeltable/exec/exec_node.py

@@ -20,7 +20,6 @@ class ExecNode(abc.ABC):
     row_builder: exprs.RowBuilder
     input: Optional[ExecNode]
     flushed_img_slots: list[int]  # idxs of image slots of our output_exprs dependencies
-    stored_img_cols: list[exprs.ColumnSlotIdx]
     ctx: Optional[ExecContext]
 
     def __init__(
@@ -40,7 +39,6 @@ class ExecNode(abc.ABC):
         self.flushed_img_slots = [
             e.slot_idx for e in output_dependencies if e.col_type.is_image_type() and e.slot_idx not in output_slot_idxs
         ]
-        self.stored_img_cols = []
         self.ctx = None  # all nodes of a tree share the same context
 
     def set_ctx(self, ctx: ExecContext) -> None:
@@ -48,12 +46,6 @@ class ExecNode(abc.ABC):
         if self.input is not None:
             self.input.set_ctx(ctx)
 
-    def set_stored_img_cols(self, stored_img_cols: list[exprs.ColumnSlotIdx]) -> None:
-        self.stored_img_cols = stored_img_cols
-        # propagate batch size to the source
-        if self.input is not None:
-            self.input.set_stored_img_cols(stored_img_cols)
-
     @abc.abstractmethod
     def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         pass

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/globals.py

@@ -56,6 +56,7 @@ class Scheduler(abc.ABC):
         request: FnCallArgs
         num_retries: int
         exec_ctx: ExecCtx
+        retry_after: Optional[float] = None  # time.monotonic()
 
         def __lt__(self, other: Scheduler.QueueItem) -> bool:
             # prioritize by number of retries (more retries = higher priority)

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()
@@ -249,6 +270,7 @@ class RequestRateScheduler(Scheduler):
     num_in_flight: int
     total_requests: int
     total_retried: int
+    total_errors: int
 
     TIME_FORMAT = '%H:%M.%S %f'
     MAX_RETRIES = 3
@@ -273,6 +295,7 @@ class RequestRateScheduler(Scheduler):
         self.num_in_flight = 0
         self.total_requests = 0
         self.total_retried = 0
+        self.total_errors = 0
 
         # try to get the rate limit from the config
         elems = resource_pool.split(':')
@@ -291,6 +314,7 @@ class RequestRateScheduler(Scheduler):
             key = model
         requests_per_min = Config.get().get_int_value(key, section=section)
         requests_per_min = requests_per_min or self.DEFAULT_RATE_LIMIT
+        _logger.debug(f'rate limit for {self.resource_pool}: {requests_per_min} RPM')
         self.secs_per_request = 1 / (requests_per_min / 60)
 
     @classmethod
@@ -304,8 +328,12 @@ class RequestRateScheduler(Scheduler):
             if item.num_retries > 0:
                 self.total_retried += 1
             now = time.monotonic()
+            wait_duration = 0.0
+            if item.retry_after is not None:
+                wait_duration = item.retry_after - now
             if now - last_request_ts < self.secs_per_request:
-                wait_duration = self.secs_per_request - (now - last_request_ts)
+                wait_duration = max(wait_duration, self.secs_per_request - (now - last_request_ts))
+            if wait_duration > 0:
                 _logger.debug(f'waiting for {wait_duration} for {self.resource_pool}')
                 await asyncio.sleep(wait_duration)
 
@@ -351,15 +379,20 @@ class RequestRateScheduler(Scheduler):
 
         except Exception as exc:
             _logger.debug(f'exception for {self.resource_pool}: type={type(exc)}\n{exc}')
+            if hasattr(exc, 'response') and hasattr(exc.response, 'headers'):
+                _logger.debug(f'scheduler {self.resource_pool}: exception headers: {exc.response.headers}')
             is_rate_limit_error, retry_after = self._is_rate_limit_error(exc)
             if is_rate_limit_error and num_retries < self.MAX_RETRIES:
                 retry_delay = self._compute_retry_delay(num_retries, retry_after)
                 _logger.debug(f'scheduler {self.resource_pool}: retrying after {retry_delay}')
-                await asyncio.sleep(retry_delay)
-                self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx))
+                now = time.monotonic()
+                # put the request back in the queue right away, which prevents new requests from being generated until
+                # this one succeeds or exceeds its retry limit
+                self.queue.put_nowait(self.QueueItem(request, num_retries + 1, exec_ctx, retry_after=now + retry_delay))
                 return
 
             # record the exception
+            self.total_errors += 1
             _, _, exc_tb = sys.exc_info()
             for row in request.rows:
                 row.set_exc(request.fn_call.slot_idx, exc)
@@ -367,7 +400,7 @@ class RequestRateScheduler(Scheduler):
         finally:
             _logger.debug(
                 f'Scheduler stats: #in-flight={self.num_in_flight} #requests={self.total_requests}, '
-                f'#retried={self.total_retried}'
+                f'#retried={self.total_retried} #errors={self.total_errors}'
             )
             if is_task:
                 self.num_in_flight -= 1

pixeltable/exec/in_memory_data_node.py

@@ -2,7 +2,7 @@ import logging
 from typing import Any, AsyncIterator, Optional
 
 from pixeltable import catalog, exprs
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.media_store import TempStore
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -67,8 +67,7 @@ class InMemoryDataNode(ExecNode):
                 col = col_info.col
                 if col.col_type.is_image_type() and isinstance(val, bytes):
                     # this is a literal media file, ie, a sequence of bytes; save it as a binary file and store the path
-                    assert col.tbl.id == self.tbl.id
-                    filepath, _ = MediaStore.save_media_object(val, col, format=None)
+                    filepath, _ = TempStore.save_media_object(val, col, format=None)
                     output_row[col_info.slot_idx] = str(filepath)
                 else:
                     output_row[col_info.slot_idx] = val

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

@@ -14,7 +14,7 @@ import PIL.Image
 import sqlalchemy as sql
 
 from pixeltable import catalog, env
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.media_store import MediaStore, TempStore
 
 
 class DataRow:
@@ -270,7 +270,7 @@ class DataRow:
                     # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
                     # In that case, use WebP instead.
                     format = 'webp' if image.has_transparency_data else 'jpeg'
-                filepath, url = MediaStore.save_media_object(image, col, format=format)
+                filepath, url = MediaStore.get().save_media_object(image, col, format=format)
                 self.file_paths[index] = str(filepath)
                 self.file_urls[index] = url
             else:
@@ -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 a MediaStore"""
+        if self.file_urls[index] is None:
+            return
+        assert self.excs[index] is None
+        assert col.col_type.is_media_type()
+        src_path = TempStore.resolve_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.get().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