Method write_datasink

python/ray/data/dataset.py:5853–5937

Writes the dataset to a custom :class:`~ray.data.Datasink`. Time complexity: O(dataset size / parallelism) Args: datasink: The :class:`~ray.data.Datasink` to write to. ray_remote_args: Kwargs passed to :func:`ray.remote` in the write tasks. concu

(
        self,
        datasink: Datasink,
        *,
        ray_remote_args: Dict[str, Any] = None,
        concurrency: Optional[int] = None,
    )

    @ConsumptionAPI(pattern="Time complexity:")
    def write_datasink(
        self,
        datasink: Datasink,
        *,
        ray_remote_args: Dict[str, Any] = None,
        concurrency: Optional[int] = None,
    ) -> None:
        """Writes the dataset to a custom :class:`~ray.data.Datasink`.

        Time complexity: O(dataset size / parallelism)

        Args:
            datasink: The :class:`~ray.data.Datasink` to write to.
            ray_remote_args: Kwargs passed to :func:`ray.remote` in the write tasks.
            concurrency: The maximum number of Ray tasks to run concurrently. Set this
                to control number of tasks to run concurrently. This doesn't change the
                total number of tasks run. By default, concurrency is dynamically
                decided based on the available resources.
        """  # noqa: E501
        if ray_remote_args is None:
            ray_remote_args = {}

        if not datasink.supports_distributed_writes:
            if ray.util.client.ray.is_connected():
                raise ValueError(
                    "If you're using Ray Client, Ray Data won't schedule write tasks "
                    "on the driver's node."
                )
            label_selector = ray_remote_args.get("label_selector", {})
            label_selector[
                ray._raylet.RAY_NODE_ID_KEY
            ] = ray.get_runtime_context().get_node_id()
            ray_remote_args["label_selector"] = label_selector
            ray_remote_args.pop("scheduling_strategy", None)

            _validate_head_node_resources_for_local_scheduling(
                ray_remote_args,
                op_description="Writing to a local:// path",
            )

        write_op = Write(
            datasink,
            input_dependencies=[self._logical_plan.dag],
            ray_remote_args=ray_remote_args,
            compute=TaskPoolStrategy(concurrency),
        )
        logical_plan = LogicalPlan(write_op, self.context)

        try:
            # Call on_write_start for _FileDatasink before execution to handle
            # SaveMode checks (ERROR raises, OVERWRITE deletes contents, IGNORE skips)
            # and directory creation. For other datasinks, on_write_start is called
            # automatically by the Write operator when the first input bundle arrives.
            if isinstance(datasink, _FileDatasink):
                datasink.on_write_start()
                # TODO (https://github.com/ray-project/ray/issues/59326): There should be no special handling for skipping writes.
                if datasink._skip_write:
                    logger.info(