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

pixeltable/exec/cell_materialization_node.py

@@ -0,0 +1,268 @@
+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_BINARY_SIZE = 512  # max size of binary data 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.supports_file_offloading()
+        }
+        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)
+                    elif col.col_type.is_array_type():
+                        assert isinstance(val, np.ndarray)
+                        self._materialize_array_cell(row, col, val)
+                    else:
+                        assert col.col_type.is_binary_type()
+                        assert isinstance(val, bytes)
+                        self._materialize_binary_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_BINARY_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 _materialize_binary_cell(self, row: exprs.DataRow, col: catalog.Column, val: bytes) -> None:
+        if len(val) <= self.MAX_DB_BINARY_SIZE:
+            # this `bytes` object is small enough to store in the db column (type: binary) directly
+            row.cell_vals[col.id] = val
+            row.cell_md[col.id] = None
+        else:
+            self.init_writer()
+            start = self.buffered_writer.tell()
+            self.buffered_writer.write(val)
+            end = self.buffered_writer.tell()
+            row.cell_vals[col.id] = None
+            cell_md = exprs.CellMd(
+                file_urls=[self.inlined_obj_files[-1].as_uri()], binary_md=exprs.BinaryMd(start=start, end=end)
+            )
+            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, bytes))
+
+    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()}
+        if isinstance(element, bytes):
+            obj_md = self._write_inlined_bytes(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 _write_inlined_bytes(self, data: bytes) -> InlinedObjectMd:
+        """Write raw bytes 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()
+        self.buffered_writer.write(data)
+        end = self.buffered_writer.tell()
+        self._flush_buffer()
+        return InlinedObjectMd(
+            type=ts.ColumnType.Type.BINARY.name, url_idx=url_idx, binary_md=exprs.BinaryMd(start, 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,168 @@
+from __future__ import annotations
+
+import io
+import logging
+from pathlib import Path
+from types import NoneType
+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
+            elif obj_md.type == ts.ColumnType.Type.IMAGE.name:
+                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:
+                assert obj_md.type == ts.ColumnType.Type.BINARY.name
+                assert obj_md.binary_md is not None
+                fp.seek(obj_md.binary_md.start)
+                data = fp.read(obj_md.binary_md.end - obj_md.binary_md.start)
+                assert fp.tell() == obj_md.binary_md.end, (
+                    f'{fp.tell()} != {obj_md.binary_md.end} ({obj_md.binary_md.start})'
+                )
+                return data
+        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]
+    binary_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],
+        binary_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.binary_refs = binary_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
+                        row[col_ref.slot_idx] = self._reconstruct_array(cell_md)
+                    else:
+                        assert isinstance(row[col_ref.slot_idx], (NoneType, np.ndarray))
+
+                for col_ref in self.binary_refs:
+                    cell_md = row.slot_md.get(col_ref.slot_idx)
+                    if cell_md is not None and cell_md.binary_md is not None:
+                        assert row[col_ref.slot_idx] is None
+                        row[col_ref.slot_idx] = self._reconstruct_binary(cell_md)
+                    else:
+                        assert isinstance(row[col_ref.slot_idx], (NoneType, bytes))
+
+            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
+        assert cell_md.file_urls is not None and len(cell_md.file_urls) == 1
+        fp = self.__get_file_pointer(cell_md.file_urls[0])
+        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
+
+    def _reconstruct_binary(self, cell_md: exprs.CellMd) -> bytes:
+        assert cell_md.binary_md is not None
+        assert cell_md.file_urls is not None and len(cell_md.file_urls) == 1
+        fp = self.__get_file_pointer(cell_md.file_urls[0])
+        fp.seek(cell_md.binary_md.start)
+        data = fp.read(cell_md.binary_md.end - cell_md.binary_md.start)
+        assert fp.tell() == cell_md.binary_md.end
+        return data
+
+    def __get_file_pointer(self, file_url: str) -> io.BufferedReader:
+        local_path = parse_local_file_path(file_url)
+        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
+        return self.file_handles[local_path]

pixeltable/exec/component_iteration_node.py

@@ -1,9 +1,6 @@
-import inspect
-from typing import Iterator, Optional
+from typing import AsyncIterator
 
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.exprs as exprs
+from pixeltable import catalog, exceptions as excs, exprs
 
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
@@ -14,34 +11,37 @@ class ComponentIterationNode(ExecNode):
 
     Returns row batches of OUTPUT_BATCH_SIZE size.
     """
+
+    view: catalog.TableVersionHandle
+
     __OUTPUT_BATCH_SIZE = 1024
 
-    def __init__(self, view: catalog.TableVersion, input: ExecNode):
-        assert view.is_component_view()
+    def __init__(self, view: catalog.TableVersionHandle, input: ExecNode):
+        assert view.get().is_component_view
         super().__init__(input.row_builder, [], [], input)
         self.view = view
-        iterator_args = [view.iterator_args.copy()]
+        iterator_args = [view.get().iterator_args.copy()]
         self.row_builder.set_slot_idxs(iterator_args)
         self.iterator_args = iterator_args[0]
         assert isinstance(self.iterator_args, exprs.InlineDict)
         self.iterator_args_ctx = self.row_builder.create_eval_ctx([self.iterator_args])
-        self.iterator_output_schema, self.unstored_column_names = (
-            self.view.iterator_cls.output_schema(**self.iterator_args.to_kwargs())
+        self.iterator_output_schema, self.unstored_column_names = self.view.get().iterator_cls.output_schema(
+            **self.iterator_args.to_kwargs()
         )
         self.iterator_output_fields = list(self.iterator_output_schema.keys())
         self.iterator_output_cols = {
-            field_name: self.view.cols_by_name[field_name] for field_name in self.iterator_output_fields
+            field_name: self.view.get().cols_by_name[field_name] for field_name in self.iterator_output_fields
         }
         # referenced iterator output fields
         self.refd_output_slot_idxs = {
-            e.col.name: e.slot_idx for e in self.row_builder.unique_exprs
+            e.col.name: e.slot_idx
+            for e in self.row_builder.unique_exprs
             if isinstance(e, exprs.ColumnRef) and e.col.name in self.iterator_output_fields
         }
-        self.__output: Optional[Iterator[DataRowBatch]] = None
 
-    def __output_batches(self) -> Iterator[DataRowBatch]:
-        output_batch = DataRowBatch(self.view, self.row_builder)
-        for input_batch in self.input:
+    async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
+        output_batch = DataRowBatch(self.row_builder)
+        async for input_batch in self.input:
             for input_row in input_batch:
                 self.row_builder.eval(input_row, self.iterator_args_ctx)
                 iterator_args = input_row[self.iterator_args.slot_idx]
@@ -50,15 +50,16 @@ class ComponentIterationNode(ExecNode):
                 # specified and are not null. If any of them are null, then we skip this row (i.e., we emit 0
                 # output rows for this input row).
                 if self.__non_nullable_args_specified(iterator_args):
-                    iterator = self.view.iterator_cls(**iterator_args)
+                    iterator = self.view.get().iterator_cls(**iterator_args)
                     for pos, component_dict in enumerate(iterator):
-                        output_row = output_batch.add_row()
+                        output_row = self.row_builder.make_row()
                         input_row.copy(output_row)
                         # we're expanding the input and need to add the iterator position to the pk
                         self.__populate_output_row(output_row, pos, component_dict)
+                        output_batch.add_row(output_row)
                         if len(output_batch) == self.__OUTPUT_BATCH_SIZE:
                             yield output_batch
-                            output_batch = DataRowBatch(self.view, self.row_builder)
+                            output_batch = DataRowBatch(self.row_builder)
 
         if len(output_batch) > 0:
             yield output_batch
@@ -67,7 +68,7 @@ class ComponentIterationNode(ExecNode):
         """
         Returns true if all non-nullable iterator arguments are not `None`.
         """
-        input_schema = self.view.iterator_cls.input_schema()
+        input_schema = self.view.get().iterator_cls.input_schema()
         for arg_name, arg_value in iterator_args.items():
             col_type = input_schema[arg_name]
             if arg_value is None and not col_type.nullable:
@@ -81,7 +82,8 @@ class ComponentIterationNode(ExecNode):
         for field_name, field_val in component_dict.items():
             if field_name not in self.iterator_output_fields:
                 raise excs.Error(
-                    f'Invalid field name {field_name} in output of {self.view.iterator_cls.__name__}')
+                    f'Invalid field name {field_name} in output of {self.view.get().iterator_cls.__name__}'
+                )
             if field_name not in self.refd_output_slot_idxs:
                 # we can ignore this
                 continue
@@ -91,10 +93,5 @@ class ComponentIterationNode(ExecNode):
         if len(component_dict) != len(self.iterator_output_fields):
             missing_fields = set(self.refd_output_slot_idxs.keys()) - set(component_dict.keys())
             raise excs.Error(
-                f'Invalid output of {self.view.iterator_cls.__name__}: '
-                f'missing fields {", ".join(missing_fields)}')
-
-    def __next__(self) -> DataRowBatch:
-        if self.__output is None:
-            self.__output = self.__output_batches()
-        return next(self.__output)
+                f'Invalid output of {self.view.get().iterator_cls.__name__}: missing fields {", ".join(missing_fields)}'
+            )

