pixeltable 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

pixeltable/func/function.py

@@ -3,16 +3,20 @@ from __future__ import annotations
 import abc
 import importlib
 import inspect
-from typing import Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import sqlalchemy as sql
 
 import pixeltable as pxt
+import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
 from .globals import resolve_symbol
 from .signature import Signature
 
+if TYPE_CHECKING:
+    from .expr_template_function import ExprTemplateFunction
+
 
 class Function(abc.ABC):
     """Base class for Pixeltable's function interface.
@@ -99,6 +103,38 @@ class Function(abc.ABC):
         self._conditional_return_type = fn
         return fn
 
+    def using(self, **kwargs: Any) -> 'ExprTemplateFunction':
+        from pixeltable import exprs
+
+        from .expr_template_function import ExprTemplateFunction
+
+        # Resolve each kwarg into a parameter binding
+        bindings: dict[str, exprs.Expr] = {}
+        for k, v in kwargs.items():
+            if k not in self.signature.parameters:
+                raise excs.Error(f'Unknown parameter: {k}')
+            param = self.signature.parameters[k]
+            expr = exprs.Expr.from_object(v)
+            if not param.col_type.is_supertype_of(expr.col_type):
+                raise excs.Error(f'Expected type `{param.col_type}` for parameter `{k}`; got `{expr.col_type}`')
+            bindings[k] = v  # Use the original value, not the Expr (The Expr is only for validation)
+
+        residual_params = [
+            p for p in self.signature.parameters.values() if p.name not in bindings
+        ]
+
+        # Bind each remaining parameter to a like-named variable
+        for param in residual_params:
+            bindings[param.name] = exprs.Variable(param.name, param.col_type)
+
+        call = exprs.FunctionCall(self, bindings)
+
+        # Construct the (n-k)-ary signature of the new function. We use `call.col_type` for this, rather than
+        # `self.signature.return_type`, because the return type of the new function may be specialized via a
+        # conditional return type.
+        new_signature = Signature(call.col_type, residual_params, self.signature.is_batched)
+        return ExprTemplateFunction(call, new_signature)
+
     @abc.abstractmethod
     def exec(self, *args: Any, **kwargs: Any) -> Any:
         """Execute the function with the given arguments and return the result."""

pixeltable/func/signature.py

@@ -91,6 +91,7 @@ class Signature:
         self.parameters_by_pos = parameters.copy()
         self.constant_parameters = [p for p in parameters if not p.is_batched]
         self.batched_parameters = [p for p in parameters if p.is_batched]
+        self.required_parameters = [p for p in parameters if not p.has_default()]
         self.py_signature = inspect.Signature([p.to_py_param() for p in self.parameters_by_pos])
 
     def get_return_type(self) -> ts.ColumnType:

pixeltable/functions/mistralai.py

@@ -36,7 +36,6 @@ def chat_completions(
     temperature: Optional[float] = 0.7,
     top_p: Optional[float] = 1.0,
     max_tokens: Optional[int] = None,
-    min_tokens: Optional[int] = None,
     stop: Optional[list[str]] = None,
     random_seed: Optional[int] = None,
     response_format: Optional[dict] = None,
@@ -75,7 +74,6 @@ def chat_completions(
         temperature=temperature,
         top_p=top_p,
         max_tokens=_opt(max_tokens),
-        min_tokens=_opt(min_tokens),
         stop=stop,
         random_seed=_opt(random_seed),
         response_format=response_format,  # type: ignore[arg-type]

pixeltable/functions/ollama.py

@@ -68,7 +68,7 @@ def generate(
         raw=raw,
         format=format,
         options=options,
-    )  # type: ignore[call-overload]
+    ).dict()  # type: ignore[call-overload]
 
 
 @pxt.udf
@@ -103,7 +103,7 @@ def chat(
         tools=tools,
         format=format,
         options=options,
-    )  # type: ignore[call-overload]
+    ).dict()  # type: ignore[call-overload]
 
 
 @pxt.udf(batch_size=16)
@@ -135,8 +135,8 @@ def embed(
         model=model,
         input=input,
         truncate=truncate,
-        options=options,  # type: ignore[arg-type]
-    )
+        options=options,
+    ).dict()
     return [np.array(data, dtype=np.float64) for data in results['embeddings']]
 
 

pixeltable/globals.py

@@ -46,6 +46,7 @@ def create_table(
         num_retained_versions: Number of versions of the table to retain.
         comment: An optional comment; its meaning is user-defined.
         media_validation: Media validation policy for the table.
+
             - `'on_read'`: validate media files at query time
             - `'on_write'`: validate media files during insert/update operations
 
@@ -149,7 +150,9 @@ def create_view(
         tbl_version_path = base._tbl_version_path
     elif isinstance(base, DataFrame):
         base._validate_mutable('create_view')
-        tbl_version_path = base.tbl
+        if len(base._from_clause.tbls) > 1:
+            raise excs.Error('Cannot create a view of a join')
+        tbl_version_path = base._from_clause.tbls[0]
         where = base.where_clause
     else:
         raise excs.Error('`base` must be an instance of `Table` or `DataFrame`')
@@ -296,31 +299,42 @@ def move(path: str, new_path: str) -> None:
     obj._move(new_p.name, new_dir._id)
 
 
-def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> None:
+def drop_table(table: Union[str, catalog.Table], force: bool = False, ignore_errors: bool = False) -> None:
     """Drop a table, view, or snapshot.
 
     Args:
