pixeltable 0.4.15__py3-none-any.whl → 0.4.16__py3-none-any.whl

pixeltable/exec/cell_materialization_node.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import io
+import logging
+import os
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+import numpy as np
+import pgvector.sqlalchemy  # type: ignore[import-untyped]
+import PIL.Image
+import sqlalchemy as sql
+
+import pixeltable.type_system as ts
+import pixeltable.utils.image as image_utils
+from pixeltable import catalog, exprs
+from pixeltable.env import Env
+from pixeltable.utils.local_store import LocalStore
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+from .globals import INLINED_OBJECT_MD_KEY, InlinedObjectMd
+
+_logger = logging.getLogger('pixeltable')
+
+
+class CellMaterializationNode(ExecNode):
+    """
+    Node to populate DataRow.cell_vals/cell_md.
+
+    For now, the scope is limited to populating DataRow.cells_vals for json and array columns.
+
+    Array values:
+    - Arrays < MAX_DB_ARRAY_SIZE are stored inline in the db column
+    - Larger arrays are written to inlined_obj_files
+    - Bool arrays are stored as packed bits (uint8)
+    - cell_md: holds the url of the file, plus start and end offsets, plus bool flag and shape for bool arrays
+      (this allows us to query cell_md to get the total external storage size of an array column)
+
+    Json values:
+    - Inlined images and arrays are written to inlined_obj_files and replaced with a dict containing the object
+      location
+    - Bool arrays are also stored as packed bits; the dict also contains the shape and bool flag
+    - cell_md contains the list of urls for the inlined objects.
+
+    TODO:
+    - execute file IO via asyncio Tasks in a thread pool?
+      (we already seem to be getting 90% of hardware IO throughput)
+    - subsume all cell materialization
+    """
+
+    output_col_info: dict[catalog.Column, int]  # value: slot idx
+
+    # execution state
+    inlined_obj_files: list[Path]  # only [-1] is open for writing
+    buffered_writer: io.BufferedWriter | None  # BufferedWriter for inlined_obj_files[-1]
+
+    MIN_FILE_SIZE = 8 * 2**20  # 8MB
+    MAX_DB_ARRAY_SIZE = 512  # max size of array stored in table column; in bytes
+
+    def __init__(self, input: ExecNode):
+        super().__init__(input.row_builder, [], [], input)
+        self.output_col_info = {
+            col: slot_idx
+            for col, slot_idx in input.row_builder.table_columns.items()
+            if slot_idx is not None and (col.col_type.is_json_type() or col.col_type.is_array_type())
+        }
+        self.inlined_obj_files = []
+        self.buffered_writer = None
+
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        async for batch in self.input:
+            for row in batch:
+                for col, slot_idx in self.output_col_info.items():
+                    if row.has_exc(slot_idx):
+                        # Nulls in JSONB columns need to be stored as sql.sql.null(), otherwise it stores a json 'null'
+                        row.cell_vals[col.id] = sql.sql.null() if col.col_type.is_json_type() else None
+                        exc = row.get_exc(slot_idx)
+                        row.cell_md[col.id] = exprs.CellMd(errortype=type(exc).__name__, errormsg=str(exc))
+                        continue
+
+                    val = row[slot_idx]
+                    if val is None:
+                        row.cell_vals[col.id] = sql.sql.null() if col.col_type.is_json_type() else None
+                        row.cell_md[col.id] = None
+                        continue
+
+                    if col.col_type.is_json_type():
+                        self._materialize_json_cell(row, col, val)
+                    else:
+                        assert col.col_type.is_array_type()
+                        assert isinstance(val, np.ndarray)
+                        self._materialize_array_cell(row, col, val)
+
+                    # continue with only the currently open file
+                    self.inlined_obj_files = self.inlined_obj_files[-1:]
+
+            yield batch
+
+        self._flush_buffer(finalize=True)
+
+    def init_writer(self) -> None:
+        if self.buffered_writer is None:
+            self._reset_buffer()
+            assert self.buffered_writer is not None
+
+    def close(self) -> None:
+        if self.buffered_writer is not None:
+            # there must have been an error, otherwise _flush_full_buffer(finalize=True) would have set this to None
+            self.buffered_writer.close()
+            self.buffered_writer = None
+
+    def _materialize_json_cell(self, row: exprs.DataRow, col: catalog.Column, val: Any) -> None:
+        if self._json_has_inlined_objs(val):
+            row.cell_vals[col.id] = self._rewrite_json(val)
+            row.cell_md[col.id] = exprs.CellMd(file_urls=[local_path.as_uri() for local_path in self.inlined_obj_files])
+        else:
+            row.cell_vals[col.id] = val
+            row.cell_md[col.id] = None
+
+    def _materialize_array_cell(self, row: exprs.DataRow, col: catalog.Column, val: np.ndarray) -> None:
+        if isinstance(col.sa_col_type, pgvector.sqlalchemy.Vector):
+            # this is a vector column (ie, used for a vector index): store the array itself
+            row.cell_vals[col.id] = val
+            row.cell_md[col.id] = None
+        elif val.nbytes <= self.MAX_DB_ARRAY_SIZE:
+            # this array is small enough to store in the db column (type: binary) directly
+            buffer = io.BytesIO()
+            np.save(buffer, val, allow_pickle=False)
+            row.cell_vals[col.id] = buffer.getvalue()
+            row.cell_md[col.id] = None
+        else:
+            # append this array to the buffer and store its location in the cell md
+            ar: np.ndarray
+            if np.issubdtype(val.dtype, np.bool_):
+                # for bool arrays, store as packed bits, otherwise it's 1 byte per element
+                ar = np.packbits(val)
+            else:
+                ar = val
+            self.init_writer()
+            start = self.buffered_writer.tell()
+            np.save(self.buffered_writer, ar, allow_pickle=False)
+            end = self.buffered_writer.tell()
+            row.cell_vals[col.id] = None
+            cell_md = exprs.CellMd(
+                file_urls=[self.inlined_obj_files[-1].as_uri()], array_md=exprs.ArrayMd(start=start, end=end)
+            )
+            if np.issubdtype(val.dtype, np.bool_):
+                cell_md.array_md.is_bool = True
+                cell_md.array_md.shape = val.shape
+            row.cell_md[col.id] = cell_md
+            self._flush_buffer()
+
+        assert row.cell_vals[col.id] is not None or row.cell_md[col.id] is not None
+
+    def _json_has_inlined_objs(self, element: Any) -> bool:
+        if isinstance(element, list):
+            return any(self._json_has_inlined_objs(v) for v in element)
+        if isinstance(element, dict):
+            return any(self._json_has_inlined_objs(v) for v in element.values())
+        return isinstance(element, (np.ndarray, PIL.Image.Image))
+
+    def _rewrite_json(self, element: Any) -> Any:
+        """Recursively rewrites a JSON structure by writing any inlined arrays or images to self.buffered_writer."""
+        if isinstance(element, list):
+            return [self._rewrite_json(v) for v in element]
+        if isinstance(element, dict):
+            return {k: self._rewrite_json(v) for k, v in element.items()}
+        if isinstance(element, np.ndarray):
+            obj_md = self._write_inlined_array(element)
+            return {INLINED_OBJECT_MD_KEY: obj_md.as_dict()}
+        if isinstance(element, PIL.Image.Image):
+            obj_md = self._write_inlined_image(element)
+            return {INLINED_OBJECT_MD_KEY: obj_md.as_dict()}
+        return element
+
+    def _write_inlined_array(self, ar: np.ndarray) -> InlinedObjectMd:
+        """Write an ndarray to buffered_writer and return its metadata."""
+        self.init_writer()
+        url_idx = len(self.inlined_obj_files) - 1
+        start = self.buffered_writer.tell()
+        shape: tuple[int, ...] | None
+        is_bool_array: bool
+        if np.issubdtype(ar.dtype, np.bool_):
+            shape = ar.shape
+            ar = np.packbits(ar)
+            is_bool_array = True
+        else:
+            shape = None
+            is_bool_array = False
+        np.save(self.buffered_writer, ar, allow_pickle=False)
+        end = self.buffered_writer.tell()
+        self._flush_buffer()
+        return InlinedObjectMd(
+            type=ts.ColumnType.Type.ARRAY.name,
+            url_idx=url_idx,
+            array_md=exprs.ArrayMd(start=start, end=end, is_bool=is_bool_array, shape=shape),
+        )
+
+    def _write_inlined_image(self, img: PIL.Image.Image) -> InlinedObjectMd:
+        """Write a PIL image to buffered_writer and return: index into inlined_obj_files, start offset, end offset"""
+        self.init_writer()
+        url_idx = len(self.inlined_obj_files) - 1
+        start = self.buffered_writer.tell()
+        img.save(self.buffered_writer, format=image_utils.default_format(img))
+        end = self.buffered_writer.tell()
+        self._flush_buffer()
+        return InlinedObjectMd(type=ts.ColumnType.Type.IMAGE.name, url_idx=url_idx, img_start=start, img_end=end)
+
+    def _reset_buffer(self) -> None:
+        local_path = LocalStore(Env.get().media_dir)._prepare_path_raw(
+            self.row_builder.tbl.id, 0, self.row_builder.tbl.version
+        )
+        self.inlined_obj_files.append(local_path)
+        fh = open(local_path, 'wb', buffering=self.MIN_FILE_SIZE * 2)  # noqa: SIM115
+        assert isinstance(fh, io.BufferedWriter)
+        self.buffered_writer = fh
+
+    def _flush_buffer(self, finalize: bool = False) -> None:
+        """Flush buffered_writer to storage if it exceeds its minimum size or finalize is True."""
+        if self.buffered_writer is None:
+            return
+        if self.buffered_writer.tell() < self.MIN_FILE_SIZE and not finalize:
+            return
+        self.buffered_writer.flush()
+        os.fsync(self.buffered_writer.fileno())  # needed to force bytes cached by OS to storage
+        self.buffered_writer.close()
+        if finalize:
+            self.buffered_writer = None
+        else:
+            self._reset_buffer()

