def python_random_split(data, ratio=0.75, seed=42):
"""Pandas random splitter
The splitter randomly splits the input data.
Args:
data (pd.DataFrame): Pandas DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halves and the ratio argument indicates the ratio
of training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
Returns:
list: Splits of the input data as pd.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
if multi_split:
splits = split_pandas_data_with_ratios(data, ratio, shuffle=True, seed=seed)
splits_new = [x.drop('split_index', axis=1) for x in splits]
return splits_new
else:
return sk_split(data, test_size=None, train_size=ratio, random_state=seed)
def python_random_split(data, ratio=0.75, seed=42):
"""Pandas random splitter
The splitter randomly splits the input data.
Args:
data (pd.DataFrame): Pandas DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halves and the ratio argument indicates the ratio
of training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
Returns:
list: Splits of the input data as pd.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
if multi_split:
splits = split_pandas_data_with_ratios(data,
ratio,
shuffle=True,
seed=seed)
splits_new = [x.drop('split_index', axis=1) for x in splits]
return splits_new
else:
return sk_split(data,
test_size=None,
train_size=ratio,
random_state=seed)
def spark_timestamp_split(
data,
ratio=0.75,
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Spark timestamp based splitter
The splitter splits the data into sets by timestamps without stratification on either
user or item.
The ratios are applied on the timestamp column which is divided accordingly into
several partitions.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two sets and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
Earlier indexed splits will have earlier times
(e.g the latest time in split[0] <= the earliest time in split[1])
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps. Float number represented in
seconds since Epoch.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
ratio = ratio if multi_split else [ratio, 1 - ratio]
ratio_index = np.cumsum(ratio)
window_spec = Window.orderBy(col(col_timestamp))
rating = data.withColumn("rank", row_number().over(window_spec))
data_count = rating.count()
rating_rank = rating.withColumn(
"rank",
row_number().over(window_spec) / data_count)
splits = []
for i, _ in enumerate(ratio_index):
if i == 0:
rating_split = rating_rank.filter(
col("rank") <= ratio_index[i]).drop("rank")
else:
rating_split = rating_rank.filter(
(col("rank") <= ratio_index[i])
& (col("rank") > ratio_index[i - 1])).drop("rank")
splits.append(rating_split)
return splits
def spark_timestamp_split(
data,
ratio=0.75,
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Spark timestamp based splitter
The splitter splits the data into sets by timestamps without stratification on either
user or item.
The ratios are applied on the timestamp column which is divided accordingly into
several partitions.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two sets and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
Earlier indexed splits will have earlier times
(e.g the latest time in split[0] <= the earliest time in split[1])
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps. Float number represented in
seconds since Epoch.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
ratio = ratio if multi_split else [ratio, 1 - ratio]
ratio_index = np.cumsum(ratio)
window_spec = Window.orderBy(col(col_timestamp))
rating = data.withColumn("rank", row_number().over(window_spec))
data_count = rating.count()
rating_rank = rating.withColumn("rank", row_number().over(window_spec) / data_count)
splits = []
for i, _ in enumerate(ratio_index):
if i == 0:
rating_split = rating_rank.filter(col("rank") <= ratio_index[i]).drop(
"rank"
)
else:
rating_split = rating_rank.filter(
(col("rank") <= ratio_index[i]) & (col("rank") > ratio_index[i - 1])
).drop("rank")
splits.append(rating_split)
return splits
def spark_random_split(data, ratio=0.75, seed=123):
"""Spark random splitter
Randomly split the data into several splits.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halfs and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list
is provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
if multi_split:
return data.randomSplit(ratio, seed=seed)
else:
return data.randomSplit([ratio, 1 - ratio], seed=seed)
def spark_random_split(data, ratio=0.75, seed=42):
"""Spark random splitter
Randomly split the data into several splits.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halves and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list
is provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
multi_split, ratio = process_split_ratio(ratio)
if multi_split:
return data.randomSplit(ratio, seed=seed)
else:
return data.randomSplit([ratio, 1 - ratio], seed=seed)
def _do_stratification(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
is_random=True,
seed=42,
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
# A few preliminary checks.
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError(
"min_rating should be integer and larger than or equal to 1.")
if col_user not in data.columns:
raise ValueError("Schema of data not valid. Missing User Col")
if col_item not in data.columns:
raise ValueError("Schema of data not valid. Missing Item Col")
if not is_random:
if col_timestamp not in data.columns:
raise ValueError("Schema of data not valid. Missing Timestamp Col")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
ratio = ratio if multi_split else [ratio, 1 - ratio]
if min_rating > 1:
data = min_rating_filter_pandas(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
# Split by each group and aggregate splits together.
splits = []
# If it is for chronological splitting, the split will be performed in a random way.
df_grouped = (data.sort_values(col_timestamp).groupby(split_by_column)
if is_random is False else data.groupby(split_by_column))
for name, group in df_grouped:
group_splits = split_pandas_data_with_ratios(
df_grouped.get_group(name), ratio, shuffle=is_random, seed=seed)
# Concatenate the list of split dataframes.
concat_group_splits = pd.concat(group_splits)
splits.append(concat_group_splits)
# Concatenate splits for all the groups together.
splits_all = pd.concat(splits)
# Take split by split_index
splits_list = [
splits_all[splits_all["split_index"] == x].drop("split_index", axis=1)
for x in range(len(ratio))
]
return splits_list
def python_chrono_split(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Pandas chronological splitter
This function splits data in a chronological manner. That is, for each user / item, the
split function takes proportions of ratings which is specified by the split ratio(s).
The split is stratified.
Args:
data (pd.DataFrame): Pandas DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halfs and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to
filter with min_rating.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps.
Returns:
list: Splits of the input data as pd.DataFrame.
"""
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError(
"min_rating should be integer and larger than or equal to 1.")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
# Sort data by timestamp.
data = data.sort_values(by=[split_by_column, col_timestamp],
axis=0,
ascending=False)
ratio = ratio if multi_split else [ratio, 1 - ratio]
if min_rating > 1:
data = min_rating_filter_pandas(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
num_of_splits = len(ratio)
splits = [pd.DataFrame({})] * num_of_splits
df_grouped = data.sort_values(col_timestamp).groupby(split_by_column)
for name, group in df_grouped:
group_splits = split_pandas_data_with_ratios(
df_grouped.get_group(name), ratio, resample=False)
for x in range(num_of_splits):
splits[x] = pd.concat([splits[x], group_splits[x]])
return splits
def spark_chrono_split(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Spark chronological splitter
This function splits data in a chronological manner. That is, for each user / item, the
split function takes proportions of ratings which is specified by the split ratio(s).
The split is stratified.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two sets and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to filter
with min_rating.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError("min_rating should be integer and larger than or equal to 1.")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
if min_rating > 1:
data = min_rating_filter_spark(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
ratio = ratio if multi_split else [ratio, 1 - ratio]
ratio_index = np.cumsum(ratio)
window_spec = Window.partitionBy(split_by_column).orderBy(col(col_timestamp))
rating_grouped = (
data.groupBy(split_by_column)
.agg({col_timestamp: "count"})
.withColumnRenamed("count(" + col_timestamp + ")", "count")
)
rating_all = data.join(broadcast(rating_grouped), on=split_by_column)
rating_rank = rating_all.withColumn(
"rank", row_number().over(window_spec) / col("count")
)
splits = []
for i, _ in enumerate(ratio_index):
if i == 0:
rating_split = rating_rank.filter(col("rank") <= ratio_index[i])
else:
rating_split = rating_rank.filter(
(col("rank") <= ratio_index[i]) & (col("rank") > ratio_index[i - 1])
)
splits.append(rating_split)
return splits
def spark_chrono_split(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Spark chronological splitter
This function splits data in a chronological manner. That is, for each user / item, the
split function takes proportions of ratings which is specified by the split ratio(s).
The split is stratified.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two sets and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to filter
with min_rating.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError("min_rating should be integer and larger than or equal to 1.")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
if min_rating > 1:
data = min_rating_filter_spark(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
ratio = ratio if multi_split else [ratio, 1 - ratio]
ratio_index = np.cumsum(ratio)
window_spec = Window.partitionBy(split_by_column).orderBy(col(col_timestamp))
rating_grouped = (
data.groupBy(split_by_column)
.agg({col_timestamp: "count"})
.withColumnRenamed("count(" + col_timestamp + ")", "count")
)
rating_all = data.join(broadcast(rating_grouped), on=split_by_column)
rating_rank = rating_all.withColumn(
"rank", row_number().over(window_spec) / col("count")
)
splits = []
for i, _ in enumerate(ratio_index):
if i == 0:
rating_split = rating_rank.filter(col("rank") <= ratio_index[i])
else:
rating_split = rating_rank.filter(
(col("rank") <= ratio_index[i]) & (col("rank") > ratio_index[i - 1])
)
splits.append(rating_split)
return splits
def spark_stratified_split(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_rating=DEFAULT_RATING_COL,
seed=42,
):
"""Spark stratified splitter.
For each user / item, the split function takes proportions of ratings which is
specified by the split ratio(s). The split is stratified.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halves and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
Earlier indexed splits will have earlier times
(e.g the latest time per user or item in split[0] <= the earliest time per user or item in split[1])
seed (int): Seed.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to filter
with min_rating.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_rating (str): column name of ratings.
Returns:
list: Splits of the input data as spark.DataFrame.
"""
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError("min_rating should be integer and larger than or equal to 1.")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
if min_rating > 1:
data = min_rating_filter_spark(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
ratio = ratio if multi_split else [ratio, 1 - ratio]
ratio_index = np.cumsum(ratio)
window_count = Window.partitionBy(split_by_column)
window_spec = Window.partitionBy(split_by_column).orderBy(rand(seed=seed))
rating_all = data.withColumn(
"count", size(collect_list(col_rating).over(window_count))
)
rating_rank = rating_all.withColumn(
"rank", row_number().over(window_spec) / col("count")
)
splits = []
for i, _ in enumerate(ratio_index):
if i == 0:
rating_split = rating_rank.filter(col("rank") <= ratio_index[i])
else:
rating_split = rating_rank.filter(
(col("rank") <= ratio_index[i]) & (col("rank") > ratio_index[i - 1])
)
splits.append(rating_split)
return splits
def _do_stratification(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
is_random=True,
seed=42,
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
# A few preliminary checks.
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError("min_rating should be integer and larger than or equal to 1.")
if col_user not in data.columns:
raise ValueError("Schema of data not valid. Missing User Col")
if col_item not in data.columns:
raise ValueError("Schema of data not valid. Missing Item Col")
if not is_random:
if col_timestamp not in data.columns:
raise ValueError("Schema of data not valid. Missing Timestamp Col")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
ratio = ratio if multi_split else [ratio, 1 - ratio]
if min_rating > 1:
data = min_rating_filter_pandas(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
# Split by each group and aggregate splits together.
splits = []
# If it is for chronological splitting, the split will be performed in a random way.
df_grouped = (
data.sort_values(col_timestamp).groupby(split_by_column)
if is_random is False
else data.groupby(split_by_column)
)
for name, group in df_grouped:
group_splits = split_pandas_data_with_ratios(
df_grouped.get_group(name), ratio, shuffle=is_random, seed=seed
)
# Concatenate the list of split dataframes.
concat_group_splits = pd.concat(group_splits)
splits.append(concat_group_splits)
# Concatenate splits for all the groups together.
splits_all = pd.concat(splits)
# Take split by split_index
splits_list = [
splits_all[splits_all["split_index"] == x].drop("split_index", axis=1)
for x in range(len(ratio))
]
return splits_list
def python_chrono_split(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Pandas chronological splitter
This function splits data in a chronological manner. That is, for each user / item, the
split function takes proportions of ratings which is specified by the split ratio(s).
The split is stratified.
Args:
data (pd.DataFrame): Pandas DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two halfs and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
seed (int): Seed.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to
filter with min_rating.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps.
Returns:
list: Splits of the input data as pd.DataFrame.
"""
# A few preliminary checks.
if not (filter_by == "user" or filter_by == "item"):
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError(
"min_rating should be integer and larger than or equal to 1.")
if col_user not in data.columns:
raise ValueError("Schema of data not valid. Missing User Col")
if col_item not in data.columns:
raise ValueError("Schema of data not valid. Missing Item Col")
if col_timestamp not in data.columns:
raise ValueError("Schema of data not valid. Missing Timestamp Col")
multi_split, ratio = process_split_ratio(ratio)
split_by_column = col_user if filter_by == "user" else col_item
ratio = ratio if multi_split else [ratio, 1 - ratio]
if min_rating > 1:
data = min_rating_filter_pandas(
data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
# Split by each group and aggregate splits together.
splits = []
df_grouped = data.sort_values(col_timestamp).groupby(split_by_column)
for name, group in df_grouped:
group_splits = split_pandas_data_with_ratios(
df_grouped.get_group(name), ratio, shuffle=False)
# Concatenate the list of split dataframes.
concat_group_splits = pd.concat(group_splits)
splits.append(concat_group_splits)
# Concatenate splits for all the groups together.
splits_all = pd.concat(splits)
# Take split by split_index
splits_list = [
splits_all[splits_all["split_index"] == x].drop("split_index", axis=1)
for x in range(len(ratio))
]
return splits_list
def _do_stratification_spark(
data,
ratio=0.75,
min_rating=1,
filter_by="user",
is_partitioned=True,
is_random=True,
seed=42,
col_user=DEFAULT_USER_COL,
col_item=DEFAULT_ITEM_COL,
col_timestamp=DEFAULT_TIMESTAMP_COL,
):
"""Helper function to perform stratified splits.
This function splits data in a stratified manner. That is, the same values for the
filter_by column are retained in each split, but the corresponding set of entries
are divided according to the ratio provided.
Args:
data (spark.DataFrame): Spark DataFrame to be split.
ratio (float or list): Ratio for splitting data. If it is a single float number
it splits data into two sets and the ratio argument indicates the ratio of
training data set; if it is a list of float numbers, the splitter splits
data into several portions corresponding to the split ratios. If a list is
provided and the ratios are not summed to 1, they will be normalized.
min_rating (int): minimum number of ratings for user or item.
filter_by (str): either "user" or "item", depending on which of the two is to filter
with min_rating.
is_partitioned (bool): flag to partition data by filter_by column
is_random (bool): flag to make split randomly or use timestamp column
seed (int): Seed.
col_user (str): column name of user IDs.
col_item (str): column name of item IDs.
col_timestamp (str): column name of timestamps.
Args:
Returns:
"""
# A few preliminary checks.
if filter_by not in ["user", "item"]:
raise ValueError("filter_by should be either 'user' or 'item'.")
if min_rating < 1:
raise ValueError(
"min_rating should be integer and larger than or equal to 1.")
if col_user not in data.columns:
raise ValueError("Schema of data not valid. Missing User Col")
if col_item not in data.columns:
raise ValueError("Schema of data not valid. Missing Item Col")
if not is_random:
if col_timestamp not in data.columns:
raise ValueError("Schema of data not valid. Missing Timestamp Col")
if min_rating > 1:
data = min_rating_filter_spark(
data=data,
min_rating=min_rating,
filter_by=filter_by,
col_user=col_user,
col_item=col_item,
)
split_by = col_user if filter_by == "user" else col_item
partition_by = split_by if is_partitioned else []
order_by = F.rand(seed=seed) if is_random else F.col(col_timestamp)
window_count = Window.partitionBy(partition_by)
window_spec = Window.partitionBy(partition_by).orderBy(order_by)
data = (data.withColumn("_count",
F.count(split_by).over(window_count)).withColumn(
"_rank",
F.row_number().over(window_spec) /
F.col("_count")).drop("_count"))
multi_split, ratio = process_split_ratio(ratio)
ratio = ratio if multi_split else [ratio, 1 - ratio]
splits = []
prev_split = None
for split in np.cumsum(ratio):
condition = F.col("_rank") <= split
if prev_split is not None:
condition &= F.col("_rank") > prev_split
splits.append(data.filter(condition).drop("_rank"))
prev_split = split
return splits
