pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exprs/data_row.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import dataclasses
 import datetime
 import io
 import urllib.parse
 import urllib.request
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import pgvector.sqlalchemy  # type: ignore[import-untyped]
@@ -13,14 +14,81 @@ import PIL
 import PIL.Image
 import sqlalchemy as sql
 
-from pixeltable import env
+import pixeltable.utils.image as image_utils
+from pixeltable import catalog, env
+from pixeltable.utils.local_store import TempStore
+from pixeltable.utils.misc import non_none_dict_factory
+
+
+@dataclasses.dataclass
+class ArrayMd:
+    """
+    Metadata for array cells that are stored externally.
+    """
+
+    start: int
+    end: int
+
+    # we store bool arrays as packed bits (uint8 arrays), and need to record the shape to reconstruct the array
+    is_bool: bool = False
+    shape: tuple[int, ...] | None = None
+
+    def as_dict(self) -> dict:
+        # dict_factory: suppress Nones
+        x = dataclasses.asdict(self, dict_factory=non_none_dict_factory)
+        return x
+
+
+@dataclasses.dataclass
+class BinaryMd:
+    """
+    Metadata for binary cells that are stored externally.
+    """
+
+    start: int
+    end: int
+
+
+@dataclasses.dataclass
+class CellMd:
+    """
+    Content of the cellmd column.
+
+    All fields are optional, to minimize storage.
+    """
+
+    errortype: str | None = None
+    errormsg: str | None = None
+
+    # a list of file urls that are used to store images and arrays; only set for json and array columns
+    # for json columns: a list of all urls referenced in the column value
+    # for array columns: a single url
+    file_urls: list[str] | None = None
+
+    array_md: ArrayMd | None = None
+    binary_md: BinaryMd | None = None
+
+    @classmethod
+    def from_dict(cls, d: dict) -> CellMd:
+        d = d.copy()
+        if 'array_md' in d:
+            d['array_md'] = ArrayMd(**d['array_md'])
+        if 'binary_md' in d:
+            d['binary_md'] = BinaryMd(**d['binary_md'])
+        return cls(**d)
+
+    def as_dict(self) -> dict:
+        x = dataclasses.asdict(self, dict_factory=non_none_dict_factory)
+        return x
 
 
 class DataRow:
     """
     Encapsulates all data and execution state needed by RowBuilder and DataRowBatch:
     - state for in-memory computation
-    - state for storing the data
+    - state needed for expression evaluation
+    - containers for output column values
+
     This is not meant to be a black-box abstraction.
 
     In-memory representations by column type:
@@ -30,6 +98,8 @@ class DataRow:
     - BoolType: bool
     - TimestampType: datetime.datetime
     - DateType: datetime.date
+    - UUIDType: uuid.UUID
+    - BinaryType: bytes
     - JsonType: json-serializable object
     - ArrayType: numpy.ndarray
     - ImageType: PIL.Image.Image
@@ -38,37 +108,49 @@ class DataRow:
     - DocumentType: local path if available, otherwise url
     """
 
+    # expr evaluation state; indexed by slot idx
     vals: np.ndarray  # of object
     has_val: np.ndarray  # of bool
     excs: np.ndarray  # of object
-
-    # expr evaluation state; indexed by slot idx
     missing_slots: np.ndarray  # of bool; number of missing dependencies
     missing_dependents: np.ndarray  # of int16; number of missing dependents
     is_scheduled: np.ndarray  # of bool; True if this slot is scheduled for evaluation
 
-    # control structures that are shared across all DataRows in a batch
-    img_slot_idxs: list[int]
-    media_slot_idxs: list[int]
-    array_slot_idxs: list[int]
-
-    # the primary key of a store row is a sequence of ints (the number is different for table vs view)
-    pk: Optional[tuple[int, ...]]
+    # CellMd needed for query execution; needs to be indexed by slot idx, not column id, to work for joins
+    slot_md: dict[int, CellMd]
 
     # file_urls:
     # - stored url of file for media in vals[i]
     # - None if vals[i] is not media type
     # - not None if file_paths[i] is not None
