pixeltable 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/exec/cache_prefetch_node.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+from typing import List, Optional, Any, Tuple, Dict
+import threading
+from collections import defaultdict
+from uuid import UUID
+import concurrent
+import logging
+import urllib
+from pathlib import Path
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+import pixeltable.exprs as exprs
+from pixeltable.utils.filecache import FileCache
+import pixeltable.env as env
+import pixeltable.exceptions as excs
+
+_logger = logging.getLogger('pixeltable')
+
+class CachePrefetchNode(ExecNode):
+    """Brings files with external URLs into the cache
+
+    TODO:
+    - maintain a queue of row batches, in order to overlap download and evaluation
+    - adapting the number of download threads at runtime to maximize throughput
+    """
+    def __init__(self, tbl_id: UUID, file_col_info: List[exprs.ColumnSlotIdx], input: ExecNode):
+        # []: we don't have anything to evaluate
+        super().__init__(input.row_builder, [], [], input)
+        self.tbl_id = tbl_id
+        self.file_col_info = file_col_info
+
+        # clients for specific services are constructed as needed, because it's time-consuming
+        self.boto_client: Optional[Any] = None
+        self.boto_client_lock = threading.Lock()
+
+    def __next__(self) -> DataRowBatch:
+        input_batch = next(self.input)
+
+        # collect external URLs that aren't already cached, and set DataRow.file_paths for those that are
+        file_cache = FileCache.get()
+        cache_misses: List[Tuple[exprs.DataRow, exprs.ColumnSlotIdx]] = []
+        missing_url_rows: Dict[str, List[exprs.DataRow]] = defaultdict(list)  # URL -> rows in which it's missing
+        for row in input_batch:
+            for info in self.file_col_info:
+                url = row.file_urls[info.slot_idx]
+                if url is None or row.file_paths[info.slot_idx] is not None:
+                    # nothing to do
+                    continue
+                if url in missing_url_rows:
+                    missing_url_rows[url].append(row)
+                    continue
+                local_path = file_cache.lookup(url)
+                if local_path is None:
+                    cache_misses.append((row, info))
+                    missing_url_rows[url].append(row)
+                else:
+                    row.set_file_path(info.slot_idx, str(local_path))
+
+        # download the cache misses in parallel
+        # TODO: set max_workers to maximize throughput
+        futures: Dict[concurrent.futures.Future, Tuple[exprs.DataRow, exprs.ColumnSlotIdx]] = {}
+        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
+            for row, info in cache_misses:
+                futures[executor.submit(self._fetch_url, row, info.slot_idx)] = (row, info)
+            for future in concurrent.futures.as_completed(futures):
+                # TODO:  does this need to deal with recoverable errors (such as retry after throttling)?
+                tmp_path = future.result()
+                if tmp_path is None:
+                    continue
+                row, info = futures[future]
+                url = row.file_urls[info.slot_idx]
+                local_path = file_cache.add(self.tbl_id, info.col.id, url, tmp_path)
+                _logger.debug(f'PrefetchNode: cached {url} as {local_path}')
+                for row in missing_url_rows[url]:
+                    row.set_file_path(info.slot_idx, str(local_path))
+
+        return input_batch
+
+    def _fetch_url(self, row: exprs.DataRow, slot_idx: int) -> Optional[str]:
+        """Fetches a remote URL into Env.tmp_dir and returns its path"""
+        url = row.file_urls[slot_idx]
+        parsed = urllib.parse.urlparse(url)
+        assert parsed.scheme != '' and parsed.scheme != 'file'
+        # preserve the file extension, if there is one
+        extension = ''
+        if parsed.path != '':
+            p = Path(urllib.parse.unquote(parsed.path))
+            extension = p.suffix
+        tmp_path = env.Env.get().create_tmp_path(extension=extension)
+        try:
+            if parsed.scheme == 's3':
+                from pixeltable.utils.s3 import get_client
+                with self.boto_client_lock:
+                    if self.boto_client is None:
+                        self.boto_client = get_client()
+                    self.boto_client.download_file(parsed.netloc, parsed.path.lstrip('/'), str(tmp_path))
+                    return tmp_path
+            elif parsed.scheme == 'http' or parsed.scheme == 'https':
+                with urllib.request.urlopen(url) as resp, open(tmp_path, 'wb') as f:
+                    data = resp.read()
+                    f.write(data)
+            else:
+                assert False, f'Unsupported URL scheme: {parsed.scheme}'
+            return tmp_path
+        except Exception as e:
+            # we want to add the file url to the exception message
+            exc = excs.Error(f'Failed to download {url}: {e}')
+            self.row_builder.set_exc(row, slot_idx, exc)
+            if not self.ctx.ignore_errors:
+                raise exc from None  # suppress original exception
+        return None
+

