code: def _split_dataset_spark(
      dataset,
      split_type: Literal["n_instances", "fraction", "condition"],
      split: int | float | pd.Series | str,
      time_axis: str,
      partition_cols: list[str] | None = None,
      persist: bool = True,
  ):
      """
      Split a dataset into a reference and remaining part based on split params using PySpark
      (see `split_dataset`)
  
      :param split_type: meaning of split parameter
      :param split: split parameter
      :param time_axis: time axis
      :param partition_cols: cols to partition by (for performance)
      :param persist: persist or not, enabled by default since we are checking for empty dataframes
      """
      from pyspark.sql import functions as F  # noqa: N812
      from pyspark.sql.window import Window
  
      if split_type in ["n_instances", "fraction"]:
          data_win = Window.orderBy(time_axis)
          if partition_cols is not None:
              data_win = data_win.partitionBy(partition_cols)
  
          if split_type == "n_instances":
              dt_rank = dataset.withColumn(
                  "tmpcol", F.row_number().over(data_win)
              )
          else:
              dt_rank = dataset.withColumn(
                  "tmpcol", F.percent_rank().over(data_win)
              )
  
          reference = dt_rank.filter(F.col("tmpcol") <= split).drop(
              "tmpcol"
          )
          df = dt_rank.filter(F.col("tmpcol") > split).drop("tmpcol")
      else:
          reference = dataset.filter(split)
          df = dataset.filter(f"NOT ({split})")
  
      if persist:
          reference.persist()
          df.persist()
  
      if len(reference.head(1)) == 0:
          raise ValueError(
              "Reference is empty. Please adjust the `split` argument"
          )
      if len(df.head(1)) == 0:
          raise ValueError(
              "Returned dataframe is empty. Please adjust the `split` argument"
          )
  
      return reference, df
  
imports: from __future__ import annotations
  from math import floor
  import pandas
  from typing_extensions import Literal
label_desc: Use image node_id 2 for calling a global function with example usage: _split_dataset_spark(dataset, split_type, split, time_axis, partition_cols, persist) and returns: reference,  df