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

Code before:
        return any(len(q) > 0 for q in self.input_queues)

    def total_enqueued_input_blocks_bytes(self) -> int:
        """Total number of bytes occupied by input bundles currently enqueued among:
        1. Input queue(s) pending dispatching (``OpState.input_queues``)
        2. Operator's internal queues (like ``MapOperator``s ref-bundler, etc)
        """
        internal_queue_size_bytes = (
            self.op.internal_input_queue_num_bytes()
            if isinstance(self.op, InternalQueueOperatorMixin)
            else 0
        )
        return self.input_queue_bytes() + internal_queue_size_bytes

    def add_output(self, ref: RefBundle) -> None:
        """Move a bundle produced by the operator to its outqueue."""

        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

Review comment (suggestion):
I think you should be able to use the `diverged` variable, like if `True` then the schema has changed. I checked the implementation of `dedupe_schemas_with_validation` -- you'll need to set `diverged = True` only on `enforce_schemas` to make this work

Code after:
        1. Input queue(s) pending dispatching (``OpState.input_queues``)
        2. Operator's internal queues (like ``MapOperator``s ref-bundler, etc)
        """
        internal_queue_size_bytes = (
            self.op.internal_input_queue_num_bytes()
            if isinstance(self.op, InternalQueueOperatorMixin)
            else 0
        )
        return self.input_queue_bytes() + internal_queue_size_bytes

    def add_output(self, ref: RefBundle) -> None:
        """Move a bundle produced by the operator to its outqueue."""

        ref, diverged = dedupe_schemas_with_validation(
            self._schema,
            ref,
            warn=not self._warned_on_schema_divergence,
            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