-        path: Path to the [`Table`][pixeltable.Table].
+        table: Fully qualified name, or handle, of the table to be dropped.
         force: If `True`, will also drop all views and sub-views of this table.
         ignore_errors: If `True`, return silently if the table does not exist (without throwing an exception).
 
     Raises:
-        Error: If the path does not exist or does not designate a table object, and `ignore_errors=False`.
+        Error: If the name does not exist or does not designate a table object, and `ignore_errors=False`.
 
     Examples:
-        >>> pxt.drop_table('my_table')
+        Drop a table by its fully qualified name:
+        >>> pxt.drop_table('subdir.my_table')
+
+        Drop a table by its handle:
+        >>> t = pxt.get_table('subdir.my_table')
+        ... pxt.drop_table(t)
+
     """
     cat = Catalog.get()
-    path_obj = catalog.Path(path)
-    try:
-        cat.paths.check_is_valid(path_obj, expected=catalog.Table)
-    except Exception as e:
-        if ignore_errors or force:
-            _logger.info(f'Skipped table `{path}` (does not exist).')
-            return
-        else:
-            raise e
-    tbl = cat.paths[path_obj]
+    if isinstance(table, str):
+        tbl_path_obj = catalog.Path(table)
+        try:
+            cat.paths.check_is_valid(tbl_path_obj, expected=catalog.Table)
+        except Exception as e:
+            if ignore_errors or force:
+                _logger.info(f'Skipped table `{table}` (does not exist).')
+                return
+            else:
+                raise e
+        tbl = cat.paths[tbl_path_obj]
+    else:
+        tbl = table
+        tbl_path_obj = catalog.Path(tbl._path)
+
     assert isinstance(tbl, catalog.Table)
     if len(cat.tbl_dependents[tbl._id]) > 0:
         dependent_paths = [dep._path for dep in cat.tbl_dependents[tbl._id]]
@@ -328,10 +342,10 @@ def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> N
             for dependent_path in dependent_paths:
                 drop_table(dependent_path, force=True)
         else:
-            raise excs.Error(f'Table {path} has dependents: {", ".join(dependent_paths)}')
+            raise excs.Error(f'Table {tbl._path} has dependents: {", ".join(dependent_paths)}')
     tbl._drop()
-    del cat.paths[path_obj]
-    _logger.info(f'Dropped table `{path}`.')
+    del cat.paths[tbl_path_obj]
+    _logger.info(f'Dropped table `{tbl._path}`.')
 
 
 def list_tables(dir_path: str = '', recursive: bool = True) -> list[str]:

pixeltable/index/embedding_index.py

@@ -136,7 +136,12 @@ class EmbeddingIndex(IndexBase):
         """Validate the signature"""
         assert isinstance(embed_fn, func.Function)
         sig = embed_fn.signature
-        if len(sig.parameters) != 1 or sig.parameters_by_pos[0].col_type.type_enum != expected_type:
+
+        # The embedding function must be a 1-ary function of the correct type. But it's ok if the function signature
+        # has more than one parameter, as long as it has at most one *required* parameter.
+        if (len(sig.parameters) == 0
+            or len(sig.required_parameters) > 1
+            or sig.parameters_by_pos[0].col_type.type_enum != expected_type):
             raise excs.Error(
                 f'{name} must take a single {expected_type.name.lower()} parameter, but has signature {sig}')
 

pixeltable/io/__init__.py

@@ -2,7 +2,7 @@ from .external_store import ExternalStore, SyncStatus
 from .globals import create_label_studio_project, export_images_as_fo_dataset, import_json, import_rows
 from .hf_datasets import import_huggingface_dataset
 from .pandas import import_csv, import_excel, import_pandas
-from .parquet import import_parquet
+from .parquet import import_parquet, export_parquet
 
 __default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
 __removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet'}

pixeltable/io/parquet.py

@@ -7,11 +7,14 @@ import random
 import typing
 from collections import deque
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import numpy as np
 import PIL.Image
+import datetime
 
+import pixeltable as pxt
+from pixeltable.env import Env
 import pixeltable.exceptions as exc
 import pixeltable.type_system as ts
 from pixeltable.utils.transactional_directory import transactional_directory
@@ -39,28 +42,44 @@ def _write_batch(value_batch: dict[str, deque], schema: pa.Schema, output_path:
     parquet.write_table(tab, str(output_path))
 
 
-def save_parquet(df: pxt.DataFrame, dest_path: Path, partition_size_bytes: int = 100_000_000) -> None:
+def export_parquet(
+            table_or_df: Union[pxt.Table, pxt.DataFrame],
+            parquet_path: Path,
+            partition_size_bytes: int = 100_000_000,
+            inline_images: bool = False
+            ) -> None:
     """
