class_docstr: pandas-on-Spark specific features for DataFrame.
docstr: Transform chunks with a function that takes pandas DataFrame and outputs pandas DataFrame.
  The pandas DataFrame given to the function is of a batch used internally. The length of
  each input and output should be the same.
  
  See also `Transform and apply a function
  <https://spark.apache.org/docs/latest/api/python/user_guide/pandas_on_spark/transform_apply.html>`_.
  
  .. note:: the `func` is unable to access the whole input frame. pandas-on-Spark
      internally splits the input series into multiple batches and calls `func` with each
      batch multiple times. Therefore, operations such as global aggregations are impossible.
      See the example below.
  
      >>> # This case does not return the length of whole frame but of the batch internally
      ... # used.
      ... def length(pdf) -> ps.DataFrame[int]:
      ...     return pd.DataFrame([len(pdf)] * len(pdf))
      ...
      >>> df = ps.DataFrame({'A': range(1000)})
      >>> df.pandas_on_spark.transform_batch(length)  # doctest: +SKIP
          c0
      0   83
      1   83
      2   83
      ...
  
  .. note:: this API executes the function once to infer the type which is
      potentially expensive, for instance, when the dataset is created after
      aggregations or sorting.
  
      To avoid this, specify return type in ``func``, for instance, as below:
  
      >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]:
      ...     return x + 1
  
      If the return type is specified, the output column names become
      `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
      DataFrame in ``func``.
  
      To specify the column names, you can assign them in a NumPy compound type style
      as below:
  
      >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
      ...     return x + 1
  
      >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
      >>> def plus_one(x) -> ps.DataFrame[
      ...         (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]:
      ...     return x + 1
  
  Parameters
  ----------
  func : function
      Function to transform each pandas frame.
  *args
      Positional arguments to pass to func.
  **kwargs
      Keyword arguments to pass to func.
  
  Returns
  -------
  DataFrame or Series
  
  See Also
  --------
  DataFrame.pandas_on_spark.apply_batch: For row/columnwise operations.
  Series.pandas_on_spark.transform_batch: transform the search as each pandas chunks.
  
  Examples
  --------
  >>> df = ps.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B'])
  >>> df
     A  B
  0  1  2
  1  3  4
  2  5  6
  
  >>> def plus_one_func(pdf) -> ps.DataFrame[int, [int, int]]:
  ...     return pdf + 1
  >>> df.pandas_on_spark.transform_batch(plus_one_func)
     c0  c1
  0   2   3
  1   4   5
  2   6   7
  
  >>> def plus_one_func(pdf) -> ps.DataFrame[("index", int), [('A', int), ('B', int)]]:
  ...     return pdf + 1
  >>> df.pandas_on_spark.transform_batch(plus_one_func)  # doctest: +NORMALIZE_WHITESPACE
         A  B
  index
  0      2  3
  1      4  5
  2      6  7
  
  >>> def plus_one_func(pdf) -> ps.Series[int]:
  ...     return pdf.B + 1
  >>> df.pandas_on_spark.transform_batch(plus_one_func)
  0    3
  1    5
  2    7
  dtype: int64
  
  You can also omit the type hints so pandas-on-Spark infers the return schema as below:
  
  >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf + 1)
     A  B
  0  2  3
  1  4  5
  2  6  7
  
  >>> (df * -1).pandas_on_spark.transform_batch(abs)
     A  B
  0  1  2
  1  3  4
  2  5  6
  
  Note that you should not transform the index. The index information will not change.
  
  >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf.B + 1)
  0    3
  1    5
  2    7
  Name: B, dtype: int64
  
  You can also specify extra arguments as below.
  
  >>> df.pandas_on_spark.transform_batch(lambda pdf, a, b, c: pdf.B + a + b + c, 1, 2, c=3)
  0     8
  1    10
  2    12
  Name: B, dtype: int64
docstr_tok: ["Transform","chunks","with","a","function","that","takes","pandas","DataFrame","and","outputs","pandas","DataFrame",".","The","pandas","DataFrame","given","to","the","function","is","of","a","batch","used","internally",".","The","length","of","each","input","and","output","should","be","the","same",".","See","also","`","Transform","and","apply","a","function","<","https",":","\/\/spark.apache.org\/docs\/latest\/api\/python\/user_guide\/pandas_on_spark\/transform_apply.html",">","`","_",".","..","note",":",":","the","`","func","`","is","unable","to","access","the","whole","input","frame",".","pandas-on-Spark","internally","splits","the","input","series","into","multiple","batches","and","calls","`","func","`","with","each","batch","multiple","times",".","Therefore",",","operations","such","as","global","aggregations","are","impossible",".","See","the","example","below",".",">",">",">","#","This","case","does","not","return","the","length","of","whole","frame","but","of","the","batch","internally","...","#","used",".","...","def","length","(","pdf",")","-",">","ps.DataFrame","[","int","]",":","...","return","pd.DataFrame","(","[","len","(","pdf",")","]","*","len","(","pdf",")",")","...",">",">",">","df","=","ps.DataFrame","(","{","'","A","'",":","range","(","1000",")","}",")",">",">",">","df.pandas_on_spark.transform_batch","(","length",")","#","doctest",":","+SKIP","c0","0","83","1","83","2","83","...","..","note",":",":","this","API","executes","the","function","once","to","infer","the","type","which","is","potentially","expensive",",","for","instance",",","when","the","dataset","is","created","after","aggregations","or","sorting",".","To","avoid","this",",","specify","return","type","in","``","func","``",",","for","instance",",","as","below",":",">",">",">","def","plus_one","(","x",")","-",">","ps.DataFrame","[","int",",","[","float",",","float","]","]",":","...","return","x","+","1","If","the","return","type","is","specified",",","the","output","column","names","become","`","c0",",","c1",",","c2","...","cn","`",".","These","names","are","positionally","mapped","to","the","returned","DataFrame","in","``","func","``",".","To","specify","the","column","names",",","you","can","assign","them","in","a","NumPy","compound","type","style","as","below",":",">",">",">","def","plus_one","(","x",")","-",">","ps.DataFrame","[","(","``","index","''",",","int",")",",","[","(","``","a","''",",","float",")",",","(","``","b","''",",","float",")","]","]",":","...","return","x","+","1",">",">",">","pdf","=","pd.DataFrame","(","{","'","a","'",":","[","1",",","2",",","3","]",",","'","b","'",":","[","3",",","4",",","5","]","}",")",">",">",">","def","plus_one","(","x",")","-",">","ps.DataFrame","[","...","(","pdf.index.name",",","pdf.index.dtype",")",",","zip","(","pdf.dtypes",",","pdf.columns",")","]",":","...","return","x","+","1","Parameters","--","--","--","--","--","func",":","function","Function","to","transform","each","pandas","frame",".","*","args","Positional","arguments","to","pass","to","func",".","*","*","kwargs","Keyword","arguments","to","pass","to","func",".","Returns","--","--","--","-","DataFrame","or","Series","See","Also","--","--","--","--","DataFrame.pandas_on_spark.apply_batch",":","For","row\/columnwise","operations",".","Series.pandas_on_spark.transform_batch",":","transform","the","search","as","each","pandas","chunks",".","Examples","--","--","--","--",">",">",">","df","=","ps.DataFrame","(","[","(","1",",","2",")",",","(","3",",","4",")",",","(","5",",","6",")","]",",","columns=","[","'","A","'",",","'","B","'","]",")",">",">",">","df","A","B","0","1","2","1","3","4","2","5","6",">",">",">","def","plus_one_func","(","pdf",")","-",">","ps.DataFrame","[","int",",","[","int",",","int","]","]",":","...","return","pdf","+","1",">",">",">","df.pandas_on_spark.transform_batch","(","plus_one_func",")","c0","c1","0","2","3","1","4","5","2","6","7",">",">",">","def","plus_one_func","(","pdf",")","-",">","ps.DataFrame","[","(","``","index","''",",","int",")",",","[","(","'","A","'",",","int",")",",","(","'","B","'",",","int",")","]","]",":","...","return","pdf","+","1",">",">",">","df.pandas_on_spark.transform_batch","(","plus_one_func",")","#","doctest",":","+NORMALIZE_WHITESPACE","A","B","index","0","2","3","1","4","5","2","6","7",">",">",">","def","plus_one_func","(","pdf",")","-",">","ps.Series","[","int","]",":","...","return","pdf.B","+","1",">",">",">","df.pandas_on_spark.transform_batch","(","plus_one_func",")","0","3","1","5","2","7","dtype",":","int64","You","can","also","omit","the","type","hints","so","pandas-on-Spark","infers","the","return","schema","as","below",":",">",">",">","df.pandas_on_spark.transform_batch","(","lambda","pdf",":","pdf","+","1",")","A","B","0","2","3","1","4","5","2","6","7",">",">",">","(","df","*","-1",")",".pandas_on_spark.transform_batch","(","abs",")","A","B","0","1","2","1","3","4","2","5","6","Note","that","you","should","not","transform","the","index",".","The","index","information","will","not","change",".",">",">",">","df.pandas_on_spark.transform_batch","(","lambda","pdf",":","pdf.B","+","1",")","0","3","1","5","2","7","Name",":","B",",","dtype",":","int64","You","can","also","specify","extra","arguments","as","below",".",">",">",">","df.pandas_on_spark.transform_batch","(","lambda","pdf",",","a",",","b",",","c",":","pdf.B","+","a","+","b","+","c",",","1",",","2",",","c=3",")","0","8","1","10","2","12","Name",":","B",",","dtype",":","int64"]
code: def transform_batch(
      self,
      func: Callable[..., Union[pd.DataFrame, pd.Series]],
      *args: Any,
      **kwargs: Any,
  ) -> DataFrameOrSeries:
      """
      Transform chunks with a function that takes pandas DataFrame and outputs pandas DataFrame.
      The pandas DataFrame given to the function is of a batch used internally. The length of
      each input and output should be the same.
  
      See also `Transform and apply a function
      <https://spark.apache.org/docs/latest/api/python/user_guide/pandas_on_spark/transform_apply.html>`_.
  
      .. note:: the `func` is unable to access the whole input frame. pandas-on-Spark
          internally splits the input series into multiple batches and calls `func` with each
          batch multiple times. Therefore, operations such as global aggregations are impossible.
          See the example below.
  
          >>> # This case does not return the length of whole frame but of the batch internally
          ... # used.
          ... def length(pdf) -> ps.DataFrame[int]:
          ...     return pd.DataFrame([len(pdf)] * len(pdf))
          ...
          >>> df = ps.DataFrame({'A': range(1000)})
          >>> df.pandas_on_spark.transform_batch(length)  # doctest: +SKIP
              c0
          0   83
          1   83
          2   83
          ...
  
      .. note:: this API executes the function once to infer the type which is
          potentially expensive, for instance, when the dataset is created after
          aggregations or sorting.
  
          To avoid this, specify return type in ``func``, for instance, as below:
  
          >>> def plus_one(x) -> ps.DataFrame[int, [float, float]]:
          ...     return x + 1
  
          If the return type is specified, the output column names become
          `c0, c1, c2 ... cn`. These names are positionally mapped to the returned
          DataFrame in ``func``.
  
          To specify the column names, you can assign them in a NumPy compound type style
          as below:
  
          >>> def plus_one(x) -> ps.DataFrame[("index", int), [("a", float), ("b", float)]]:
          ...     return x + 1
  
          >>> pdf = pd.DataFrame({'a': [1, 2, 3], 'b': [3, 4, 5]})
          >>> def plus_one(x) -> ps.DataFrame[
          ...         (pdf.index.name, pdf.index.dtype), zip(pdf.dtypes, pdf.columns)]:
          ...     return x + 1
  
      Parameters
      ----------
      func : function
          Function to transform each pandas frame.
      *args
          Positional arguments to pass to func.
      **kwargs
          Keyword arguments to pass to func.
  
      Returns
      -------
      DataFrame or Series
  
      See Also
      --------
      DataFrame.pandas_on_spark.apply_batch: For row/columnwise operations.
      Series.pandas_on_spark.transform_batch: transform the search as each pandas chunks.
  
      Examples
      --------
      >>> df = ps.DataFrame([(1, 2), (3, 4), (5, 6)], columns=['A', 'B'])
      >>> df
         A  B
      0  1  2
      1  3  4
      2  5  6
  
      >>> def plus_one_func(pdf) -> ps.DataFrame[int, [int, int]]:
      ...     return pdf + 1
      >>> df.pandas_on_spark.transform_batch(plus_one_func)
         c0  c1
      0   2   3
      1   4   5
      2   6   7
  
      >>> def plus_one_func(pdf) -> ps.DataFrame[("index", int), [('A', int), ('B', int)]]:
      ...     return pdf + 1
      >>> df.pandas_on_spark.transform_batch(plus_one_func)  # doctest: +NORMALIZE_WHITESPACE
             A  B
      index
      0      2  3
      1      4  5
      2      6  7
  
      >>> def plus_one_func(pdf) -> ps.Series[int]:
      ...     return pdf.B + 1
      >>> df.pandas_on_spark.transform_batch(plus_one_func)
      0    3
      1    5
      2    7
      dtype: int64
  
      You can also omit the type hints so pandas-on-Spark infers the return schema as below:
  
      >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf + 1)
         A  B
      0  2  3
      1  4  5
      2  6  7
  
      >>> (df * -1).pandas_on_spark.transform_batch(abs)
         A  B
      0  1  2
      1  3  4
      2  5  6
  
      Note that you should not transform the index. The index information will not change.
  
      >>> df.pandas_on_spark.transform_batch(lambda pdf: pdf.B + 1)
      0    3
      1    5
      2    7
      Name: B, dtype: int64
  
      You can also specify extra arguments as below.
  
      >>> df.pandas_on_spark.transform_batch(lambda pdf, a, b, c: pdf.B + a + b + c, 1, 2, c=3)
      0     8
      1    10
      2    12
      Name: B, dtype: int64
      """
      from pyspark.pandas.groupby import GroupBy
      from pyspark.pandas.frame import DataFrame
      from pyspark.pandas.series import first_series
      from pyspark import pandas as ps
  
      assert callable(
          func
      ), "the first argument should be a callable function."
      spec = inspect.getfullargspec(func)
      return_sig = spec.annotations.get("return", None)
      should_infer_schema = return_sig is None
      should_retain_index = should_infer_schema
      original_func = func
  
      def new_func(o: Any) -> Union[pd.DataFrame, pd.Series]:
          return original_func(o, *args, **kwargs)
  
      def apply_func(pdf: pd.DataFrame) -> pd.DataFrame:
          return new_func(pdf).to_frame()
  
      def pandas_series_func(
          f: Callable[[pd.DataFrame], pd.DataFrame],
          return_type: DataType,
      ) -> "UserDefinedFunctionLike":
          ff = f
  
          @pandas_udf(returnType=return_type)  # type: ignore[call-overload]
          def udf(pdf: pd.DataFrame) -> pd.Series:
              return first_series(ff(pdf))
  
          return udf
  
      if should_infer_schema:
          # Here we execute with the first 1000 to get the return type.
          # If the records were less than 1000, it uses pandas API directly for a shortcut.
          log_advice(
              "If the type hints is not specified for `transform_batch`, "
              "it is expensive to infer the data type internally."
          )
          limit = ps.get_option("compute.shortcut_limit")
          pdf = self._psdf.head(limit + 1)._to_internal_pandas()
          transformed = new_func(pdf)
          if not isinstance(transformed, (pd.DataFrame, pd.Series)):
              raise ValueError(
                  "The given function should return a frame; however, "
                  "the return type was %s." % type(transformed)
              )
          if len(transformed) != len(pdf):
              raise ValueError(
                  "transform_batch cannot produce aggregated results"
              )
          psdf_or_psser = ps.from_pandas(transformed)
  
          if isinstance(psdf_or_psser, ps.Series):
              psser = psdf_or_psser
  
              field = psser._internal.data_fields[
                  0
              ].normalize_spark_type()
  
              return_schema = StructType([field.struct_field])
              output_func = GroupBy._make_pandas_df_builder_func(
                  self._psdf,
                  apply_func,
                  return_schema,
                  retain_index=False,
              )
  
              pudf = pandas_series_func(
                  output_func, return_type=field.spark_type
              )
              columns = self._psdf._internal.spark_columns
              # TODO: Index will be lost in this case.
              internal = self._psdf._internal.copy(
                  column_labels=psser._internal.column_labels,
                  data_spark_columns=[
                      pudf(F.struct(*columns)).alias(field.name)
                  ],
                  data_fields=[field],
                  column_label_names=psser._internal.column_label_names,
              )
              return first_series(DataFrame(internal))
          else:
              psdf = cast(DataFrame, psdf_or_psser)
              if len(pdf) <= limit:
                  # only do the short cut when it returns a frame to avoid
                  # operations on different dataframes in case of series.
                  return psdf
  
              index_fields = [
                  field.normalize_spark_type()
                  for field in psdf._internal.index_fields
              ]
              data_fields = [
                  field.normalize_spark_type()
                  for field in psdf._internal.data_fields
              ]
  
              return_schema = StructType(
                  [
                      field.struct_field
                      for field in index_fields + data_fields
                  ]
              )
  
              self_applied: DataFrame = DataFrame(
                  self._psdf._internal.resolved_copy
              )
  
              output_func = GroupBy._make_pandas_df_builder_func(
                  self_applied,
                  new_func,  # type: ignore[arg-type]
                  return_schema,
                  retain_index=True,
              )
              columns = self_applied._internal.spark_columns
  
              pudf = pandas_udf(  # type: ignore[call-overload]
                  output_func, returnType=return_schema
              )
              temp_struct_column = verify_temp_column_name(
                  self_applied._internal.spark_frame, "__temp_struct__"
              )
              applied = pudf(F.struct(*columns)).alias(
                  temp_struct_column
              )
              sdf = self_applied._internal.spark_frame.select(applied)
              sdf = sdf.selectExpr("%s.*" % temp_struct_column)
  
              return DataFrame(
                  psdf._internal.with_new_sdf(
                      spark_frame=sdf,
                      index_fields=index_fields,
                      data_fields=data_fields,
                  )
              )
      else:
          return_type = infer_return_type(original_func)
          is_return_series = isinstance(return_type, SeriesType)
          is_return_dataframe = isinstance(return_type, DataFrameType)
          if not is_return_dataframe and not is_return_series:
              raise TypeError(
                  "The given function should specify a frame or series as its type "
                  "hints; however, the return type was %s." % return_sig
              )
          if is_return_series:
              field = InternalField(
                  dtype=cast(SeriesType, return_type).dtype,
                  struct_field=StructField(
                      name=SPARK_DEFAULT_SERIES_NAME,
                      dataType=cast(SeriesType, return_type).spark_type,
                  ),
              ).normalize_spark_type()
  
              return_schema = StructType([field.struct_field])
              output_func = GroupBy._make_pandas_df_builder_func(
                  self._psdf,
                  apply_func,
                  return_schema,
                  retain_index=False,
              )
  
              pudf = pandas_series_func(
                  output_func, return_type=field.spark_type
              )
              columns = self._psdf._internal.spark_columns
              internal = self._psdf._internal.copy(
                  column_labels=[None],
                  data_spark_columns=[
                      pudf(F.struct(*columns)).alias(field.name)
                  ],
                  data_fields=[field],
                  column_label_names=None,
              )
              return first_series(DataFrame(internal))
          else:
              index_fields = cast(
                  DataFrameType, return_type
              ).index_fields
              index_fields = [
                  index_field.normalize_spark_type()
                  for index_field in index_fields
              ]
              data_fields = [
                  field.normalize_spark_type()
                  for field in cast(
                      DataFrameType, return_type
                  ).data_fields
              ]
              normalized_fields = index_fields + data_fields
              return_schema = StructType(
                  [field.struct_field for field in normalized_fields]
              )
              should_retain_index = len(index_fields) > 0
  
              self_applied = DataFrame(
                  self._psdf._internal.resolved_copy
              )
  
              output_func = GroupBy._make_pandas_df_builder_func(
                  self_applied,
                  new_func,  # type: ignore[arg-type]
                  return_schema,
                  retain_index=should_retain_index,
              )
              columns = self_applied._internal.spark_columns
  
              pudf = pandas_udf(  # type: ignore[call-overload]
                  output_func, returnType=return_schema
              )
              temp_struct_column = verify_temp_column_name(
                  self_applied._internal.spark_frame, "__temp_struct__"
              )
              applied = pudf(F.struct(*columns)).alias(
                  temp_struct_column
              )
              sdf = self_applied._internal.spark_frame.select(applied)
              sdf = sdf.selectExpr("%s.*" % temp_struct_column)
  
              index_spark_columns = None
              index_names: Optional[
                  List[Optional[Tuple[Any, ...]]]
              ] = None
  
              if should_retain_index:
                  index_spark_columns = [
                      scol_for(sdf, index_field.struct_field.name)
                      for index_field in index_fields
                  ]
  
                  if not any(
                      [
                          SPARK_INDEX_NAME_PATTERN.match(
                              index_field.struct_field.name
                          )
                          for index_field in index_fields
                      ]
                  ):
                      index_names = [
                          (index_field.struct_field.name,)
                          for index_field in index_fields
                      ]
              internal = InternalFrame(
                  spark_frame=sdf,
                  index_names=index_names,
                  index_spark_columns=index_spark_columns,
                  index_fields=index_fields,
                  data_fields=data_fields,
              )
              return DataFrame(internal)
  
imports: import inspect
  from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING, cast, List
  from types import FunctionType
  import numpy
  import pandas
  from pyspark.sql import functions
  from pyspark.sql.functions import pandas_udf
  from pyspark.sql.types import DataType, LongType, StructField, StructType
  from pyspark.pandas._typing import DataFrameOrSeries, Name
  from pyspark.pandas.internal import InternalField, InternalFrame, SPARK_INDEX_NAME_FORMAT, SPARK_DEFAULT_SERIES_NAME, SPARK_INDEX_NAME_PATTERN
  from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType
  from pyspark.pandas.utils import is_name_like_value, is_name_like_tuple, name_like_string, scol_for, verify_temp_column_name, log_advice
label_desc: Use image node_id 4 for calling the PandasOnSparkFrameMethods obj's underlying member method code with example usage: obj.transform_batch(func) and returns: original_func, new_func, udf, first_series, first_series, DataFrame, first_series, DataFrame, psdf