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

pixeltable/functions/huggingface.py

@@ -63,13 +63,10 @@ def sentence_transformer(
 
 @sentence_transformer.conditional_return_type
 def _(model_id: str) -> ts.ArrayType:
-    try:
-        from sentence_transformers import SentenceTransformer
+    from sentence_transformers import SentenceTransformer
 
-        model = _lookup_model(model_id, SentenceTransformer)
-        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
-    except ImportError:
-        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+    model = _lookup_model(model_id, SentenceTransformer)
+    return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf
@@ -201,13 +198,10 @@ def _(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[pxt.Array[(None,
 
 @clip.conditional_return_type
 def _(model_id: str) -> ts.ArrayType:
-    try:
-        from transformers import CLIPModel
+    from transformers import CLIPModel
 
-        model = _lookup_model(model_id, CLIPModel.from_pretrained)
-        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
-    except ImportError:
-        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf(batch_size=4)

pixeltable/functions/openai.py

@@ -23,6 +23,7 @@ import pixeltable as pxt
 from pixeltable import env, exprs, type_system as ts
 from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.media_store import TempStore
 
 if TYPE_CHECKING:
     import openai
@@ -216,7 +217,7 @@ async def speech(input: str, *, model: str, voice: str, model_kwargs: Optional[d
 
     content = await _openai_client().audio.speech.create(input=input, model=model, voice=voice, **model_kwargs)
     ext = model_kwargs.get('response_format', 'mp3')
-    output_filename = str(env.Env.get().create_tmp_path(f'.{ext}'))
+    output_filename = str(TempStore.create_path(extension=f'.{ext}'))
     content.write_to_file(output_filename)
     return output_filename
 

pixeltable/functions/video.py

@@ -9,8 +9,8 @@ import numpy as np
 import PIL.Image
 
 import pixeltable as pxt
-from pixeltable import env
 from pixeltable.utils.code import local_public_names
+from pixeltable.utils.media_store import TempStore
 
 _format_defaults: dict[str, tuple[str, str]] = {  # format -> (codec, ext)
     'wav': ('pcm_s16le', 'wav'),
@@ -109,7 +109,7 @@ class make_video(pxt.Aggregator):
         if frame is None:
             return
         if self.container is None:
-            self.out_file = env.Env.get().create_tmp_path('.mp4')
+            self.out_file = TempStore.create_path(extension='.mp4')
             self.container = av.open(str(self.out_file), mode='w')
             self.stream = self.container.add_stream('h264', rate=self.fps)
             self.stream.pix_fmt = 'yuv420p'
@@ -158,16 +158,16 @@ def extract_audio(
             return None
         audio_stream = container.streams.audio[stream_idx]
         # create this in our tmp directory, so it'll get cleaned up if it's being generated as part of a query
-        output_filename = str(env.Env.get().create_tmp_path(f'.{ext}'))
+        output_path = str(TempStore.create_path(extension=f'.{ext}'))
 
-        with av.open(output_filename, 'w', format=format) as output_container:
+        with av.open(output_path, 'w', format=format) as output_container:
             output_stream = output_container.add_stream(codec or default_codec)
             assert isinstance(output_stream, av.audio.stream.AudioStream)
             for packet in container.demux(audio_stream):
                 for frame in packet.decode():
                     output_container.mux(output_stream.encode(frame))  # type: ignore[arg-type]
 
-        return output_filename
+        return output_path
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/whisperx.py

@@ -0,0 +1,177 @@
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+
+import pixeltable as pxt
+from pixeltable.config import Config
+from pixeltable.functions.util import resolve_torch_device
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from transformers import Wav2Vec2Model
+    from whisperx.asr import FasterWhisperPipeline  # type: ignore[import-untyped]
+    from whisperx.diarize import DiarizationPipeline  # type: ignore[import-untyped]
+
+
+@pxt.udf
+def transcribe(
+    audio: pxt.Audio,
+    *,
+    model: str,
+    diarize: bool = False,
+    compute_type: Optional[str] = None,
+    language: Optional[str] = None,
+    task: Optional[str] = None,
+    chunk_size: Optional[int] = None,
+    alignment_model_name: Optional[str] = None,
+    interpolate_method: Optional[str] = None,
+    return_char_alignments: Optional[bool] = None,
+    diarization_model_name: Optional[str] = None,
+    num_speakers: Optional[int] = None,
+    min_speakers: Optional[int] = None,
+    max_speakers: Optional[int] = None,
+) -> dict:
+    """
+    Transcribe an audio file using WhisperX.
+
+    This UDF runs a transcription model _locally_ using the WhisperX library,
+    equivalent to the WhisperX `transcribe` function, as described in the
+    [WhisperX library documentation](https://github.com/m-bain/whisperX).
+
+    If `diarize=True`, then speaker diarization will also be performed. Several of the UDF parameters are only valid if
+    `diarize=True`, as documented in the parameters list below.
+
+    __Requirements:__
+
+    - `pip install whisperx`
+
+    Args:
+        audio: The audio file to transcribe.
+        model: The name of the model to use for transcription.
+        diarize: Whether to perform speaker diarization.
+        compute_type: The compute type to use for the model (e.g., `'int8'`, `'float16'`). If `None`,
+            defaults to `'float16'` on CUDA devices and `'int8'` otherwise.
+        language: The language code for the transcription (e.g., `'en'` for English).
+        task: The task to perform (e.g., `'transcribe'` or `'translate'`). Defaults to `'transcribe'`.
+        chunk_size: The size of the audio chunks to process, in seconds. Defaults to `30`.
+        alignment_model_name: The name of the alignment model to use. If `None`, uses the default model for the given
+            language. Only valid if `diarize=True`.
+        interpolate_method: The method to use for interpolation of the alignment results. If not specified, uses the
+            WhisperX default (`'nearest'`). Only valid if `diarize=True`.
+        return_char_alignments: Whether to return character-level alignments. Defaults to `False`.
+            Only valid if `diarize=True`.
+        diarization_model_name: The name of the diarization model to use. Defaults to
+            `pyannote/speaker-diarization-3.1`. Only valid if `diarize=True`.
+        num_speakers: The number of speakers to expect in the audio. By default, the model with try to detect the
+            number of speakers. Only valid if `diarize=True`.
+        min_speakers: If specified, the minimum number of speakers to expect in the audio.
+            Only valid if `diarize=True`.
+        max_speakers: If specified, the maximum number of speakers to expect in the audio.
+            Only valid if `diarize=True`.
+
+    Returns:
+        A dictionary containing the audio transcription, diarization (if enabled), and various other metadata.
+
+    Examples:
+        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
+        of the table `tbl`:
+
+        >>> tbl.add_computed_column(result=transcribe(tbl.audio, model='tiny.en'))
+
+        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
+        of the table `tbl`, with speaker diarization enabled, expecting at least 2 speakers:
+
+        >>> tbl.add_computed_column(
+        ...     result=transcribe(
+        ...         tbl.audio, model='tiny.en', diarize=True, min_speakers=2
+        ...     )
+        ... )
+    """
+    import whisperx  # type: ignore[import-untyped]
+
+    if not diarize:
+        args = locals()
+        for param in (
+            'alignment_model_name',
+            'interpolate_method',
+            'return_char_alignments',
+            'diarization_model_name',
+            'num_speakers',
+            'min_speakers',
+            'max_speakers',
+        ):
+            if args[param] is not None:
+                raise pxt.Error(f'`{param}` can only be set if `diarize=True`')
+
+    device = resolve_torch_device('auto', allow_mps=False)
+    compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
+    transcription_model = _lookup_transcription_model(model, device, compute_type)
+    audio_array: np.ndarray = whisperx.load_audio(audio)
+    kwargs: dict[str, Any] = {'language': language, 'task': task}
+    if chunk_size is not None:
+        kwargs['chunk_size'] = chunk_size
+    result: dict[str, Any] = transcription_model.transcribe(audio_array, batch_size=16, **kwargs)
+
+    if diarize:
+        # Alignment
+        alignment_model, metadata = _lookup_alignment_model(result['language'], device, alignment_model_name)
+        kwargs = {}
+        if interpolate_method is not None:
+            kwargs['interpolate_method'] = interpolate_method
+        if return_char_alignments is not None:
+            kwargs['return_char_alignments'] = return_char_alignments
+        result = whisperx.align(result['segments'], alignment_model, metadata, audio_array, device, **kwargs)
+
+        # Diarization
+        diarization_model = _lookup_diarization_model(device, diarization_model_name)
+        diarization_segments = diarization_model(
+            audio_array, num_speakers=num_speakers, min_speakers=min_speakers, max_speakers=max_speakers
+        )
+        result = whisperx.assign_word_speakers(diarization_segments, result)
+
+    return result
+
+
+def _lookup_transcription_model(model: str, device: str, compute_type: str) -> 'FasterWhisperPipeline':
+    import whisperx
+
+    key = (model, device, compute_type)
+    if key not in _model_cache:
+        transcription_model = whisperx.load_model(model, device, compute_type=compute_type)
+        _model_cache[key] = transcription_model
+    return _model_cache[key]
+
+
+def _lookup_alignment_model(language_code: str, device: str, model_name: Optional[str]) -> tuple['Wav2Vec2Model', dict]:
+    import whisperx
+
+    key = (language_code, device, model_name)
+    if key not in _alignment_model_cache:
+        model, metadata = whisperx.load_align_model(language_code=language_code, device=device, model_name=model_name)
+        _alignment_model_cache[key] = (model, metadata)
+    return _alignment_model_cache[key]
+
+
+def _lookup_diarization_model(device: str, model_name: Optional[str]) -> 'DiarizationPipeline':
+    from whisperx.diarize import DiarizationPipeline
+
+    key = (device, model_name)
+    if key not in _diarization_model_cache:
+        auth_token = Config.get().get_string_value('auth_token', section='hf')
+        kwargs: dict[str, Any] = {'device': device, 'use_auth_token': auth_token}
+        if model_name is not None:
+            kwargs['model_name'] = model_name
+        _diarization_model_cache[key] = DiarizationPipeline(**kwargs)
+    return _diarization_model_cache[key]
+
+
+_model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
+_alignment_model_cache: dict[tuple[str, str, Optional[str]], tuple['Wav2Vec2Model', dict]] = {}
+_diarization_model_cache: dict[tuple[str, Optional[str]], 'DiarizationPipeline'] = {}
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/{ext/functions → functions}/yolox.py

@@ -20,8 +20,6 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
     defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
 
-    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
-
     __Requirements__:
 
     - `pip install pixeltable-yolox`
@@ -55,8 +53,6 @@ def yolo_to_coco(detections: dict) -> list:
     """
     Converts the output of a YOLOX object detection model to COCO format.
 
-    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
-
     Args:
         detections: The output of a YOLOX object detection model, as returned by `yolox`.
 

pixeltable/globals.py

@@ -3,9 +3,10 @@ from __future__ import annotations
 import logging
 import os
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, NamedTuple, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, NamedTuple, Optional, Sequence, Union
 
 import pandas as pd
+import pydantic
 from pandas.io.formats.style import Styler
 
 from pixeltable import DataFrame, catalog, exceptions as excs, exprs, func, share, type_system as ts
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
         Path,  # OS paths, filenames, URLs
         Iterator[dict[str, Any]],  # iterator producing dictionaries of values
         RowData,  # list of dictionaries
+        Sequence[pydantic.BaseModel],  # list of Pydantic models
         DataFrame,  # Pixeltable DataFrame
         pd.DataFrame,  # pandas DataFrame
         datasets.Dataset,
@@ -396,7 +398,12 @@ def create_snapshot(
     )
 
 
-def create_replica(destination: str, source: str | catalog.Table) -> Optional[catalog.Table]:
+def create_replica(
+    destination: str,
+    source: str | catalog.Table,
+    bucket_name: str | None = None,
+    access: Literal['public', 'private'] = 'private',
+) -> Optional[catalog.Table]:
     """
     Create a replica of a table. Can be used either to create a remote replica of a local table, or to create a local
     replica of a remote table. A given table can have at most one replica per Pixeltable instance.
@@ -405,6 +412,12 @@ def create_replica(destination: str, source: str | catalog.Table) -> Optional[ca
         destination: Path where the replica will be created. Can be either a local path such as `'my_dir.my_table'`, or
             a remote URI such as `'pxt://username/mydir.my_table'`.
         source: Path to the source table, or (if the source table is a local table) a handle to the source table.
+        bucket_name: The name of the pixeltable cloud-registered bucket to use to store replica's data.
+            If no `bucket_name` is provided, the default Pixeltable storage bucket will be used.
+        access: Access control for the replica.
+
+            - `'public'`: Anyone can access this replica.
+            - `'private'`: Only the owner can access.
     """
     remote_dest = destination.startswith('pxt://')
     remote_source = isinstance(source, str) and source.startswith('pxt://')
@@ -414,7 +427,7 @@ def create_replica(destination: str, source: str | catalog.Table) -> Optional[ca
     if remote_dest:
         if isinstance(source, str):
             source = get_table(source)
-        share.push_replica(destination, source)
+        share.push_replica(destination, source, bucket_name, access)
         return None
     else:
         assert isinstance(source, str)

pixeltable/io/fiftyone.py

@@ -9,7 +9,7 @@ import puremagic
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 from pixeltable import exprs
-from pixeltable.env import Env
+from pixeltable.utils.media_store import TempStore
 
 
 class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
@@ -100,7 +100,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             assert isinstance(file, str)
         else:
             # Write the dynamically created image to a temp file
-            file = str(Env.get().create_tmp_path(f'.{self.__image_format}'))
+            file = TempStore.create_path(extension=f'.{self.__image_format}')
             img.save(file, format=self.__image_format)
 
         metadata = fo.ImageMetadata(
@@ -108,7 +108,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             mime_type=puremagic.from_file(file, mime=True),
             width=img.width,
             height=img.height,
-            filepath=file,
+            filepath=str(file),
             num_channels=len(img.getbands()),
         )
 

pixeltable/io/label_studio.py

@@ -19,6 +19,7 @@ from pixeltable.config import Config
 from pixeltable.exprs import ColumnRef, DataRow, Expr
 from pixeltable.io.external_store import Project
 from pixeltable.utils import coco
+from pixeltable.utils.media_store import TempStore
 
 # label_studio_sdk>=1 and label_studio_sdk<1 are not compatible, so we need to try
 # the import two different ways to insure intercompatibility
@@ -215,7 +216,7 @@ class LabelStudioProject(Project):
                 else:
                     # No localpath; create a temp file and upload it
                     assert isinstance(row[media_col_idx], PIL.Image.Image)
-                    file = env.Env.get().create_tmp_path(extension='.png')
+                    file = TempStore.create_path(extension='.png')
                     row[media_col_idx].save(file, format='png')
                     task_id = self.project.import_tasks(file)[0]
                     os.remove(file)

pixeltable/iterators/audio.py

@@ -5,7 +5,8 @@ from typing import Any, ClassVar, Optional
 
 import av
 
-from pixeltable import env, exceptions as excs, type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
+from pixeltable.utils.media_store import TempStore
 
 from .base import ComponentIterator
 
@@ -149,7 +150,7 @@ class AudioSplitter(ComponentIterator):
         target_chunk_start, target_chunk_end = self.chunks_to_extract_in_pts[self.next_pos]
         chunk_start_pts = 0
         chunk_end_pts = 0
-        chunk_file = str(env.Env.get().create_tmp_path(self.audio_path.suffix))
+        chunk_file = str(TempStore.create_path(extension=self.audio_path.suffix))
         output_container = av.open(chunk_file, mode='w')
         input_stream = self.container.streams.audio[0]
         codec_name = AudioSplitter.__codec_map.get(input_stream.codec_context.name, input_stream.codec_context.name)

pixeltable/iterators/document.py

@@ -213,12 +213,6 @@ class DocumentSplitter(ComponentIterator):
             if kwargs.get('limit') is None:
                 raise Error('limit is required with "token_limit"/"char_limit" separators')
 
-        # check dependencies at the end
-        if Separator.SENTENCE in separators:
-            _ = Env.get().spacy_nlp
-        if Separator.TOKEN_LIMIT in separators:
-            Env.get().require_package('tiktoken')
-
         return schema, []
 
     def __next__(self) -> dict[str, Any]:

pixeltable/metadata/__init__.py

@@ -25,6 +25,7 @@ def create_system_info(engine: sql.engine.Engine) -> None:
     """Create the system metadata record"""
     system_md = SystemInfoMd(schema_version=VERSION)
     record = SystemInfo(md=dataclasses.asdict(system_md))
+    _logger.debug(f'Creating pixeltable system info record {record}')
     with orm.Session(engine, future=True) as session:
         # Write system metadata only once for idempotency
         if session.query(SystemInfo).count() == 0:
@@ -54,7 +55,8 @@ for _, modname, _ in pkgutil.iter_modules([os.path.dirname(__file__) + '/convert
 def upgrade_md(engine: sql.engine.Engine) -> None:
     """Upgrade the metadata schema to the current version"""
     with orm.Session(engine) as session:
-        system_info = session.query(SystemInfo).one().md
+        # Get exclusive lock on SystemInfo row
+        system_info = session.query(SystemInfo).with_for_update().one().md
         md_version = system_info['schema_version']
         assert isinstance(md_version, int)
         _logger.info(f'Current database version: {md_version}, installed version: {VERSION}')

pixeltable/mypy/__init__.py

@@ -0,0 +1,3 @@
+from .mypy_plugin import plugin
+
+__all__ = ['plugin']

pixeltable/mypy/mypy_plugin.py

@@ -0,0 +1,123 @@
+from typing import Callable, ClassVar, Optional
+
+from mypy import nodes
+from mypy.plugin import AnalyzeTypeContext, ClassDefContext, FunctionContext, MethodSigContext, Plugin
+from mypy.plugins.common import add_attribute_to_class, add_method_to_class
+from mypy.types import AnyType, FunctionLike, Instance, NoneType, Type, TypeOfAny
+
+import pixeltable as pxt
+from pixeltable import exprs
+
+
+class PxtPlugin(Plugin):
+    __UDA_FULLNAME = f'{pxt.uda.__module__}.{pxt.uda.__name__}'
+    __ARRAY_GETITEM_FULLNAME = f'{pxt.Array.__module__}.{pxt.Array.__name__}.__class_getitem__'
+    __ADD_COLUMN_FULLNAME = f'{pxt.Table.__module__}.{pxt.Table.__name__}.{pxt.Table.add_column.__name__}'
+    __ADD_COMPUTED_COLUMN_FULLNAME = (
+        f'{pxt.Table.__module__}.{pxt.Table.__name__}.{pxt.Table.add_computed_column.__name__}'
+    )
+    __TYPE_MAP: ClassVar[dict] = {
+        pxt.Json: 'typing.Any',
+        pxt.Array: 'numpy.ndarray',
+        pxt.Image: 'PIL.Image.Image',
+        pxt.Video: 'builtins.str',
+        pxt.Audio: 'builtins.str',
+        pxt.Document: 'builtins.str',
+    }
+    __FULLNAME_MAP: ClassVar[dict] = {f'{k.__module__}.{k.__name__}': v for k, v in __TYPE_MAP.items()}
+
+    def get_function_hook(self, fullname: str) -> Optional[Callable[[FunctionContext], Type]]:
+        return adjust_uda_type
+
+    def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeTypeContext], Type]]:
+        if fullname in self.__FULLNAME_MAP:
+            subst_name = self.__FULLNAME_MAP[fullname]
+            return lambda ctx: adjust_pxt_type(ctx, subst_name)
+        return None
+
+    def get_method_signature_hook(self, fullname: str) -> Optional[Callable[[MethodSigContext], FunctionLike]]:
+        if fullname in (self.__ADD_COLUMN_FULLNAME, self.__ADD_COMPUTED_COLUMN_FULLNAME):
+            return adjust_kwargs
+        return None
+
+    def get_class_decorator_hook_2(self, fullname: str) -> Optional[Callable[[ClassDefContext], bool]]:
+        if fullname == self.__UDA_FULLNAME:
+            return adjust_uda_methods
+        return None
+
+
+def plugin(version: str) -> type:
+    return PxtPlugin
+
+
+_AGGREGATOR_FULLNAME = f'{pxt.Aggregator.__module__}.{pxt.Aggregator.__name__}'
+_FN_CALL_FULLNAME = f'{exprs.Expr.__module__}.{exprs.Expr.__name__}'
+
+
+def adjust_uda_type(ctx: FunctionContext) -> Type:
+    """
+    Mypy doesn't understand that a class with a @uda decorator isn't actually a class, so it assumes
+    that sum(expr), for example, actually returns an instance of sum. We correct this by changing the
+    return type of any subclass of `Aggregator` to `FunctionCall`.
+    """
+    ret_type = ctx.default_return_type
+    if isinstance(ret_type, Instance) and (
+        ret_type.type.fullname == _AGGREGATOR_FULLNAME
+        or any(base.type.fullname == _AGGREGATOR_FULLNAME for base in ret_type.type.bases)
+    ):
+        ret_type = AnyType(TypeOfAny.special_form)
+    return ret_type
+
+
+def adjust_pxt_type(ctx: AnalyzeTypeContext, subst_name: str) -> Type:
+    """
+    Replaces the special Pixeltable classes (such as pxt.Array) with their standard equivalents (such as np.ndarray).
+    """
+    if subst_name == 'typing.Any':
+        return AnyType(TypeOfAny.special_form)
+    return ctx.api.named_type(subst_name, [])
+
+
+def adjust_kwargs(ctx: MethodSigContext) -> FunctionLike:
+    """
+    Mypy has a "feature" where it will spit out multiple warnings if a method with signature
+    ```
+    def my_func(*, arg1: int, arg2: str, **kwargs: Expr)
+    ```
+    (for example) is called with bare kwargs:
+    ```
+    my_func(my_kwarg=value)
+    ```
+    This is a disaster for type-checking of add_column and add_computed_column. Here we adjust the signature so
+    that mypy thinks it is simply
+    ```
+    def my_func(**kwargs: Any)
+    ```
+    thereby avoiding any type-checking errors. For details, see: <https://github.com/python/mypy/issues/18481>
+    """
+    sig = ctx.default_signature
+    new_arg_names = sig.arg_names[-1:]
+    new_arg_types = [AnyType(TypeOfAny.special_form)]
+    new_arg_kinds = sig.arg_kinds[-1:]
+    return sig.copy_modified(arg_names=new_arg_names, arg_types=new_arg_types, arg_kinds=new_arg_kinds)
+
+
+def adjust_uda_methods(ctx: ClassDefContext) -> bool:
+    """
+    Mypy does not handle the `@pxt.uda` aggregator well; it continues to treat the decorated class as a class,
+    even though it has been replaced by an `AggregateFunction`. Here we add static methods to the class that
+    imitate various (instance) methods of `AggregateFunction` so that they can be properly type-checked.
+    """
+    list_type = ctx.api.named_type('builtins.list', [AnyType(TypeOfAny.special_form)])
+    fn_arg = nodes.Argument(nodes.Var('fn'), AnyType(TypeOfAny.special_form), None, nodes.ARG_POS)
+    args_arg = nodes.Argument(nodes.Var('args'), AnyType(TypeOfAny.special_form), None, nodes.ARG_STAR)
+    kwargs_arg = nodes.Argument(nodes.Var('kwargs'), AnyType(TypeOfAny.special_form), None, nodes.ARG_STAR2)
+    add_method_to_class(ctx.api, ctx.cls, '__init__', args=[args_arg, kwargs_arg], return_type=NoneType())
+    add_method_to_class(
+        ctx.api, ctx.cls, 'to_sql', args=[fn_arg], return_type=AnyType(TypeOfAny.special_form), is_staticmethod=True
+    )
+    add_method_to_class(
+        ctx.api, ctx.cls, 'overload', args=[fn_arg], return_type=AnyType(TypeOfAny.special_form), is_staticmethod=True
+    )
+    add_attribute_to_class(ctx.api, ctx.cls, 'signatures', typ=list_type, is_classvar=True)
+    return True

pixeltable/plan.py

@@ -394,9 +394,6 @@ class Planner:
                 row_builder, computed_exprs, plan.output_exprs, input=plan, maintain_input_order=False
             )
 
-        stored_col_info = row_builder.output_slot_idxs()
-        stored_img_col_info = [info for info in stored_col_info if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
         plan.set_ctx(
             exec.ExecContext(
                 row_builder,
@@ -428,10 +425,6 @@ class Planner:
             col = tbl.cols_by_name[col_name]
             plan.row_builder.add_table_column(col, expr.slot_idx)
 
-        stored_col_info = plan.row_builder.output_slot_idxs()
-        stored_img_col_info = [info for info in stored_col_info if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
-
         plan.set_ctx(
             exec.ExecContext(
                 plan.row_builder, batch_size=0, show_pbar=True, num_computed_exprs=0, ignore_errors=ignore_errors
@@ -657,10 +650,6 @@ class Planner:
         for i, col in enumerate(copied_cols + list(recomputed_cols)):  # same order as select_list
             plan.row_builder.add_table_column(col, select_list[i].slot_idx)
         # TODO: avoid duplication with view_load_plan() logic (where does this belong?)
-        stored_img_col_info = [
-            info for info in plan.row_builder.output_slot_idxs() if info.col.col_type.is_image_type()
-        ]
-        plan.set_stored_img_cols(stored_img_col_info)
         return plan
 
     @classmethod
@@ -727,8 +716,6 @@ class Planner:
                 row_builder, output_exprs=view_output_exprs, input_exprs=base_output_exprs, input=plan
             )
 
-        stored_img_col_info = [info for info in row_builder.output_slot_idxs() if info.col.col_type.is_image_type()]
-        plan.set_stored_img_cols(stored_img_col_info)
         exec_ctx.ignore_errors = True
         plan.set_ctx(exec_ctx)
         return plan, len(row_builder.default_eval_ctx.target_exprs)
@@ -1053,7 +1040,4 @@ class Planner:
         computed_exprs = row_builder.output_exprs - row_builder.input_exprs
         plan.ctx.num_computed_exprs = len(computed_exprs)  # we are adding a computed column, so we need to evaluate it
 
-        # we want to flush images
-        if col.is_computed and col.is_stored and col.col_type.is_image_type():
-            plan.set_stored_img_cols(row_builder.output_slot_idxs())
         return plan

pixeltable/share/packager.py

@@ -24,7 +24,7 @@ from pixeltable.env import Env
 from pixeltable.metadata import schema
 from pixeltable.utils import sha256sum
 from pixeltable.utils.formatter import Formatter
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.media_store import MediaStore, TempStore
 
 _logger = logging.getLogger('pixeltable')
 
@@ -57,7 +57,7 @@ class TablePackager:
 
     def __init__(self, table: catalog.Table, additional_md: Optional[dict[str, Any]] = None) -> None:
         self.table = table
-        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
         # Load metadata
@@ -92,10 +92,10 @@ class TablePackager:
         self.bundle_path = self.__build_tarball()
 
         _logger.info('Extracting preview data.')
-        self.md['count'] = self.table.count()
+        self.md['row_count'] = self.table.count()
         preview_header, preview = self.__extract_preview_data()
         self.md['preview_header'] = preview_header
-        self.md['preview'] = preview
+        self.md['preview_data'] = preview
 
         _logger.info(f'Packaging complete: {self.bundle_path}')
         return self.bundle_path
@@ -335,7 +335,7 @@ class TableRestorer:
     def __init__(self, tbl_path: str, md: Optional[dict[str, Any]] = None) -> None:
         self.tbl_path = tbl_path
         self.md = md
-        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
     def restore(self, bundle_path: Path) -> pxt.Table:
@@ -619,7 +619,7 @@ class TableRestorer:
                 # in self.media_files.
                 src_path = self.tmp_dir / 'media' / parsed_url.netloc
                 # Move the file to the media store and update the URL.
-                self.media_files[url] = MediaStore.relocate_local_media_file(src_path, media_col)
+                self.media_files[url] = MediaStore.get().relocate_local_media_file(src_path, media_col)
             return self.media_files[url]
         # For any type of URL other than a local file, just return the URL as-is.
         return url