code: def generate_write_fn(
      datasink_or_legacy_datasource: Union[Datasink, Datasource],
      **write_args,
  ) -> Callable[[Iterator[Block], TaskContext], Iterator[Block]]:
      # If the write op succeeds, the resulting Dataset is a list of
      # WriteResult (one element per write task). Otherwise, an error will
      # be raised. The Datasource can handle execution outcomes with the
      # on_write_complete() and on_write_failed().
      def fn(blocks: Iterator[Block], ctx) -> Iterator[Block]:
          if isinstance(datasink_or_legacy_datasource, Datasink):
              write_result = datasink_or_legacy_datasource.write(
                  blocks, ctx
              )
          else:
              write_result = datasink_or_legacy_datasource.write(
                  blocks, ctx, **write_args
              )
  
          # NOTE: `WriteResult` isn't a valid block type, so we need to wrap it up.
          import pandas as pd
  
          block = pd.DataFrame({"write_result": [write_result]})
          return [block]
  
      return fn
  
file_path: ray/python/ray/data/_internal/planner/plan_write_op.py
imports: from typing import Callable, Iterator, Union
  from ray.data._internal.execution.interfaces import PhysicalOperator
  from ray.data._internal.execution.interfaces.task_context import TaskContext
  from ray.data._internal.execution.operators.map_operator import MapOperator
  from ray.data._internal.execution.operators.map_transformer import BlockMapTransformFn, MapTransformer
  from ray.data._internal.logical.operators.write_operator import Write
  from ray.data.block import Block
  from ray.data.datasource.datasink import Datasink
  from ray.data.datasource.datasource import Datasource
label_desc: Use image node_id 1 for calling a global function with example usage: generate_write_fn(datasink_or_legacy_datasource) and returns: fn, list