Spark如何标记同ID下绝对值相等、符号相反的行对？

需求说明 #

给定包含id和value两列的Spark Dataset，需要标记出满足以下条件的行对：

• 同一id
• value的绝对值相等、符号相反
• 每个正数最多与一个负数配对，反之亦然

解决方案思路 #

核心逻辑是按id和value的绝对值分组，对组内的正数、负数分别生成递增序号，序号相同的正负行即为有效配对，标记为true；未找到对应配对的行标记为false。

代码实现 #

Scala 版本

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

// 构造示例数据
val df = spark.createDataFrame(Seq(
  ("a", 5), ("b", 10), ("a", -5), ("b", 10), ("a", 5),
  ("b", 10), ("a", -5), ("b", 5), ("a", 5), ("b", 1)
)).toDF("id", "value")

// 步骤1：计算绝对值、标记正负符号
val withAbsAndSign = df.withColumn("abs_value", abs(col("value")))
  .withColumn("sign", when(col("value") > 0, 1).otherwise(-1))

// 步骤2：对每个(id, abs_value, sign)组内的行生成递增序号
val windowSpec = Window.partitionBy("id", "abs_value", "sign").orderBy("value")
val withRank = withAbsAndSign.withColumn("rank", row_number().over(windowSpec))

// 步骤3：统计每个(id, abs_value, rank)下的正负符号数量，判断是否存在配对
val countSigns = withRank.groupBy("id", "abs_value", "rank")
  .agg(countDistinct("sign").alias("sign_count"))

// 步骤4：关联回原数据，标记flag
val result = withRank.join(countSigns, Seq("id", "abs_value", "rank"), "left")
  .withColumn("flag", col("sign_count") === 2)
  .select("id", "value", "flag")

result.show()

Python 版本

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import abs, when, row_number, countDistinct

# 初始化SparkSession
spark = SparkSession.builder.appName("PairMarking").getOrCreate()

# 构造示例数据
data = [
    ("a", 5), ("b", 10), ("a", -5), ("b", 10), ("a", 5),
    ("b", 10), ("a", -5), ("b", 5), ("a", 5), ("b", 1)
]
df = spark.createDataFrame(data, ["id", "value"])

# 步骤1：计算绝对值、标记正负符号
with_abs_and_sign = df.withColumn("abs_value", abs(df["value"])) \
    .withColumn("sign", when(df["value"] > 0, 1).otherwise(-1))

# 步骤2：对每个(id, abs_value, sign)组内的行生成递增序号
window_spec = Window.partitionBy("id", "abs_value", "sign").orderBy("value")
with_rank = with_abs_and_sign.withColumn("rank", row_number().over(window_spec))

# 步骤3：统计每个(id, abs_value, rank)下的正负符号数量
count_signs = with_rank.groupBy("id", "abs_value", "rank") \
    .agg(countDistinct("sign").alias("sign_count"))

# 步骤4：关联并标记flag
result = with_rank.join(count_signs, ["id", "abs_value", "rank"], "left") \
    .withColumn("flag", (count_signs["sign_count"] == 2)) \
    .select("id", "value", "flag")

result.show()

结果说明 #

运行代码后输出结果将与需求示例一致：

• 序号匹配的正负行被标记为true
• 超出配对数量的冗余行（如id=b的第三个10、id=a的第三个5）标记为false
• 无对应正负值的行（如id=b的5、1）标记为false

内容的提问来源于stack exchange，提问作者sanketd617