+    # TODO: this is a sparse vector; should it be a dict[int, str]?
     file_urls: np.ndarray  # of str
 
     # file_paths:
     # - local path of media file in vals[i]; points to the file cache if file_urls[i] is remote
     # - None if vals[i] is not a media type or if there is no local file yet for file_urls[i]
+    # TODO: this is a sparse vector; should it be a dict[int, str]?
     file_paths: np.ndarray  # of str
 
+    # If `may_have_exc` is False, then we guarantee that no slot has an exception set. This is used to optimize
+    # exception handling under normal operation.
+    _may_have_exc: bool
+
+    # the primary key of a store row is a sequence of ints (the number is different for table vs view)
+    pk: tuple[int, ...] | None
     # for nested rows (ie, those produced by JsonMapperDispatcher)
-    parent_row: Optional[DataRow]
-    parent_slot_idx: Optional[int]
+    parent_row: DataRow | None
+    parent_slot_idx: int | None
+
+    # state for table output (insert()/update()); key: column id
+    cell_vals: dict[int, Any]  # materialized values of output columns, in the format required for the column
+    cell_md: dict[int, CellMd]
+
+    # control structures that are shared across all DataRows in a batch
+    img_slot_idxs: list[int]
+    media_slot_idxs: list[int]
+    array_slot_idxs: list[int]
+    json_slot_idxs: list[int]
 
     def __init__(
         self,
@@ -76,36 +158,42 @@ class DataRow:
         img_slot_idxs: list[int],
         media_slot_idxs: list[int],
         array_slot_idxs: list[int],
-        parent_row: Optional[DataRow] = None,
-        parent_slot_idx: Optional[int] = None,
+        json_slot_idxs: list[int],
+        parent_row: DataRow | None = None,
+        parent_slot_idx: int | None = None,
     ):
-        self.img_slot_idxs = img_slot_idxs
-        self.media_slot_idxs = media_slot_idxs
-        self.array_slot_idxs = array_slot_idxs
         self.init(size)
         self.parent_row = parent_row
         self.parent_slot_idx = parent_slot_idx
-
-    def init(self, num_slots: int) -> None:
-        self.vals = np.full(num_slots, None, dtype=object)
-        self.has_val = np.zeros(num_slots, dtype=bool)
-        self.excs = np.full(num_slots, None, dtype=object)
-        self.missing_slots = np.zeros(num_slots, dtype=bool)
-        self.missing_dependents = np.zeros(num_slots, dtype=np.int16)
-        self.is_scheduled = np.zeros(num_slots, dtype=bool)
+        self.img_slot_idxs = img_slot_idxs
+        self.media_slot_idxs = media_slot_idxs
+        self.array_slot_idxs = array_slot_idxs
+        self.json_slot_idxs = json_slot_idxs
+
+    def init(self, size: int) -> None:
+        self.vals = np.full(size, None, dtype=object)
+        self.has_val = np.zeros(size, dtype=bool)
+        self.excs = np.full(size, None, dtype=object)
+        self.missing_slots = np.zeros(size, dtype=bool)
+        self.missing_dependents = np.zeros(size, dtype=np.int16)
+        self.is_scheduled = np.zeros(size, dtype=bool)
+        self.slot_md = {}
+        self.file_urls = np.full(size, None, dtype=object)
+        self.file_paths = np.full(size, None, dtype=object)
+        self._may_have_exc = False
+        self.cell_vals = {}
+        self.cell_md = {}
         self.pk = None
-        self.file_urls = np.full(num_slots, None, dtype=object)
-        self.file_paths = np.full(num_slots, None, dtype=object)
         self.parent_row = None
         self.parent_slot_idx = None
 
-    def clear(self, idxs: Optional[np.ndarray] = None) -> None:
-        if idxs is not None:
-            self.has_val[idxs] = False
-            self.vals[idxs] = None
-            self.excs[idxs] = None
-            self.file_urls[idxs] = None
-            self.file_paths[idxs] = None
+    def clear(self, slot_idxs: np.ndarray | None = None) -> None:
+        if slot_idxs is not None:
+            self.has_val[slot_idxs] = False
+            self.vals[slot_idxs] = None
+            self.excs[slot_idxs] = None
+            self.file_urls[slot_idxs] = None
+            self.file_paths[slot_idxs] = None
         else:
             self.init(len(self.vals))
 