pixeltable/exec/component_iteration_node.py

@@ -0,0 +1,79 @@
+from typing import Generator, Optional
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
+import pixeltable.exceptions as excs
+
+
+class ComponentIterationNode(ExecNode):
+    """Expands each row from a base table into one row per component returned by an iterator
+
+    Returns row batches of OUTPUT_BATCH_SIZE size.
+    """
+    OUTPUT_BATCH_SIZE = 1024
+
+    def __init__(self, view: catalog.TableVersion, input: ExecNode):
+        assert view.is_component_view()
+        super().__init__(input.row_builder, [], [], input)
+        self.view = view
+        iterator_args = [view.iterator_args.copy()]
+        self.row_builder.substitute_exprs(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_dict())
+        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}
+        # referenced iterator output fields
+        self.refd_output_slot_idxs = {
+            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[Generator[DataRowBatch, None, None]] = None
+
+    def _output_batches(self) -> Generator[DataRowBatch, None, None]:
+        output_batch = DataRowBatch(self.view, self.row_builder)
+        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]
+                iterator = self.view.iterator_cls(**iterator_args)
+                for pos, component_dict in enumerate(iterator):
+                    output_row = output_batch.add_row()
+                    input_row.copy(output_row)
+                    # we're expanding the input and need to add the iterator position to the pk
+                    pk = output_row.pk[:-1] + (pos,) + output_row.pk[-1:]
+                    output_row.set_pk(pk)
+
+                    # verify and copy component_dict fields to their respective slots in output_row
+                    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__}')
+                        if field_name not in self.refd_output_slot_idxs:
+                            # we can ignore this
+                            continue
+                        output_col = self.iterator_output_cols[field_name]
+                        output_col.col_type.validate_literal(field_val)
+                        output_row[self.refd_output_slot_idxs[field_name]] = field_val
+                    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)}')
+
+                    if len(output_batch) == self.OUTPUT_BATCH_SIZE:
+                        yield output_batch
+                        output_batch = DataRowBatch(self.view, self.row_builder)
+
+        if len(output_batch) > 0:
+            yield output_batch
+
+    def __next__(self) -> DataRowBatch:
+        if self._output is None:
+            self._output = self._output_batches()
+        return next(self._output)