pixeltable/exec/data_row_batch.py

@@ -1,45 +1,31 @@
 from __future__ import annotations
-from typing import Iterator, Optional
-import logging
 
-import pixeltable.exprs as exprs
-import pixeltable.catalog as catalog
-from pixeltable.utils.media_store import MediaStore
+import logging
+from typing import Iterator
 
+from pixeltable import exprs
 
 _logger = logging.getLogger('pixeltable')
 
+
 class DataRowBatch:
     """Set of DataRows, indexed by rowid.
 
     Contains the metadata needed to initialize DataRows.
+
+    Requires either num_rows or rows to be specified, but not both.
     """
-    tbl: Optional[catalog.TableVersion]
+
     row_builder: exprs.RowBuilder
-    img_slot_idxs: list[int]
-    media_slot_idxs: list[int]  # non-image media slots
-    array_slot_idxs: list[int]
     rows: list[exprs.DataRow]
 
-    def __init__(self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, len: int = 0):
-        self.tbl = tbl
+    def __init__(self, row_builder: exprs.RowBuilder, rows: list[exprs.DataRow] | None = None):
         self.row_builder = row_builder
-        self.img_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_image_type()]
-        # non-image media slots
-        self.media_slot_idxs = [
-            e.slot_idx for e in row_builder.unique_exprs
-            if e.col_type.is_media_type() and not e.col_type.is_image_type()
-        ]
-        self.array_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_array_type()]
-        self.rows = [
-            exprs.DataRow(row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
-            for _ in range(len)
-        ]
+        self.rows = [] if rows is None else rows
 
-    def add_row(self, row: Optional[exprs.DataRow] = None) -> exprs.DataRow:
+    def add_row(self, row: exprs.DataRow | None) -> exprs.DataRow:
         if row is None:
-            row = exprs.DataRow(
-                self.row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)
+            row = self.row_builder.make_row()
         self.rows.append(row)
         return row
 
@@ -52,26 +38,5 @@ class DataRowBatch:
     def __getitem__(self, index: int) -> exprs.DataRow:
         return self.rows[index]
 
-    def flush_imgs(
-            self, idx_range: Optional[slice] = None, stored_img_info: Optional[list[exprs.ColumnSlotIdx]] = None,
-            flushed_slot_idxs: Optional[list[int]] = None
-    ) -> None:
-        """Flushes images in the given range of rows."""
-        assert self.tbl is not None
-        if stored_img_info is None:
-            stored_img_info = []
-        if flushed_slot_idxs is None:
-            flushed_slot_idxs = []
-        if len(stored_img_info) == 0 and len(flushed_slot_idxs) == 0:
-            return
-        if idx_range is None:
-            idx_range = slice(0, len(self.rows))
-        for row in self.rows[idx_range]:
-            for info in stored_img_info:
-                filepath = str(MediaStore.prepare_media_path(self.tbl.id, info.col.id, self.tbl.version))
-                row.flush_img(info.slot_idx, filepath)
-            for slot_idx in flushed_slot_idxs:
-                row.flush_img(slot_idx)
-
     def __iter__(self) -> Iterator[exprs.DataRow]:
         return iter(self.rows)