pixeltable/exec/cell_reconstruction_node.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import io
+import logging
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+import numpy as np
+import PIL.Image
+
+import pixeltable.type_system as ts
+from pixeltable import exprs
+from pixeltable.utils import parse_local_file_path
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+from .globals import INLINED_OBJECT_MD_KEY, InlinedObjectMd
+
+_logger = logging.getLogger('pixeltable')
+
+
+def json_has_inlined_objs(element: Any) -> bool:
+    """Returns True if element contains inlined objects produced by CellMaterializationNode."""
+    if isinstance(element, list):
+        return any(json_has_inlined_objs(v) for v in element)
+    if isinstance(element, dict):
+        if INLINED_OBJECT_MD_KEY in element:
+            return True
+        return any(json_has_inlined_objs(v) for v in element.values())
+    return False
+
+
+def reconstruct_json(element: Any, urls: list[str], file_handles: dict[Path, io.BufferedReader]) -> Any:
+    """Recursively reconstructs inlined objects in a json structure."""
+    if isinstance(element, list):
+        return [reconstruct_json(v, urls, file_handles) for v in element]
+    if isinstance(element, dict):
+        if INLINED_OBJECT_MD_KEY in element:
+            obj_md = InlinedObjectMd.from_dict(element[INLINED_OBJECT_MD_KEY])
+            url = urls[obj_md.url_idx]
+            local_path = parse_local_file_path(url)
+            if local_path not in file_handles:
+                file_handles[local_path] = open(local_path, 'rb')  # noqa: SIM115
+            fp = file_handles[local_path]
+
+            if obj_md.type == ts.ColumnType.Type.ARRAY.name:
+                fp.seek(obj_md.array_md.start)
+                ar = load_array(
+                    fp, obj_md.array_md.start, obj_md.array_md.end, obj_md.array_md.is_bool, obj_md.array_md.shape
+                )
+                return ar
+            else:
+                fp.seek(obj_md.img_start)
+                bytesio = io.BytesIO(fp.read(obj_md.img_end - obj_md.img_start))
+                img = PIL.Image.open(bytesio)
+                img.load()
+                assert fp.tell() == obj_md.img_end, f'{fp.tell()} != {obj_md.img_end} / {obj_md.img_start}'
+                return img
+        else:
+            return {k: reconstruct_json(v, urls, file_handles) for k, v in element.items()}
+    return element
+
+
+def load_array(
+    fh: io.BufferedReader, start: int, end: int, is_bool_array: bool, shape: tuple[int, ...] | None
+) -> np.ndarray:
+    """Loads an array from a section of a file."""
+    fh.seek(start)
+    ar = np.load(fh, allow_pickle=False)
+    assert fh.tell() == end
+    if is_bool_array:
+        assert shape is not None
+        ar = np.unpackbits(ar, count=np.prod(shape)).reshape(shape).astype(bool)
+    return ar
+
+
+class CellReconstructionNode(ExecNode):
+    """
+    Reconstruction of stored json and array cells that were produced by CellMaterializationNode.
+    """
+
+    json_refs: list[exprs.ColumnRef]
+    array_refs: list[exprs.ColumnRef]
+    file_handles: dict[Path, io.BufferedReader]  # key: file path
+
+    def __init__(
+        self,
+        json_refs: list[exprs.ColumnRef],
+        array_refs: list[exprs.ColumnRef],
+        row_builder: exprs.RowBuilder,
+        input: ExecNode | None = None,
+    ):
+        super().__init__(row_builder, [], [], input)
+        self.json_refs = json_refs
+        self.array_refs = array_refs
+        self.file_handles = {}
+
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        async for batch in self.input:
+            for row in batch:
+                for col_ref in self.json_refs:
+                    val = row[col_ref.slot_idx]
+                    if val is None:
+                        continue
+                    cell_md = row.slot_md.get(col_ref.slot_idx)
+                    if cell_md is None or cell_md.file_urls is None or not json_has_inlined_objs(row[col_ref.slot_idx]):
+                        continue
+                    row[col_ref.slot_idx] = reconstruct_json(val, cell_md.file_urls, self.file_handles)
+
+                for col_ref in self.array_refs:
+                    cell_md = row.slot_md.get(col_ref.slot_idx)
+                    if cell_md is not None and cell_md.array_md is not None:
+                        assert row[col_ref.slot_idx] is None
+                        assert cell_md.file_urls is not None and len(cell_md.file_urls) == 1
+                        row[col_ref.slot_idx] = self._reconstruct_array(cell_md)
+                    else:
+                        assert row[col_ref.slot_idx] is None or isinstance(row[col_ref.slot_idx], np.ndarray)
+
+            yield batch
+
+    def close(self) -> None:
+        for fp in self.file_handles.values():
+            fp.close()
+
+    def _reconstruct_array(self, cell_md: exprs.CellMd) -> np.ndarray:
+        assert cell_md.array_md is not None
+        local_path = parse_local_file_path(cell_md.file_urls[0])
+        assert local_path is not None
+        if local_path not in self.file_handles:
+            self.file_handles[local_path] = open(str(local_path), 'rb')  # noqa: SIM115
+        fp = self.file_handles[local_path]
+        ar = load_array(
+            fp, cell_md.array_md.start, cell_md.array_md.end, bool(cell_md.array_md.is_bool), cell_md.array_md.shape
+        )
+        return ar