@@ -132,20 +220,23 @@ class DataRow:
     def set_pk(self, pk: tuple[int, ...]) -> None:
         self.pk = pk
 
-    def has_exc(self, slot_idx: Optional[int] = None) -> bool:
+    def has_exc(self, slot_idx: int | None = None) -> bool:
         """
         Returns True if an exception has been set for the given slot index, or for any slot index if slot_idx is None
         """
+        if not self._may_have_exc:
+            return False
+
         if slot_idx is not None:
             return self.excs[slot_idx] is not None
         return (self.excs != None).any()
 
-    def get_exc(self, slot_idx: int) -> Optional[Exception]:
+    def get_exc(self, slot_idx: int) -> Exception | None:
         exc = self.excs[slot_idx]
         assert exc is None or isinstance(exc, Exception)
         return exc
 
-    def get_first_exc(self) -> Optional[Exception]:
+    def get_first_exc(self) -> Exception | None:
         mask = self.excs != None
         if not mask.any():
             return None
@@ -154,6 +245,7 @@ class DataRow:
     def set_exc(self, slot_idx: int, exc: Exception) -> None:
         assert self.excs[slot_idx] is None
         self.excs[slot_idx] = exc
+        self._may_have_exc = True
 
         # an exception means the value is None
         self.has_val[slot_idx] = True
@@ -179,7 +271,7 @@ class DataRow:
 
         return self.vals[index]
 
-    def get_stored_val(self, index: int, sa_col_type: Optional[sql.types.TypeEngine] = None) -> Any:
+    def get_stored_val(self, index: int, sa_col_type: sql.types.TypeEngine | None = None) -> Any:
         """Return the value that gets stored in the db"""
         assert self.excs[index] is None
         if not self.has_val[index]:
@@ -214,6 +306,7 @@ class DataRow:
         """Assign in-memory cell value
         This allows overwriting
         """
+        assert isinstance(idx, int)
         assert self.excs[idx] is None
 
         if (idx in self.img_slot_idxs or idx in self.media_slot_idxs) and isinstance(val, str):
@@ -246,29 +339,46 @@ class DataRow:
             self.vals[idx] = val
         self.has_val[idx] = True
 
-    def flush_img(self, index: int, filepath: Optional[str] = None) -> None:
-        """Discard the in-memory value and save it to a local file, if filepath is not None"""
+    def prepare_col_val_for_save(self, index: int, col: catalog.Column | None = None) -> bool:
+        """
+        Prepare to save a column's value into the appropriate store. Discard unneeded values.
+
+        Return:
+            True if the media object in the column needs to be saved.
+        """
         if self.vals[index] is None:
-            return
+            return False
+
+        if self.file_urls[index] is not None:
+            return False
+
         assert self.excs[index] is None
         if self.file_paths[index] is None:
-            if filepath is not None:
-                # we want to save this to a file
-                self.file_paths[index] = filepath
-                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
-                image = self.vals[index]
-                assert isinstance(image, PIL.Image.Image)
-                # 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'
-                image.save(filepath, format=format)
+            if col is not None:
+                # This is a media object that needs to be saved
+                return True
             else:
-                # we discard the content of this cell
+                # This is a media object that we don't care about, so we discard it
                 self.has_val[index] = False
         else:
             # we already have a file for this image, nothing left to do
             pass
+
+        self.vals[index] = None
+        return False
+
+    def save_media_to_temp(self, index: int, col: catalog.Column) -> str:
+        """Save the media object in the column to the TempStore.
+        Objects cannot be saved directly to general destinations."""
+        assert col.col_type.is_media_type()
+        val = self.vals[index]
+        format = None
+        if isinstance(val, PIL.Image.Image):
+            format = image_utils.default_format(val)
+        filepath, url = TempStore.save_media_object(val, col, format=format)
+        self.file_paths[index] = str(filepath) if filepath is not None else None
         self.vals[index] = None
+        return url
 
     @property
     def rowid(self) -> tuple[int, ...]: