class_docstr: pandas-on-Spark specific features for DataFrame.
docstr: Apply a function that takes pandas DataFrame and outputs pandas DataFrame. The pandas
  DataFrame given to the function is of a batch used internally.
  
  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, [int]]:
      ...     return pd.DataFrame([len(pdf)])
      ...
      >>> df = ps.DataFrame({'A': range(1000)})
      >>> df.pandas_on_spark.apply_batch(length)  # doctest: +SKIP
          c0
      0   83
      1   83
      2   83
      ...
      10  83
      11  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 apply to each pandas frame.
  args : tuple
      Positional arguments to pass to `func` in addition to the
      array/series.
  **kwds
      Additional keyword arguments to pass as keywords arguments to
      `func`.
  
  Returns
  -------
  DataFrame
  
  See Also
  --------
  DataFrame.apply: For row/columnwise operations.
  DataFrame.applymap: For elementwise operations.
  DataFrame.aggregate: Only perform aggregating type operations.
  DataFrame.transform: Only perform transforming type 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 query_func(pdf) -> ps.DataFrame[int, [int, int]]:
  ...     return pdf.query('A == 1')
  >>> df.pandas_on_spark.apply_batch(query_func)
     c0  c1
  0   1   2
  
  >>> def query_func(pdf) -> ps.DataFrame[("idx", int), [("A", int), ("B", int)]]:
  ...     return pdf.query('A == 1')
  >>> df.pandas_on_spark.apply_batch(query_func)  # doctest: +NORMALIZE_WHITESPACE
       A  B
  idx
  0    1  2
  
  You can also omit the type hints so pandas-on-Spark infers the return schema as below:
  
  >>> df.pandas_on_spark.apply_batch(lambda pdf: pdf.query('A == 1'))
     A  B
  0  1  2
  
  You can also specify extra arguments.
  
  >>> def calculation(pdf, y, z) -> ps.DataFrame[int, [int, int]]:
  ...     return pdf ** y + z
  >>> df.pandas_on_spark.apply_batch(calculation, args=(10,), z=20)
          c0        c1
  0       21      1044
  1    59069   1048596
  2  9765645  60466196
  
  You can also use ``np.ufunc`` and built-in functions as input.
  
  >>> df.pandas_on_spark.apply_batch(np.add, args=(10,))
      A   B
  0  11  12
  1  13  14
  2  15  16
  
  >>> (df * -1).pandas_on_spark.apply_batch(abs)
     A  B
  0  1  2
  1  3  4
  2  5  6