pixeltable/exec/exec_node.py

@@ -39,7 +39,7 @@ 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.ctx = None  # all nodes of a tree share the same context
+        self.ctx = input.ctx if input is not None else None
 
     def set_ctx(self, ctx: ExecContext) -> None:
         self.ctx = ctx

pixeltable/exec/expr_eval/evaluators.py

@@ -311,6 +311,7 @@ class JsonMapperDispatcher(Evaluator):
                     img_slot_idxs=[],
                     media_slot_idxs=[],
                     array_slot_idxs=[],
+                    json_slot_idxs=[],
                     parent_row=row,
                     parent_slot_idx=self.e.slot_idx,
                 )

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -306,6 +306,9 @@ class ExprEvalNode(ExecNode):
                     task.cancel()
             _ = await asyncio.gather(*active_tasks, return_exceptions=True)
 
+            # expr cleanup
+            exprs.Expr.release_list(self.exec_ctx.all_exprs)
+
     def dispatch_exc(
         self, rows: list[exprs.DataRow], slot_with_exc: int, exc_tb: TracebackType, exec_ctx: ExecCtx
     ) -> None:

pixeltable/exec/expr_eval/globals.py

@@ -149,6 +149,7 @@ class ExecCtx:
     gc_targets: np.ndarray  # bool per slot; True if this is an intermediate expr (ie, not part of our output)
     eval_ctx: np.ndarray  # bool per slot; EvalCtx.slot_idxs as a mask
     literals: dict[int, Any]  # key: slot idx; value: literal value for this slot; used to pre-populate rows
