Repository: ray-project/ray
File: python/ray/data/_internal/execution/streaming_executor_state.py (Python)

Code before:
        ref, diverged = dedupe_schemas_with_validation(
            self._schema,
            ref,
            warn=not self._warned_on_schema_divergence,
            enforce_schemas=self.op.data_context.enforce_schemas,
        )

        # export operator schema when it's updated
        should_export_schema = ref.schema is not None and self._schema != ref.schema
        self._schema = ref.schema
        self._warned_on_schema_divergence |= diverged

        operator_schema_exporter = get_operator_schema_exporter()
        if (
            operator_schema_exporter is not None
            and should_export_schema
            and hasattr(self._schema, "names")
            and hasattr(self._schema, "types")
        ):
            names = [str(n) for n in self._schema.names]
            types = [str(t) for t in self._schema.types]
            operator_schema = OperatorSchema(
                operator_uuid=self.op.id,
                schema_fields=dict(zip(names, types)),
            )
            operator_schema_exporter.export_operator_schema(operator_schema)

        self.output_queue.append(ref)
        self.num_completed_tasks += 1

        actor_info = self.op.get_actor_info()

        self.op.metrics.num_alive_actors = actor_info.running
        self.op.metrics.num_restarting_actors = actor_info.restarting
        self.op.metrics.num_pending_actors = actor_info.pending
        for next_op in self.op.output_dependencies:
            next_op.metrics.num_external_inqueue_blocks += len(ref.blocks)
            next_op.metrics.num_external_inqueue_bytes += ref.size_bytes()
        self.op.metrics.num_external_outqueue_blocks += len(ref.blocks)
        self.op.metrics.num_external_outqueue_bytes += ref.size_bytes()

    def dispatch_next_task(self) -> None:
        """Move a bundle from the operator inqueue to the operator itself."""
        for i, inqueue in enumerate(self.input_queues):
            ref = inqueue.pop()
            if ref is not None:
                self.op.add_input(ref, input_index=i)
                self.op.metrics.num_external_inqueue_bytes -= ref.size_bytes()
                self.op.metrics.num_external_inqueue_blocks -= len(ref.blocks)
                input_op = self.op.input_dependencies[i]
                # TODO: This needs to be cleaned up.

Review comment (refactor):
i'm not sure if this is the right place to put this, but can't think of a long-term solution. Can you think of alternative places to move this to?

Code after:
            enforce_schemas=self.op.data_context.enforce_schemas,
        )

        self._schema = ref.schema
        self._warned_on_schema_divergence |= diverged

        self.output_queue.append(ref)
        self.num_completed_tasks += 1

        actor_info = self.op.get_actor_info()

        self.op.metrics.num_alive_actors = actor_info.running
        self.op.metrics.num_restarting_actors = actor_info.restarting
        self.op.metrics.num_pending_actors = actor_info.pending
        for next_op in self.op.output_dependencies:
            next_op.metrics.num_external_inqueue_blocks += len(ref.blocks)
            next_op.metrics.num_external_inqueue_bytes += ref.size_bytes()
        self.op.metrics.num_external_outqueue_blocks += len(ref.blocks)
        self.op.metrics.num_external_outqueue_bytes += ref.size_bytes()

    def dispatch_next_task(self) -> None:
        """Move a bundle from the operator inqueue to the operator itself."""
        for i, inqueue in enumerate(self.input_queues):
            ref = inqueue.pop()
            if ref is not None:
                self.op.add_input(ref, input_index=i)
                self.op.metrics.num_external_inqueue_bytes -= ref.size_bytes()
                self.op.metrics.num_external_inqueue_blocks -= len(ref.blocks)
                input_op = self.op.input_dependencies[i]
                # TODO: This needs to be cleaned up.
                # the input_op's output queue = curr_op's input queue
                input_op.metrics.num_external_outqueue_blocks -= len(ref.blocks)
                input_op.metrics.num_external_outqueue_bytes -= ref.size_bytes()
                return

        assert False, "Nothing to dispatch"

    def get_output_blocking(self, output_split_idx: Optional[int]) -> RefBundle:
        """Get an item from this node's output queue, blocking as needed.

        Returns:
            The RefBundle from the output queue, or an error / end of stream indicator.

        Raises:
            StopIteration: If all outputs are already consumed.
            Exception: If there was an exception raised during execution.
        """
        while True:
            # Check if StreamingExecutor has caught an exception or is done execution.
            if self._exception is not None:
                raise self._exception