-    Internal method to stream dataframe data to parquet format.
-    Does not materialize the dataset to memory.
+    Exports a dataframe's data to one or more Parquet files. Requires pyarrow to be installed.
 
-    It preserves pixeltable type metadata in a json file, which would otherwise
+    It additionally writes the pixeltable metadata in a json file, which would otherwise
     not be available in the parquet format.
 
-    Images are stored inline in a compressed format in their parquet file.
-
     Args:
-        df : dataframe to save.
-        dest_path : path to directory to save the parquet files to.
-        partition_size_bytes : maximum target size for each chunk. Default 100_000_000 bytes.
+        table_or_df : Table or Dataframe to export.
+        parquet_path : Path to directory to write the parquet files to.
+        partition_size_bytes : The maximum target size for each chunk. Default 100_000_000 bytes.
+        inline_images : If True, images are stored inline in the parquet file. This is useful
+                        for small images, to be imported as pytorch dataset. But can be inefficient
+                        for large images, and cannot be imported into pixeltable.
+                        If False, will raise an error if the Dataframe has any image column.
+                        Default False.
     """
     from pixeltable.utils.arrow import to_arrow_schema
 
+    df: pxt.DataFrame
+    if isinstance(table_or_df, pxt.catalog.Table):
+        df = table_or_df._df()
+    else:
+        df = table_or_df
+
     type_dict = {k: v.as_dict() for k, v in df.schema.items()}
     arrow_schema = to_arrow_schema(df.schema)
 
+    if not inline_images and any(col_type.is_image_type() for col_type in df.schema.values()):
+        raise exc.Error('Cannot export Dataframe with image columns when inline_images is False')
+
     # store the changes atomically
-    with transactional_directory(dest_path) as temp_path:
+    with transactional_directory(parquet_path) as temp_path:
         # dump metadata json file so we can inspect what was the source of the parquet file later on.
         json.dump(df.as_dict(), (temp_path / '.pixeltable.json').open('w'))
         json.dump(type_dict, (temp_path / '.pixeltable.column_types.json').open('w'))  # keep type metadata
@@ -111,6 +130,7 @@ def save_parquet(df: pxt.DataFrame, dest_path: Path, partition_size_bytes: int =
                 elif col_type.is_bool_type():
                     length = 1
                 elif col_type.is_timestamp_type():
+                    val = val.astimezone(datetime.timezone.utc)
                     length = 8
                 else:
                     assert False, f'unknown type {col_type} for {col_name}'
@@ -139,7 +159,7 @@ def parquet_schema_to_pixeltable_schema(parquet_path: str) -> dict[str, Optional
 
 
 def import_parquet(
-    table_path: str,
+    table: str,
     *,
     parquet_path: str,
     schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
@@ -148,7 +168,7 @@ def import_parquet(
     """Creates a new base table from a Parquet file or set of files. Requires pyarrow to be installed.
 
     Args:
