pixeltable 0.2.30__py3-none-any.whl → 0.3.1__py3-none-any.whl

pixeltable/exec/expr_eval_node.py

@@ -1,232 +0,0 @@
-import logging
-import sys
-import time
-import warnings
-from dataclasses import dataclass
-from typing import Iterable, Optional
-
-from tqdm import TqdmWarning, tqdm
-
-from pixeltable import exprs
-from pixeltable.func import CallableFunction
-
-from .data_row_batch import DataRowBatch
-from .exec_node import ExecNode
-
-_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]
-        batched_fn: Optional[CallableFunction]
-        segment_ctxs: list['exprs.RowBuilder.EvalCtx']
-        target_slot_idxs: list[int]
-        batch_size: int = 8
-
-    def __init__(
-            self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr], input_exprs: Iterable[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[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)
-        # This is a temporary hack. When B-tree indices on string columns were implemented (via computed columns
-        # that invoke the `BtreeIndex.str_filter` udf), it resulted in frivolous progress bars appearing on every
-        # insertion. This special-cases the `str_filter` call to suppress the corresponding progress bar.
-        # TODO(aaron-siegel) Remove this hack once we clean up progress bars more generally.
-        is_str_filter_node = all(
-            isinstance(expr, exprs.FunctionCall) and expr.fn.name == 'str_filter' for expr in self.output_exprs
-        )
-        if self.ctx.show_pbar and not is_str_filter_node:
-            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[CallableFunction]:
-        if isinstance(expr, exprs.FunctionCall) and isinstance(expr.fn, CallableFunction) and expr.fn.is_batched:
-            return expr.fn
-        return None
-
-    def _is_batched_fn_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_batched_fn: Optional[CallableFunction] = None
-        for e in all_exprs:
-            if not self._is_batched_fn_call(e):
-                continue
-            assert isinstance(e, exprs.FunctionCall)
-            assert isinstance(e.fn, CallableFunction)
-            if current_batched_fn is None or current_batched_fn != e.fn:
-                # create a new cohort
-                cohorts.append([])
-                current_batched_fn = 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_batched_segment = self._is_batched_fn_call(cohort[0])
-            batched_fn: Optional[CallableFunction] = self._get_batched_fn(cohort[0])
-            for e in cohort[1:]:
-                if self._is_batched_fn_call(e):
-                    segments.append([e])
-                    is_batched_segment = True
-                    batched_fn = self._get_batched_fn(e)
-                else:
-                    if is_batched_segment:
-                        # start a new segment
-                        segments.append([])
-                        is_batched_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, batched_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.batched_fn.get_batch_size() if cohort.batched_fn 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_batched_fn_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]
-                    assert isinstance(fn_call, exprs.FunctionCall)
-                    # make a batched external function call
-                    arg_batches: list[list[exprs.Expr]] = [[] for _ in range(len(fn_call.args))]
-                    kwarg_batches: dict[str, list[exprs.Expr]] = {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)
-                        for i in range(len(args)):
-                            arg_batches[i].append(args[i])
-                        for k in kwargs.keys():
-                            kwarg_batches[k].append(kwargs[k])
-                    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
-                        assert isinstance(fn_call.fn, CallableFunction)
-                        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()
-                        assert isinstance(fn_call.fn, CallableFunction)
-                        result_batch = fn_call.fn.exec_batch(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
-