我有一个分配问题,我想向 SO 社区询问为我的 spark 数据框(使用 spark 3.1+)实现此功能的最佳方法。我将首先描述问题,然后再进行实施。
问题是:我有最多 N 个任务和最多 N 个个人(在这个问题的情况下,N=10)。每个人都有执行每项任务的成本,其中最低成本为 0 美元,最高成本为 10 美元。这是一个匈牙利算法问题,有一些注意事项。
- 在某些情况下,任务少于 10 个和/或个人少于 10 个,可以不为某人分配任务(或不为任务分配个人)。
- [更复杂的边缘情况/我遇到麻烦的情况] - 列表中可能有一项具有标志的任务
multiTask=True
(不能超过 1multiTask
,并且可能没有)。如果一个工人的成本低于多x
任务,他会被自动分配给多任务,并且在优化期间考虑多任务。- 我将分享几个例子。在此示例中,要分配给多任务的 x 值为 1。
- 如果 10 名工人中有 1 名在多任务上的成本为 0.25,则将他分配给多任务,然后将其他 9 名工人分配给其他 9 个任务
- 如果 10 个工作人员中有 2 个工作人员在 multiTask 上的成本 < 1,则他们都被分配到 multiTask,然后其他 8 个工作人员将被分配到其余 9 个任务中的 8 个。1 项任务不会分配给任何人。
- 如果所有 10 个工作人员在 multiTask 上的 cost < 1,则所有这些工作人员都分配给 multiTask。这是非常罕见但可能的。
- 如果在 multiTask 上没有任何 worker 的 cost < 1,则在优化过程中 multiTask 将只分配给一个人以最小化成本。
- 我将分享几个例子。在此示例中,要分配给多任务的 x 值为 1。
这是 spark 数据框的样子。注意:为了简单起见,我展示了一个 N=3(3 个任务,3 个个人)的示例。
from pyspark.sql import Row
rdd = spark.sparkContext.parallelize([
Row(date='2019-08-01', locationId='z2-NY', workerId=129, taskId=220, cost=1.50, isMultiTask=False),
Row(date='2019-08-01', locationId='z2-NY', workerId=129, taskId=110, cost=2.90, isMultiTask=True),
Row(date='2019-08-01', locationId='z2-NY', workerId=129, taskId=190, cost=0.80, isMultiTask=False),
Row(date='2019-08-01', locationId='z2-NY', workerId=990, taskId=220, cost=1.80, isMultiTask=False),
Row(date='2019-08-01', locationId='z2-NY', workerId=990, taskId=110, cost=0.90, isMultiTask=True),
Row(date='2019-08-01', locationId='z2-NY', workerId=990, taskId=190, cost=9.99, isMultiTask=False),
Row(date='2019-08-01', locationId='z2-NY', workerId=433, taskId=220, cost=1.20, isMultiTask=False),
Row(date='2019-08-01', locationId='z2-NY', workerId=433, taskId=110, cost=0.25, isMultiTask=True),
Row(date='2019-08-01', locationId='z2-NY', workerId=433, taskId=190, cost=4.99, isMultiTask=False)
])
df = spark.createDataFrame(rdd)
您会看到有一个日期/位置,因为我需要为每个日期/位置分组解决这个分配问题。我打算通过使用dense_rank()
然后使用 pandas UDF 为每个工作人员和任务分配一个“索引”来解决这个问题,然后根据索引填充 N x N numpy 数组,然后调用该linear_sum_assignment
函数。但是,我不相信这个计划会奏效,因为我在 multiTask 中布置了第二个边缘案例。
worker_order_window = Window.partitionBy("date", "locationId").orderBy("workerId")
task_order_window = Window.partitionBy("date", "locationId").orderBy("taskId")
# get the dense_rank because will use this to assign a worker ID an index for the np array for linear_sum_assignment
# dense_rank - 1 as arrays are 0 indexed
df = df.withColumn("worker_idx", dense_rank().over(worker_order_window) - 1)
df = df.withColumn("task_idx", dense_rank().over(task_order_window) - 1)
def linear_assignment_udf(pandas_df: pd.DataFrame) -> pd.DataFrame:
df_dict = pandas_df.to_dict('records')
# in case there are less than N rows/columns
N = max(pandas_df.shape[0], pandas_df.shape[1])
arr = np.zeros((N,N))
for row in df_dict:
# worker_idx will be the row number, task idx will be the col number
worker_idx = row.get('worker_idx')
task_idx = row.get('task_idx')
arr[worker_idx][task_idx] = row.get('cost')
rids, cids = linear_sum_assignment(n)
return_list = []
# now want to return a dataframe that says which task_idx a worker has
for r, c in zip(rids, cids):
for d in df_dict:
if d.get('worker_idx') == r:
d['task_assignment'] = c
return_list.append(d)
return pd.DataFrame(return_list)
schema = StructType.fromJson(df.schema.jsonValue()).add('task_assignment', 'integer')
df = df.groupBy("date", "locationId").applyInPandas(linear_assignment_udf, schema)
df = df.withColumn("isAssigned", when(col("task_assignment") == col("task_idx"), True).otherwise(False))
如您所见,这种情况根本没有涵盖多任务。我想以最有效的方式解决这个问题,所以我不依赖于 pandas udf 或 scipy。