-        table_path: Path to the table.
+        table: Fully qualified name of the table to import the data into.
         parquet_path: Path to an individual Parquet file or directory of Parquet files.
         schema_overrides: If specified, then for each (name, type) pair in `schema_overrides`, the column with
             name `name` will be given type `type`, instead of being inferred from the Parquet dataset. The keys in
@@ -157,7 +177,7 @@ def import_parquet(
         kwargs: Additional arguments to pass to `create_table`.
 
     Returns:
-        A handle to the newly created [`Table`][pixeltable.Table].
+        A handle to the newly created table.
     """
     from pyarrow import parquet
 
@@ -176,11 +196,11 @@ def import_parquet(
         if v is None:
             raise exc.Error(f'Could not infer pixeltable type for column {k} from parquet file')
 
-    if table_path in pxt.list_tables():
-        raise exc.Error(f'Table {table_path} already exists')
+    if table in pxt.list_tables():
+        raise exc.Error(f'Table {table} already exists')
 
     try:
-        tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
+        tmp_name = f'{table}_tmp_{random.randint(0, 100000000)}'
         tab = pxt.create_table(tmp_name, schema, **kwargs)
         for fragment in parquet_dataset.fragments:  # type: ignore[attr-defined]
             for batch in fragment.to_batches():
@@ -190,5 +210,5 @@ def import_parquet(
         _logger.error(f'Error while inserting Parquet file into table: {e}')
         raise e
 
-    pxt.move(tmp_name, table_path)
-    return pxt.get_table(table_path)
+    pxt.move(tmp_name, table)
+    return pxt.get_table(table)

pixeltable/iterators/__init__.py

@@ -1,5 +1,6 @@
 from .base import ComponentIterator
 from .document import DocumentSplitter
+from .image import TileIterator
 from .string import StringSplitter
 from .video import FrameIterator
 

pixeltable/iterators/image.py

@@ -0,0 +1,100 @@
+from typing import Any, Sequence
+
+import PIL.Image
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from pixeltable.iterators.base import ComponentIterator
+
+
+class TileIterator(ComponentIterator):
+    """
+    Iterator over tiles of an image. Each image will be divided into tiles of size `tile_size`, and the tiles will be
+    iterated over in row-major order (left-to-right, then top-to-bottom). An optional `overlap` parameter may be
+    specified. If the tiles do not exactly cover the image, then the rightmost and bottommost tiles will be padded with
+    blackspace, so that the output images all have the exact size `tile_size`.
+
+    Args:
+        image: Image to split into tiles.
+        tile_size: Size of each tile, as a pair of integers `[width, height]`.
+        overlap: Amount of overlap between adjacent tiles, as a pair of integers `[width, height]`.
+    """
+
+    __image: PIL.Image.Image
+    __tile_size: Sequence[int]
+    __overlap: Sequence[int]
+    __width: int
+    __height: int
+    __xlen: int
+    __ylen: int
+    __i: int
+    __j: int
+
+    def __init__(
+        self,
+        image: PIL.Image.Image,
+        *,
+        tile_size: tuple[int, int],
+        overlap: tuple[int, int] = (0, 0),
+    ):
+        if overlap[0] >= tile_size[0] or overlap[1] >= tile_size[1]:
+            raise excs.Error(f"overlap dimensions {overlap} are not strictly smaller than tile size {tile_size}")
+
+        self.__image = image
+        self.__image.load()
+        self.__tile_size = tile_size
+        self.__overlap = overlap
+        self.__width, self.__height = image.size
+        # Justification for this formula: let t = tile_size[0], o = overlap[0]. Then the values of w (= width) that
+        # exactly accommodate an integer number of tiles are t, 2t - o, 3t - 2o, 4t - 3o, ...
+        # This formula ensures that t, 2t - o, 3t - 2o, ... result in an xlen of 1, 2, 3, ...
+        # but t + 1, 2t - o + 1, 3t - 2o + 1, ... result in an xlen of 2, 3, 4, ...
+        self.__xlen = (self.__width - overlap[0] - 1) // (tile_size[0] - overlap[0]) + 1
+        self.__ylen = (self.__height - overlap[1] - 1) // (tile_size[1] - overlap[1]) + 1
+        self.__i = 0
+        self.__j = 0
+
+    def __next__(self) -> dict[str, Any]:
+        if self.__j >= self.__ylen:
+            raise StopIteration
+
+        x1 = self.__i * (self.__tile_size[0] - self.__overlap[0])
+        y1 = self.__j * (self.__tile_size[1] - self.__overlap[1])
+        # If x2 > self.__width, PIL does the right thing and pads the image with blackspace
+        x2 = x1 + self.__tile_size[0]
+        y2 = y1 + self.__tile_size[1]
+        tile = self.__image.crop((x1, y1, x2, y2))
+        result = {
+            'tile': tile,
+            'tile_coord': [self.__i, self.__j],
+            'tile_box': [x1, y1, x2, y2]
+        }
+
+        self.__i += 1
+        if self.__i >= self.__xlen:
+            self.__i = 0
+            self.__j += 1
+        return result
+
+    def close(self) -> None:
+        pass
+
+    def set_pos(self, pos: int) -> None:
+        self.__j = pos // self.__xlen
+        self.__i = pos % self.__xlen
+
+    @classmethod
+    def input_schema(cls, *args: Any, **kwargs: Any) -> dict[str, ts.ColumnType]:
+        return {
+            'image': ts.ImageType(),
+            'tile_size': ts.JsonType(),
+            'overlap': ts.JsonType(),
+        }
+
+    @classmethod
+    def output_schema(cls,  *args: Any, **kwargs: Any) -> tuple[dict[str, ts.ColumnType], list[str]]:
+        return {
+            'tile': ts.ImageType(),
+            'tile_coord': ts.JsonType(),
+            'tile_box': ts.JsonType(),
+        }, ['tile']

pixeltable/iterators/video.py

@@ -23,13 +23,13 @@ class FrameIterator(ComponentIterator):
     exact number of frames will be extracted. If neither is specified, then all frames will be extracted. The first
     frame of the video will always be extracted, and the remaining frames will be spaced as evenly as possible.
 
-        Args:
-            video: URL or path of the video to use for frame extraction.
-            fps: Number of frames to extract per second of video. This may be a fractional value, such as 0.5.
-                If omitted or set to 0.0, then the native framerate of the video will be used (all frames will be
-                extracted). If `fps` is greater than the frame rate of the video, an error will be raised.
-            num_frames: Exact number of frames to extract. The frames will be spaced as evenly as possible. If
-                `num_frames` is greater than the number of frames in the video, all frames will be extracted.
+    Args:
+        video: URL or path of the video to use for frame extraction.
+        fps: Number of frames to extract per second of video. This may be a fractional value, such as 0.5.
+            If omitted or set to 0.0, then the native framerate of the video will be used (all frames will be
+            extracted). If `fps` is greater than the frame rate of the video, an error will be raised.
+        num_frames: Exact number of frames to extract. The frames will be spaced as evenly as possible. If
+            `num_frames` is greater than the number of frames in the video, all frames will be extracted.
     """
 
     # Input parameters
@@ -180,7 +180,6 @@ class FrameIterator(ComponentIterator):
         self.container.close()
 
     def set_pos(self, pos: int) -> None:
-        """Seek to frame idx"""
         if pos == self.next_pos:
             return  # already there
 

pixeltable/metadata/__init__.py

@@ -10,7 +10,7 @@ import sqlalchemy.orm as orm
 from .schema import SystemInfo, SystemInfoMd
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 22
+VERSION = 23
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:

pixeltable/metadata/converters/convert_22.py

@@ -0,0 +1,17 @@
+from typing import Any, Optional
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=22)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, substitution_fn=__substitute_md)
+
+
+def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
+    if isinstance(v, dict) and '_classname' in v and v['_classname'] == 'DataFrame':
+        v['from_clause'] = {'tbls': [v['tbl']], 'join_clauses': []}
+        return k, v
+    return None

pixeltable/metadata/notes.py

@@ -2,6 +2,7 @@
 # rather than as a comment, so that the existence of a description can be enforced by
 # the unit tests when new versions are added.
 VERSION_NOTES = {
+    23: 'DataFrame.from_clause',
     22: 'TableMd/ColumnMd.media_validation',
     21: 'Separate InlineArray and InlineList',
     20: 'Store DB timestamps in UTC',