时间戳

首页 > 解决方案 > 如何在 PySpark UDF 中使用边缘情况解决分配问题（如匈牙利语/linear_sum_assignment）

python - 如何在 PySpark UDF 中使用边缘情况解决分配问题（如匈牙利语/linear_sum_assignment）

问题描述

我有一个分配问题，我想向 SO 社区询问为我的 spark 数据框（使用 spark 3.1+）实现此功能的最佳方法。我将首先描述问题，然后再进行实施。

问题是：我有最多 N 个任务和最多 N 个个人（在这个问题的情况下，N=10）。每个人都有执行每项任务的成本，其中最低成本为 0 美元，最高成本为 10 美元。这是一个匈牙利算法问题，有一些注意事项。

  1. 在某些情况下，任务少于 10 个和/或个人少于 10 个，可以不为某人分配任务（或不为任务分配个人）。
  2. [更复杂的边缘情况/我遇到麻烦的情况] - 列表中可能有一项具有标志的任务multiTask=True（不能超过 1 multiTask，并且可能没有）。如果一个工人的成本低于多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 将只分配给一个人以最小化成本。

这是 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。

标签： pythonapache-sparkpysparkscipy-optimizehungarian-algorithm

解决方案

我对您正在使用的库一无所知，所以我无法帮助您编写代码，但我认为您应该分两步执行此操作：

  1. 如果需要将工作人员分配给多任务，则将其分配给它。如果有人被分配到此任务，请不要将其包含在您的成本矩阵中。
  2. 像往常一样使用匈牙利算法（或其他替代算法）将工人分配给任务。

基本的匈牙利算法仅适用于平方成本矩阵，看起来您已经通过用 0 填充成本矩阵来正确处理该问题，但是对矩形矩阵的算法进行了修改。您可能想看看您是否可以使用其中一种替代方案，因为它可能会更快。

推荐阅读