+    all_exprs: list[exprs.Expr]  # all evaluated exprs; needed for cleanup
 
     def __init__(
         self,
@@ -165,6 +166,7 @@ class ExecCtx:
         self.gc_targets[[e.slot_idx for e in self.row_builder.output_exprs]] = False
 
         output_ctx = self.row_builder.create_eval_ctx(output_exprs, exclude=input_exprs)
+        self.all_exprs = output_ctx.exprs
         self.literals = {e.slot_idx: e.val for e in output_ctx.exprs if isinstance(e, exprs.Literal)}
         self.eval_ctx = np.zeros(self.row_builder.num_materialized, dtype=bool)
         non_literal_slot_idxs = [e.slot_idx for e in output_ctx.exprs if not isinstance(e, exprs.Literal)]

pixeltable/exec/globals.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import dataclasses
+
+from pixeltable.exprs import ArrayMd
+from pixeltable.utils.misc import non_none_dict_factory
+
+INLINED_OBJECT_MD_KEY = '__pxtinlinedobjmd__'
+
+
+@dataclasses.dataclass
+class InlinedObjectMd:
+    type: str  # corresponds to ts.ColumnType.Type
+    url_idx: int
+    img_start: int | None = None
+    img_end: int | None = None
+    array_md: ArrayMd | None = None
+
+    @classmethod
+    def from_dict(cls, d: dict) -> InlinedObjectMd:
+        if 'array_md' in d:
+            array_md = ArrayMd(**d['array_md'])
+            del d['array_md']
+            return cls(**d, array_md=array_md)
+        else:
+            return cls(**d)
+
+    def as_dict(self) -> dict:
+        result = dataclasses.asdict(self, dict_factory=non_none_dict_factory)
+        if self.array_md is not None:
+            result['array_md'] = self.array_md.as_dict()
+        return result

pixeltable/exec/object_store_save_node.py

@@ -7,7 +7,6 @@ from collections import defaultdict, deque
 from concurrent import futures
 from pathlib import Path
 from typing import AsyncIterator, Iterator, NamedTuple, Optional
-from uuid import UUID
 
 from pixeltable import exprs
 from pixeltable.utils.object_stores import ObjectOps, ObjectPath, StorageTarget
@@ -81,9 +80,7 @@ class ObjectStoreSaveNode(ExecNode):
         num_missing: int  # number of references to media files in this row
         delete_destinations: list[Path]  # paths to delete after all copies are complete
 
-    def __init__(
-        self, tbl_id: UUID, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True
-    ):
+    def __init__(self, file_col_info: list[exprs.ColumnSlotIdx], input: ExecNode, retain_input_order: bool = True):
         # input_/output_exprs=[]: we don't have anything to evaluate
         super().__init__(input.row_builder, [], [], input)
         self.retain_input_order = retain_input_order

pixeltable/exec/row_update_node.py

@@ -14,10 +14,18 @@ class RowUpdateNode(ExecNode):
     Update individual rows in the input batches, identified by key columns.
 
     The updates for a row are provided as a dict of column names to new values.
-    The node assumes that all update dicts contain the same keys, and it populates the slots of the columns present in
-    the update list.
+    Populates the slots of the columns present in the update list.
+    Assumptions:
+    - all update dicts contain the same keys
+    - the input node populates DataRow.cell_vals for all primary key columns
     """
 
+    updates: dict[tuple, dict[catalog.Column, Any]]
+    is_rowid_key: bool  # if True, key_vals_batch contains rowids rather than primary key values
+    col_slot_idxs: dict[catalog.Column, int]
+    pk_columns: list[catalog.Column]
+    matched_key_vals: set[tuple]
+
     def __init__(
         self,
         tbl: catalog.TableVersionPath,
@@ -37,16 +45,16 @@ class RowUpdateNode(ExecNode):
             for col_ref in row_builder.unique_exprs
             if isinstance(col_ref, exprs.ColumnRef)
         }
+        # all update target columns should have assigned slot idxs
+        assert all(col in all_col_slot_idxs for col in col_vals_batch[0])
         self.col_slot_idxs = {col: all_col_slot_idxs[col] for col in col_vals_batch[0]}
-        self.key_slot_idxs = {col: all_col_slot_idxs[col] for col in tbl.tbl_version.get().primary_key_columns()}
-        self.matched_key_vals: set[tuple] = set()
+        self.pk_columns = tbl.tbl_version.get().primary_key_columns()
+        self.matched_key_vals = set()
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
         async for batch in self.input:
             for row in batch:
-                key_vals = (
-                    row.rowid if self.is_rowid_key else tuple(row[slot_idx] for slot_idx in self.key_slot_idxs.values())
-                )
+                key_vals = row.rowid if self.is_rowid_key else tuple(row.cell_vals[col.id] for col in self.pk_columns)
                 if key_vals not in self.updates:
                     continue
                 self.matched_key_vals.add(key_vals)
@@ -59,11 +67,10 @@ class RowUpdateNode(ExecNode):
     def unmatched_rows(self) -> list[dict[str, Any]]:
         """Return rows that didn't get used in the updates as a list of dicts compatible with TableVersion.insert()."""
         result: list[dict[str, Any]] = []
-        key_cols = self.key_slot_idxs.keys()
         for key_vals, col_vals in self.updates.items():
             if key_vals in self.matched_key_vals:
                 continue
-            row = {col.name: val for col, val in zip(key_cols, key_vals)}
+            row = {col.name: val for col, val in zip(self.pk_columns, key_vals)}
             row.update({col.name: val for col, val in col_vals.items()})
             result.append(row)
         return result