pixeltable/exec/data_row_batch.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+from typing import List, Iterator, Optional
+import logging
+
+import pixeltable.exprs as exprs
+import pixeltable.catalog as catalog
+from pixeltable.utils.media_store import MediaStore
+
+
+_logger = logging.getLogger('pixeltable')
+
+class DataRowBatch:
+    """Set of DataRows, indexed by rowid.
+
+    Contains the metadata needed to initialize DataRows.
+    """
+    def __init__(self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, len: int = 0):
+        self.tbl_id = tbl.id
+        self.tbl_version = tbl.version
+        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)
+        ]
+
+    def add_row(self, row: Optional[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)
+        self.rows.append(row)
+        return row
+
+    def pop_row(self) -> exprs.DataRow:
+        return self.rows.pop()
+
+    def set_row_ids(self, row_ids: List[int]) -> None:
+        """Sets pks for rows in batch"""
+        assert len(row_ids) == len(self.rows)
+        for row, row_id in zip(self.rows, row_ids):
+            row.set_pk((row_id, self.tbl_version))
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, index: object) -> 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."""
+        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)
+        #_logger.debug(
+            #f'flushed images in range {idx_range}: slot_idxs={flushed_slot_idxs} stored_img_info={stored_img_info}')
+
+    def __iter__(self) -> Iterator[exprs.DataRow]:
+        return DataRowBatchIterator(self)
+
+
+class DataRowBatchIterator:
+    """
+    Iterator over a DataRowBatch.
+    """
+    def __init__(self, batch: DataRowBatch):
+        self.row_batch = batch
+        self.index = 0
+
+    def __next__(self) -> exprs.DataRow:
+        if self.index >= len(self.row_batch.rows):
+            raise StopIteration
+        row = self.row_batch.rows[self.index]
+        self.index += 1
+        return row
+

pixeltable/exec/exec_context.py

@@ -0,0 +1,22 @@
+from typing import Optional, List
+
+import sqlalchemy as sql
+
+import pixeltable.exprs as exprs
+
+class ExecContext:
+    """Class for execution runtime constants"""
+    def __init__(
+            self, row_builder: exprs.RowBuilder, *, show_pbar: bool = False, batch_size: int = 0,
+            pk_clause: Optional[List[sql.ClauseElement]] = None, num_computed_exprs: int = 0,
+            ignore_errors: bool = False
+    ):
+        self.show_pbar = show_pbar
+        self.batch_size = batch_size
+        self.profile = exprs.ExecProfile(row_builder)
+        # num_rows is used to compute the total number of computed cells used for the progress bar
+        self.num_rows: Optional[int] = None
+        self.conn: Optional[sql.engine.Connection] = None  # if present, use this to execute SQL queries
+        self.pk_clause = pk_clause
+        self.num_computed_exprs = num_computed_exprs
+        self.ignore_errors = ignore_errors

pixeltable/exec/exec_node.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+from typing import Iterable, Optional, List
+import abc
+
+from .data_row_batch import DataRowBatch
+from .exec_context import ExecContext
+import pixeltable.exprs as exprs
+
+class ExecNode(abc.ABC):
+    """Base class of all execution nodes"""
+    def __init__(
+            self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr],
+            input_exprs: Iterable[exprs.Expr], input: Optional[ExecNode] = None):
+        self.row_builder = row_builder
+        self.input = input
+        # we flush all image slots that aren't part of our output but are needed to create our output
+        output_slot_idxs = {e.slot_idx for e in output_exprs}
+        output_dependencies = row_builder.get_dependencies(output_exprs, exclude=input_exprs)
+        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.stored_img_cols: List[exprs.ColumnSlotIdx] = []
+        self.ctx: Optional[ExecContext] = None  # all nodes of a tree share the same context
+
+    def set_ctx(self, ctx: ExecContext) -> None:
+        self.ctx = ctx
+        if self.input is not None:
+            self.input.set_ctx(ctx)
+
+    def set_stored_img_cols(self, stored_img_cols: List[exprs.ColumnSlotIdx]) -> None:
+        self.stored_img_cols = stored_img_cols
+        # propagate batch size to the source
+        if self.input is not None:
+            self.input.set_stored_img_cols(stored_img_cols)
+
+    def __iter__(self):
+        return self
+
+    @abc.abstractmethod
+    def __next__(self) -> DataRowBatch:
+        pass
+
+    def open(self) -> None:
+        """Bottom-up initialization of nodes for execution. Must be called before __next__."""
+        if self.input is not None:
+            self.input.open()
+        self._open()
+
+    def close(self) -> None:
+        """Frees node resources top-down after execution. Must be called after final __next__."""
+        self._close()
+        if self.input is not None:
+            self.input.close()
+
+    def _open(self) -> None:
+        pass
+
+    def _close(self) -> None:
+        pass
+

pixeltable/exec/expr_eval_node.py

@@ -0,0 +1,217 @@
+import sys
+import warnings
+from typing import List, Optional, Tuple
+from dataclasses import dataclass, field
+import logging
+import time
+
+from tqdm import tqdm, TqdmWarning
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+import pixeltable.exprs as exprs
+import pixeltable.func as func
+
+
+_logger = logging.getLogger('pixeltable')
+
+class ExprEvalNode(ExecNode):
+    """Materializes expressions
+    """
+    @dataclass
+    class Cohort:
+        """List of exprs that form an evaluation context and contain calls to at most one external function"""
+        exprs: List[exprs.Expr]
+        ext_function: Optional[func.BatchedFunction]
+        segment_ctxs: List[exprs.RowBuilder.EvalCtx]
+        target_slot_idxs: List[int]
+        batch_size: int = 8
+
+    def __init__(
+            self, row_builder: exprs.RowBuilder, output_exprs: List[exprs.Expr], input_exprs: List[exprs.Expr],
+            input: ExecNode
+    ):
+        super().__init__(row_builder, output_exprs, input_exprs, input)
+        self.input_exprs = input_exprs
+        input_slot_idxs = {e.slot_idx for e in input_exprs}
+        # we're only materializing exprs that are not already in the input
+        self.target_exprs = [e for e in output_exprs if e.slot_idx not in input_slot_idxs]
+        self.pbar: Optional[tqdm] = None
+        self.cohorts: List[List[ExprEvalNode.Cohort]] = []
+        self._create_cohorts()
+
+    def __next__(self) -> DataRowBatch:
+        input_batch = next(self.input)
+        # compute target exprs
+        for cohort in self.cohorts:
+            self._exec_cohort(cohort, input_batch)
+        _logger.debug(f'ExprEvalNode: returning {len(input_batch)} rows')
+        return input_batch
+
+    def _open(self) -> None:
+        warnings.simplefilter("ignore", category=TqdmWarning)
+        if self.ctx.show_pbar:
+            self.pbar = tqdm(
+                total=len(self.target_exprs) * self.ctx.num_rows,
+                desc='Computing cells',
+                unit=' cells',
+                ncols=100,
+                file=sys.stdout
+            )
+
+    def _close(self) -> None:
+        if self.pbar is not None:
+            self.pbar.close()
+
+    def _get_batched_fn(self, expr: exprs.Expr) -> Optional[func.BatchedFunction]:
+        if not isinstance(expr, exprs.FunctionCall):
+            return None
+        return expr.fn if isinstance(expr.fn, func.BatchedFunction) else None
+
+    def _is_ext_call(self, expr: exprs.Expr) -> bool:
+        return self._get_batched_fn(expr) is not None
+
+    def _create_cohorts(self) -> None:
+        all_exprs = self.row_builder.get_dependencies(self.target_exprs)
+        # break up all_exprs into cohorts such that each cohort contains calls to at most one external function;
+        # seed the cohorts with only the ext fn calls
+        cohorts: List[List[exprs.Expr]] = []
+        current_ext_function: Optional[func.BatchedFunction] = None
+        for e in all_exprs:
+            if not self._is_ext_call(e):
+                continue
+            if current_ext_function is None or current_ext_function != e.fn:
+                # create a new cohort
+                cohorts.append([])
+                current_ext_function = e.fn
+            cohorts[-1].append(e)
+
+        # expand the cohorts to include all exprs that are in the same evaluation context as the external calls;
+        # cohorts are evaluated in order, so we can exclude the target slots from preceding cohorts and input slots
+        exclude = set([e.slot_idx for e in self.input_exprs])
+        all_target_slot_idxs = set([e.slot_idx for e in self.target_exprs])
+        target_slot_idxs: List[List[int]] = []  # the ones materialized by each cohort
+        for i in range(len(cohorts)):
+            cohorts[i] = self.row_builder.get_dependencies(
+                cohorts[i], exclude=[self.row_builder.unique_exprs[slot_idx] for slot_idx in exclude])
+            target_slot_idxs.append(
+                [e.slot_idx for e in cohorts[i] if e.slot_idx in all_target_slot_idxs])
+            exclude.update(target_slot_idxs[-1])
+
+        all_cohort_slot_idxs = set([e.slot_idx for cohort in cohorts for e in cohort])
+        remaining_slot_idxs = set(all_target_slot_idxs) - all_cohort_slot_idxs
+        if len(remaining_slot_idxs) > 0:
+            cohorts.append(self.row_builder.get_dependencies(
+                [self.row_builder.unique_exprs[slot_idx] for slot_idx in remaining_slot_idxs],
+                exclude=[self.row_builder.unique_exprs[slot_idx] for slot_idx in exclude]))
+            target_slot_idxs.append(list(remaining_slot_idxs))
+        # we need to have captured all target slots at this point
+        assert all_target_slot_idxs == set().union(*target_slot_idxs)
+
+        for i in range(len(cohorts)):
+            cohort = cohorts[i]
+            # segment the cohort into sublists that contain either a single ext. function call or no ext. function calls
+            # (i.e., only computed cols)
+            assert len(cohort) > 0
+            # create the first segment here, so we can avoid checking for an empty list in the loop
+            segments = [[cohort[0]]]
+            is_ext_segment = self._is_ext_call(cohort[0])
+            ext_fn: Optional[func.BatchedFunction] = self._get_batched_fn(cohort[0])
+            for e in cohort[1:]:
+                if self._is_ext_call(e):
+                    segments.append([e])
+                    is_ext_segment = True
+                    ext_fn = self._get_batched_fn(e)
+                else:
+                    if is_ext_segment:
+                        # start a new segment
+                        segments.append([])
+                        is_ext_segment = False
+                    segments[-1].append(e)
+
+            # we create the EvalCtxs manually because create_eval_ctx() would repeat the dependencies of each segment
+            segment_ctxs = [
+                exprs.RowBuilder.EvalCtx(
+                    slot_idxs=[e.slot_idx for e in s], exprs=s, target_slot_idxs=[], target_exprs=[])
+                for s in segments
+            ]
+            cohort_info = self.Cohort(cohort, ext_fn, segment_ctxs, target_slot_idxs[i])
+            self.cohorts.append(cohort_info)
+
+    def _exec_cohort(self, cohort: Cohort, rows: DataRowBatch) -> None:
+        """Compute the cohort for the entire input batch by dividing it up into sub-batches"""
+        batch_start_idx = 0  # start row of the current sub-batch
+        # for multi-resolution models, we re-assess the correct ext fn batch size for each input batch
+        ext_batch_size = cohort.ext_function.get_batch_size() if cohort.ext_function is not None else None
+        if ext_batch_size is not None:
+            cohort.batch_size = ext_batch_size
+
+        while batch_start_idx < len(rows):
+            num_batch_rows = min(cohort.batch_size, len(rows) - batch_start_idx)
+            for segment_ctx in cohort.segment_ctxs:
+                if not self._is_ext_call(segment_ctx.exprs[0]):
+                    # compute batch row-wise
+                    for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
+                        self.row_builder.eval(
+                            rows[row_idx], segment_ctx, self.ctx.profile, ignore_errors=self.ctx.ignore_errors)
+                else:
+                    fn_call = segment_ctx.exprs[0]
+                    # make a batched external function call
+                    arg_batches = [[] for _ in range(len(fn_call.args))]
+                    kwarg_batches = {k: [] for k in fn_call.kwargs.keys()}
+
+                    valid_batch_idxs: List[int] = []  # rows with exceptions are not valid
+                    for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
+                        row = rows[row_idx]
+                        if row.has_exc(fn_call.slot_idx):
+                            # one of our inputs had an exception, skip this row
+                            continue
+                        valid_batch_idxs.append(row_idx)
+                        args, kwargs = fn_call._make_args(row)
+                        [arg_batches[i].append(args[i]) for i in range(len(args))]
+                        [kwarg_batches[k].append(kwargs[k]) for k in kwargs.keys()]
+                    num_valid_batch_rows = len(valid_batch_idxs)
+
+                    if ext_batch_size is None:
+                        # we need to choose a batch size based on the args
+                        sample_args = [arg_batches[i][0] for i in range(len(arg_batches))]
+                        ext_batch_size = fn_call.fn.get_batch_size(*sample_args)
+
+                    num_remaining_batch_rows = num_valid_batch_rows
+                    while num_remaining_batch_rows > 0:
+                        # we make ext. fn calls in batches of ext_batch_size
+                        if ext_batch_size is None:
+                            pass
+                        num_ext_batch_rows = min(ext_batch_size, num_remaining_batch_rows)
+                        ext_batch_offset = num_valid_batch_rows - num_remaining_batch_rows  # offset into args, not rows
+                        call_args = [
+                            arg_batches[i][ext_batch_offset:ext_batch_offset + num_ext_batch_rows]
+                            for i in range(len(arg_batches))
+                        ]
+                        call_kwargs = {
+                            k: kwarg_batches[k][ext_batch_offset:ext_batch_offset + num_ext_batch_rows]
+                            for k in kwarg_batches.keys()
+                        }
+                        start_ts = time.perf_counter()
+                        result_batch = fn_call.fn.invoke(call_args, call_kwargs)
+                        self.ctx.profile.eval_time[fn_call.slot_idx] += time.perf_counter() - start_ts
+                        self.ctx.profile.eval_count[fn_call.slot_idx] += num_ext_batch_rows
+
+                        # move the result into the row batch
+                        for result_idx in range(len(result_batch)):
+                            row_idx = valid_batch_idxs[ext_batch_offset + result_idx]
+                            row = rows[row_idx]
+                            row[fn_call.slot_idx] = result_batch[result_idx]
+
+                        num_remaining_batch_rows -= num_ext_batch_rows
+
+                    # switch to the ext fn batch size
+                    cohort.batch_size = ext_batch_size
+
+            # make sure images for stored cols have been saved to files before moving on to the next batch
+            rows.flush_imgs(
+                slice(batch_start_idx, batch_start_idx + num_batch_rows), self.stored_img_cols, self.flushed_img_slots)
+            if self.pbar is not None:
+                self.pbar.update(num_batch_rows * len(cohort.target_slot_idxs))
+            batch_start_idx += num_batch_rows
+

pixeltable/exec/in_memory_data_node.py

@@ -0,0 +1,69 @@
+from typing import List, Dict, Any, Optional
+import urllib
+import logging
+import os
+
+from .data_row_batch import DataRowBatch
+from .exec_node import ExecNode
+import pixeltable.catalog as catalog
+import pixeltable.exprs as exprs
+import pixeltable.env as env
+from pixeltable.utils.media_store import MediaStore
+
+
+_logger = logging.getLogger('pixeltable')
+
+class InMemoryDataNode(ExecNode):
+    """Outputs in-memory data as a row batch of a particular table"""
+    def __init__(
+            self, tbl: catalog.TableVersionPath, rows: List[Dict[str, Any]],
+            row_builder: exprs.RowBuilder, start_row_id: int,
+    ):
+        super().__init__(row_builder, [], [], None)
+        assert tbl.is_insertable()
+        self.tbl = tbl
+        self.input_rows = rows
+        self.start_row_id = start_row_id
+        self.has_returned_data = False
+        self.output_rows: Optional[DataRowBatch] = None
+
+    def _open(self) -> None:
+        """Create row batch and populate with self.input_rows"""
+        column_info = {info.col.name: info for info in self.row_builder.output_slot_idxs()}
+        # stored columns that are not computed
+        inserted_column_names = set([
+            info.col.name for info in self.row_builder.output_slot_idxs()
+            if info.col.is_stored and not info.col.is_computed
+        ])
+
+        self.output_rows = DataRowBatch(self.tbl, self.row_builder, len(self.input_rows))
+        for row_idx, input_row in enumerate(self.input_rows):
+            # populate the output row with the values provided in the input row
+            for col_name, val in input_row.items():
+                col_info = column_info.get(col_name)
+                assert col_info is not None
+
+                if col_info.col.col_type.is_image_type() and isinstance(val, bytes):
+                    # this is a literal image, ie, a sequence of bytes; we save this as a media file and store the path
+                    path = str(MediaStore.prepare_media_path(self.tbl.id, col_info.col.id, self.tbl.version))
+                    open(path, 'wb').write(val)
+                    val = path
+                self.output_rows[row_idx][col_info.slot_idx] = val
+
+            # set the remaining stored non-computed columns to null
+            null_col_names = inserted_column_names - set(input_row.keys())
+            for col_name in null_col_names:
+                col_info = column_info.get(col_name)
+                assert col_info is not None
+                self.output_rows[row_idx][col_info.slot_idx] = None
+
+        self.output_rows.set_row_ids([self.start_row_id + i for i in range(len(self.output_rows))])
+        self.ctx.num_rows = len(self.output_rows)
+
+    def __next__(self) -> DataRowBatch:
+        if self.has_returned_data:
+            raise StopIteration
+        self.has_returned_data = True
+        _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_rows)} output_rows')
+        return self.output_rows
+