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

pixeltable/dataframe.py

@@ -23,7 +23,7 @@ from typing import (
 
 import pandas as pd
 import pydantic
-import sqlalchemy as sql
+import sqlalchemy.exc as sql_exc
 
 from pixeltable import catalog, exceptions as excs, exec, exprs, plan, type_system as ts
 from pixeltable.catalog import Catalog, is_valid_identifier
@@ -186,6 +186,8 @@ class DataFrameResultSet:
 
 
 class DataFrame:
+    """Represents a query for retrieving and transforming data from Pixeltable tables."""
+
     _from_clause: plan.FromClause
     _select_list_exprs: list[exprs.Expr]
     _schema: dict[str, ts.ColumnType]
@@ -539,20 +541,23 @@ class DataFrame:
                     yield [data_row[e.slot_idx] for e in self._select_list_exprs]
             except excs.ExprEvalError as e:
                 self._raise_expr_eval_err(e)
-            except sql.exc.DBAPIError as e:
-                raise excs.Error(f'Error during SQL execution:\n{e}') from e
+            except (sql_exc.DBAPIError, sql_exc.OperationalError, sql_exc.InternalError) as e:
+                Catalog.get().convert_sql_exc(e, tbl=(single_tbl.tbl_version if single_tbl is not None else None))
+                raise  # just re-raise if not converted to a Pixeltable error
 
     def collect(self) -> DataFrameResultSet:
         return DataFrameResultSet(list(self._output_row_iterator()), self.schema)
 
     async def _acollect(self) -> DataFrameResultSet:
+        single_tbl = self._first_tbl if len(self._from_clause.tbls) == 1 else None
         try:
             result = [[row[e.slot_idx] for e in self._select_list_exprs] async for row in self._aexec()]
             return DataFrameResultSet(result, self.schema)
         except excs.ExprEvalError as e:
             self._raise_expr_eval_err(e)
-        except sql.exc.DBAPIError as e:
-            raise excs.Error(f'Error during SQL execution:\n{e}') from e
+        except (sql_exc.DBAPIError, sql_exc.OperationalError, sql_exc.InternalError) as e:
+            Catalog.get().convert_sql_exc(e, tbl=(single_tbl.tbl_version if single_tbl is not None else None))
+            raise  # just re-raise if not converted to a Pixeltable error
 
     def count(self) -> int:
         """Return the number of rows in the DataFrame.

pixeltable/env.py

@@ -27,6 +27,7 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
 import nest_asyncio  # type: ignore[import-untyped]
 import pixeltable_pgserver
 import sqlalchemy as sql
+import tzlocal
 from pillow_heif import register_heif_opener  # type: ignore[import-untyped]
 from sqlalchemy import orm
 from tenacity import retry, stop_after_attempt, wait_exponential_jitter
@@ -71,6 +72,7 @@ class Env:
     _db_server: Optional[pixeltable_pgserver.PostgresServer]  # set only when running in local environment
     _db_url: Optional[str]
     _default_time_zone: Optional[ZoneInfo]
+    _verbosity: int
 
     # info about optional packages that are utilized by some parts of the code
     __optional_packages: dict[str, PackageInfo]
@@ -218,10 +220,18 @@ class Env:
         """
         This is not a publicly visible setter; it is only for testing purposes.
         """
-        tz_name = None if tz is None else tz.key
+        if tz is None:
+            tz_name = self._get_tz_name()
+        else:
+            assert isinstance(tz, ZoneInfo)
+            tz_name = tz.key
         self.engine.dispose()
         self._create_engine(time_zone_name=tz_name)
 
+    @property
+    def verbosity(self) -> int:
+        return self._verbosity
+
     @property
     def conn(self) -> Optional[sql.Connection]:
         assert self._current_conn is not None
@@ -237,6 +247,11 @@ class Env:
         assert self._dbms is not None
         return self._dbms
 
+    @property
+    def is_using_cockroachdb(self) -> bool:
+        assert self._dbms is not None
+        return isinstance(self._dbms, CockroachDbms)
+
     @property
     def in_xact(self) -> bool:
         return self._current_conn is not None
@@ -247,7 +262,7 @@ class Env:
         return self._db_server is not None
 
     @contextmanager
-    def begin_xact(self, for_write: bool = False) -> Iterator[sql.Connection]:
+    def begin_xact(self, *, for_write: bool = False) -> Iterator[sql.Connection]:
         """
         Call Catalog.begin_xact() instead, unless there is a specific reason to call this directly.
 
@@ -340,6 +355,8 @@ class Env:
             # accept log messages from a configured pixeltable module (at any level of the module hierarchy)
             path_parts = list(Path(record.pathname).parts)
             path_parts.reverse()
+            if 'pixeltable' not in path_parts:
+                return False
             max_idx = path_parts.index('pixeltable')
             for module_name in path_parts[:max_idx]:
                 if module_name in self._module_log_level and record.levelno >= self._module_log_level[module_name]:
@@ -350,6 +367,26 @@ class Env:
     def console_logger(self) -> ConsoleLogger:
         return self._console_logger
 
+    def _get_tz_name(self) -> str:
+        """Get the time zone name from the configuration, or the system local time zone if not specified.
+
+        Returns:
+            str: The time zone name.
+        """
+        tz_name = Config.get().get_string_value('time_zone')
+        if tz_name is not None:
+            # Validate tzname
+            if not isinstance(tz_name, str):
+                self._logger.error('Invalid time zone specified in configuration.')
+            else:
+                try:
+                    _ = ZoneInfo(tz_name)
+                except ZoneInfoNotFoundError:
+                    self._logger.error(f'Invalid time zone specified in configuration: {tz_name}')
+        else:
+            tz_name = tzlocal.get_localzone_name()
+        return tz_name
+
     def _set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
         if self._initialized:
             return
@@ -393,10 +430,12 @@ class Env:
             warnings.simplefilter('ignore', category=UserWarning)
             warnings.simplefilter('ignore', category=FutureWarning)
 
-        # Set verbose level for user visible console messages
-        verbosity = map_level(config.get_int_value('verbosity'))
+        # Set verbosity level for user visible console messages
+        self._verbosity = config.get_int_value('verbosity')
+        if self._verbosity is None:
+            self._verbosity = 1
         stdout_handler = ConsoleOutputHandler(stream=stdout)
-        stdout_handler.setLevel(verbosity)
+        stdout_handler.setLevel(map_level(self._verbosity))
         stdout_handler.addFilter(ConsoleMessageFilter())
         self._logger.addHandler(stdout_handler)
         self._console_logger = ConsoleLogger(self._logger)
@@ -430,6 +469,7 @@ class Env:
         http_logger.propagate = False
 
         self.clear_tmp_dir()
+        tz_name = self._get_tz_name()
 
         # configure pixeltable database
         self._init_db(config)
@@ -439,22 +479,10 @@ class Env:
                 'Reinitializing pixeltable database is not supported when running in non-local environment'
             )
 
