正如注释中所讨论的，基本思想是对数据进行适当的分区，以便具有相同LNAME+Address的记录保持在同一个分区中，运行Python代码在每个分区上生成单独的idx，然后将它们合并到最终的id。在

注意：我在示例记录中添加了一些新行，请参见下面显示的df_new.show()的结果。在from pyspark.sql import Window, Row

from pyspark.sql.functions import coalesce, sum as fsum, col, max as fmax, lit, broadcast

# ...skip code to initialize the dataframe

# tweak the number of repartitioning N based on actual data size

N = 5

# Python function to iterate through the sorted list of elements in the same

# partition and assign an in-partition idx based on Address and LNAME.

def func(partition_id, it):

idx, lname, address = (1, None, None)

for row in sorted(it, key=lambda x: (x.LNAME, x.Address)):

if lname and (row.LNAME != lname or row.Address != address): idx += 1

yield Row(partition_id=partition_id, idx=idx, **row.asDict())

lname = row.LNAME

address = row.Address

# Repartition based on 'LNAME' and 'Address' and then run mapPartitionsWithIndex()

# function to create in-partition idx. Adjust N so that records in each partition

# should be small enough to be loaded into the executor memory:

df1 = df.repartition(N, 'LNAME', 'Address') \

.rdd.mapPartitionsWithIndex(func) \

.toDF()

获取唯一行数cnt(基于Address+LNAME)，即max_idx，然后获取该rcnt的运行和。在

^{pr2}$

将df1与df2连接并创建最终的id idx + rcntdf_new = df1.join(broadcast(df2), on=['partition_id']).withColumn('id', col('idx')+col('rcnt'))

df_new.show()

#+ + -+ -+ + -+ + + -+ -+ + -+ + -+

#|partition_id|Address| D| DOB|FNAME|GENDER| LNAME|MNAME|idx|snapshot|cnt|rcnt| id|

#+ + -+ -+ + -+ + + -+ -+ + -+ + -+

#| 0| B| 0|1990|David| M| Lee| H M| 1|201211.0| 3| 0| 1|

#| 0| J| 3|1991|David| M| Lee| HM| 2|201211.0| 3| 0| 2|

#| 0| D| 6|2000| Marc| M|Robert| MS| 3|201211.0| 3| 0| 3|

#| 1| C| 3|2000| Marc| M|Robert| H| 1|201211.0| 1| 3| 4|

#| 1| C| 6|1988| Marc| M|Robert| M| 1|201211.0| 1| 3| 4|

#| 2| J| 6|1991| 66M| F| Rek| null| 1|201211.0| 1| 4| 5|

#| 2| J| 6|1992| 66M| F| Rek| null| 1|201211.0| 1| 4| 5|

#| 4| J| 2|1995| 66M| F| Rock| J| 1|201211.0| 1| 5| 6|

#| 4| J| 6|1990| 66M| F| Rock| null| 1|201211.0| 1| 5| 6|

#| 4| J| 6|1990| 66M| F| Rock| null| 1|201211.0| 1| 5| 6|

#+ + -+ -+ + -+ + + -+ -+ + -+ + -+

df_new = df_new.drop('partition_id', 'idx', 'rcnt', 'cnt')

注意事项：实际上，在将列LNAME和Address用作唯一性检查之前，需要清除/规范化列LNAME。例如，使用一个单独的列uniq_key，它组合了LNAME和Address作为数据帧的唯一键。下面是一些基本数据清理过程的示例：from pyspark.sql.functions import coalesce, lit, concat_ws, upper, regexp_replace, trim

#(1) convert NULL to '': coalesce(col, '')

#(2) concatenate LNAME and Address using NULL char '\x00' or '\0'

#(3) convert to uppercase: upper(text)

#(4) remove all non-[word/whitespace/NULL_char]: regexp_replace(text, r'[^\x00\w\s]', '')

#(5) convert consecutive whitespaces to a SPACE: regexp_replace(text, r'\s+', ' ')

#(6) trim leading/trailing spaces: trim(text)

df = (df.withColumn('uniq_key',

trim(

regexp_replace(

regexp_replace(

upper(

concat_ws('\0', coalesce('LNAME', lit('')), coalesce('Address', lit('')))

),

r'[^\x00\s\w]+',

''

),

r'\s+',

' '

)

)

))

然后在代码中，将'LNAME'和{}替换为uniq_key，以找到idx

正如cronoik在注释中提到的，您也可以尝试使用一个窗口等级函数来计算分区内的idx。例如：from pyspark.sql.functions import spark_partition_id, dense_rank

# use dense_rank to calculate the in-partition idx

w2 = Window.partitionBy('partition_id').orderBy('LNAME', 'Address')

df1 = df.repartition(N, 'LNAME', 'Address') \

.withColumn('partition_id', spark_partition_id()) \

.withColumn('idx', dense_rank().over(w2))

当你有了df1之后，用同样的方法来计算df2和df_new。这应该比使用mapPartitionsWithIndex()更快，后者基本上是一种基于RDD的方法。

对于实际数据，请调整N以适合您的实际数据大小。这个N只影响初始分区，在dataframe连接之后，分区将被重置为默认值(200)。您可以使用spark.sql.shuffle.partitions来调整此值，例如在初始化spark会话时：spark = SparkSession.builder \

....

.config("spark.sql.shuffle.partitions", 500) \

.getOrCreate()