def test_min_rating_filter(spark_dataset):
dfs_user = min_rating_filter_spark(spark_dataset, min_rating=5, filter_by="user")
dfs_item = min_rating_filter_spark(spark_dataset, min_rating=5, filter_by="item")
user_rating_counts = [
x["count"] >= 5 for x in dfs_user.groupBy(DEFAULT_USER_COL).count().collect()
]
item_rating_counts = [
x["count"] >= 5 for x in dfs_item.groupBy(DEFAULT_ITEM_COL).count().collect()
]
assert all(user_rating_counts)
assert all(item_rating_counts)
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 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 _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