-        tz_name = config.get_string_value('time_zone')
-        if tz_name is not None:
-            # Validate tzname
-            if not isinstance(tz_name, str):
-                self._logger.error('Invalid time zone specified in configuration.')
-            else:
-                try:
-                    _ = ZoneInfo(tz_name)
-                except ZoneInfoNotFoundError:
-                    self._logger.error(f'Invalid time zone specified in configuration: {tz_name}')
-
         if reinit_db and self._store_db_exists():
             self._drop_store_db()
 
         create_db = not self._store_db_exists()
-
         if create_db:
             self._logger.info(f'creating database at: {self.db_url}')
             self._create_store_db()
@@ -534,19 +562,28 @@ class Env:
         metadata.schema.base_metadata.create_all(self._sa_engine, checkfirst=True)
         metadata.create_system_info(self._sa_engine)
 
-    def _create_engine(self, time_zone_name: Optional[str], echo: bool = False) -> None:
-        connect_args = {} if time_zone_name is None else {'options': f'-c timezone={time_zone_name}'}
+    def _create_engine(self, time_zone_name: str, echo: bool = False) -> None:
+        connect_args = {'options': f'-c timezone={time_zone_name}'}
+        self._logger.info(f'Creating SQLAlchemy engine with connection arguments: {connect_args}')
         self._sa_engine = sql.create_engine(
             self.db_url, echo=echo, isolation_level=self._dbms.transaction_isolation_level, connect_args=connect_args
         )
 
         self._logger.info(f'Created SQLAlchemy engine at: {self.db_url}')
+        self._logger.info(f'Engine dialect: {self._sa_engine.dialect.name}')
+        self._logger.info(f'Engine driver : {self._sa_engine.dialect.driver}')
 
         with self.engine.begin() as conn:
             tz_name = conn.execute(sql.text('SHOW TIME ZONE')).scalar()
             assert isinstance(tz_name, str)
             self._logger.info(f'Database time zone is now: {tz_name}')
             self._default_time_zone = ZoneInfo(tz_name)
+            if self.is_using_cockroachdb:
+                # This could be set when the database is created, but we set it now
+                conn.execute(sql.text('SET null_ordered_last = true;'))
+                null_ordered_last = conn.execute(sql.text('SHOW null_ordered_last')).scalar()
+                assert isinstance(null_ordered_last, str)
+                self._logger.info(f'Database null_ordered_last is now: {null_ordered_last}')
 
     def _store_db_exists(self) -> bool:
         assert self._db_name is not None

pixeltable/exec/__init__.py

@@ -2,6 +2,8 @@
 
 from .aggregation_node import AggregationNode
 from .cache_prefetch_node import CachePrefetchNode
+from .cell_materialization_node import CellMaterializationNode
+from .cell_reconstruction_node import CellReconstructionNode
 from .component_iteration_node import ComponentIterationNode
 from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext

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