docstr_tok: ["Apply","a","function","that","takes","pandas","DataFrame","and","outputs","pandas","DataFrame",".","The","pandas","DataFrame","given","to","the","function","is","of","a","batch","used","internally",".","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",",","[","int","]","]",":","...","return","pd.DataFrame","(","[","len","(","pdf",")","]",")","...",">",">",">","df","=","ps.DataFrame","(","{","'","A","'",":","range","(","1000",")","}",")",">",">",">","df.pandas_on_spark.apply_batch","(","length",")","#","doctest",":","+SKIP","c0","0","83","1","83","2","83","...","10","83","11","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","apply","to","each","pandas","frame",".","args",":","tuple","Positional","arguments","to","pass","to","`","func","`","in","addition","to","the","array\/series",".","*","*","kwds","Additional","keyword","arguments","to","pass","as","keywords","arguments","to","`","func","`",".","Returns","--","--","--","-","DataFrame","See","Also","--","--","--","--","DataFrame.apply",":","For","row\/columnwise","operations",".","DataFrame.applymap",":","For","elementwise","operations",".","DataFrame.aggregate",":","Only","perform","aggregating","type","operations",".","DataFrame.transform",":","Only","perform","transforming","type","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","query_func","(","pdf",")","-",">","ps.DataFrame","[","int",",","[","int",",","int","]","]",":","...","return","pdf.query","(","'","A","==","1","'",")",">",">",">","df.pandas_on_spark.apply_batch","(","query_func",")","c0","c1","0","1","2",">",">",">","def","query_func","(","pdf",")","-",">","ps.DataFrame","[","(","``","idx","''",",","int",")",",","[","(","``","A","''",",","int",")",",","(","``","B","''",",","int",")","]","]",":","...","return","pdf.query","(","'","A","==","1","'",")",">",">",">","df.pandas_on_spark.apply_batch","(","query_func",")","#","doctest",":","+NORMALIZE_WHITESPACE","A","B","idx","0","1","2","You","can","also","omit","the","type","hints","so","pandas-on-Spark","infers","the","return","schema","as","below",":",">",">",">","df.pandas_on_spark.apply_batch","(","lambda","pdf",":","pdf.query","(","'","A","==","1","'",")",")","A","B","0","1","2","You","can","also","specify","extra","arguments",".",">",">",">","def","calculation","(","pdf",",","y",",","z",")","-",">","ps.DataFrame","[","int",",","[","int",",","int","]","]",":","...","return","pdf","*","*","y","+","z",">",">",">","df.pandas_on_spark.apply_batch","(","calculation",",","args=","(","10",",",")",",","z=20",")","c0","c1","0","21","1044","1","59069","1048596","2","9765645","60466196","You","can","also","use","``","np.ufunc","``","and","built-in","functions","as","input",".",">",">",">","df.pandas_on_spark.apply_batch","(","np.add",",","args=","(","10",",",")",")","A","B","0","11","12","1","13","14","2","15","16",">",">",">","(","df","*","-1",")",".pandas_on_spark.apply_batch","(","abs",")","A","B","0","1","2","1","3","4","2","5","6"]
code: def apply_batch(
      self,
      func: Callable[..., pd.DataFrame],
      args: Tuple = (),
      **kwds: Any,
  ) -> "DataFrame":
      """
      Apply a function that takes pandas DataFrame and outputs pandas DataFrame. The pandas
      DataFrame given to the function is of a batch used internally.
  
      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, [int]]:
          ...     return pd.DataFrame([len(pdf)])
          ...
          >>> df = ps.DataFrame({'A': range(1000)})
          >>> df.pandas_on_spark.apply_batch(length)  # doctest: +SKIP
              c0
          0   83
          1   83
          2   83
          ...
          10  83
          11  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 apply to each pandas frame.
      args : tuple
          Positional arguments to pass to `func` in addition to the
          array/series.
      **kwds
          Additional keyword arguments to pass as keywords arguments to
          `func`.
  
      Returns
      -------
      DataFrame
  
      See Also
      --------
      DataFrame.apply: For row/columnwise operations.
      DataFrame.applymap: For elementwise operations.
      DataFrame.aggregate: Only perform aggregating type operations.
      DataFrame.transform: Only perform transforming type 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 query_func(pdf) -> ps.DataFrame[int, [int, int]]:
      ...     return pdf.query('A == 1')
      >>> df.pandas_on_spark.apply_batch(query_func)
         c0  c1
      0   1   2
  
      >>> def query_func(pdf) -> ps.DataFrame[("idx", int), [("A", int), ("B", int)]]:
      ...     return pdf.query('A == 1')
      >>> df.pandas_on_spark.apply_batch(query_func)  # doctest: +NORMALIZE_WHITESPACE
           A  B
      idx
      0    1  2
  
      You can also omit the type hints so pandas-on-Spark infers the return schema as below:
  
      >>> df.pandas_on_spark.apply_batch(lambda pdf: pdf.query('A == 1'))
         A  B
      0  1  2
  
      You can also specify extra arguments.
  
      >>> def calculation(pdf, y, z) -> ps.DataFrame[int, [int, int]]:
      ...     return pdf ** y + z
      >>> df.pandas_on_spark.apply_batch(calculation, args=(10,), z=20)
              c0        c1
      0       21      1044
      1    59069   1048596
      2  9765645  60466196
  
      You can also use ``np.ufunc`` and built-in functions as input.
  
      >>> df.pandas_on_spark.apply_batch(np.add, args=(10,))
          A   B
      0  11  12
      1  13  14
      2  15  16
  
      >>> (df * -1).pandas_on_spark.apply_batch(abs)
         A  B
      0  1  2
      1  3  4
      2  5  6
  
      """
      # TODO: codes here partially duplicate `DataFrame.apply`. Can we deduplicate?
  
      from pyspark.pandas.groupby import GroupBy
      from pyspark.pandas.frame import DataFrame
      from pyspark import pandas as ps
  
      if not isinstance(func, FunctionType):
          assert callable(
              func
          ), "the first argument should be a callable function."
          f = func
          # Note that the return type hint specified here affects actual return
          # type in Spark (e.g., infer_return_type). And, MyPy does not allow
          # redefinition of a function.
          func = lambda *args, **kwargs: f(
              *args, **kwargs
          )  # noqa: E731
  
      spec = inspect.getfullargspec(func)
      return_sig = spec.annotations.get("return", None)
      should_infer_schema = return_sig is None
  
      original_func = func
  
      def new_func(o: Any) -> pd.DataFrame:
          return original_func(o, *args, **kwds)
  
      self_applied: DataFrame = DataFrame(
          self._psdf._internal.resolved_copy
      )
  
      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 `apply_batch`, "
              "it is expensive to infer the data type internally."
          )
          limit = ps.get_option("compute.shortcut_limit")
          pdf = self_applied.head(limit + 1)._to_internal_pandas()
          applied = new_func(pdf)
          if not isinstance(applied, pd.DataFrame):
              raise ValueError(
                  "The given function should return a frame; however, "
                  "the return type was %s." % type(applied)
              )
          psdf: DataFrame = DataFrame(applied)
          if len(pdf) <= limit:
              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
              ]
          )
  
          output_func = GroupBy._make_pandas_df_builder_func(
              self_applied, new_func, return_schema, retain_index=True
          )
          sdf = self_applied._internal.spark_frame.mapInPandas(
              lambda iterator: map(output_func, iterator),
              schema=return_schema,
          )
  
          # If schema is inferred, we can restore indexes too.
          internal = 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_dataframe = isinstance(return_type, DataFrameType)
          if not is_return_dataframe:
              raise TypeError(
                  "The given function should specify a frame as its type "
                  "hints; however, the return type was %s." % return_sig
              )
          index_fields = cast(DataFrameType, return_type).index_fields
          should_retain_index = len(index_fields) > 0
          return_schema = cast(DataFrameType, return_type).spark_type
  
          output_func = GroupBy._make_pandas_df_builder_func(
              self_applied,
              new_func,
              return_schema,
              retain_index=should_retain_index,
          )
          sdf = self_applied._internal.to_internal_spark_frame.mapInPandas(
              lambda iterator: map(output_func, iterator),
              schema=return_schema,
          )
  
          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=cast(DataFrameType, return_type).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 3 for calling the PandasOnSparkFrameMethods obj's underlying member method code with example usage: obj.apply_batch(func, args) and returns: DataFrame, original_func, psdf