Scala环境下Spark DataFrame数据关联与两两通道概率列添加

Got it, let's walk through how to add those pairwise channel co-occurrence probability columns to your Spark DataFrame in Scala. First, let's align on the setup—assuming your DataFrame has columns A, B, C (0/1 for absent/present) and total (sum of rows for each channel combination), we'll calculate the pairwise co-occurrence probabilities (we'll cover both joint probability and conditional probability since you mentioned "pairwise occurrence probability").

Step 1: Example Data Setup (for context) #

First, let's create a sample DataFrame matching your description to test our code:

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

val spark = SparkSession.builder()
  .appName("ChannelCooccurrence")
  .master("local[*]") // Remove this for cluster mode
  .getOrCreate()

// Simulate your 8-channel-combination DataFrame
val df = spark.createDataFrame(Seq(
  (0, 0, 0, 5),   // No channels present
  (0, 0, 1, 10),  // Only C present
  (0, 1, 0, 15),  // Only B present
  (0, 1, 1, 20),  // B + C present
  (1, 0, 0, 25),  // Only A present
  (1, 0, 1, 30),  // A + C present
  (1, 1, 0, 35),  // A + B present
  (1, 1, 1, 40)   // All channels present
)).toDF("A", "B", "C", "total")

Step 2: Define Probability Logic #

We'll focus on two common pairwise probability definitions:

• Joint Probability: Probability that both channels are present (e.g., P(A∩B) = sum of total where A=1 and B=1 / total sum of all total values)
• Conditional Probability: Probability one channel is present given the other is (e.g., P(A|B) = sum of total where A=1 and B=1 / sum of total where B=1)

Step 3: Implementation Using Window Functions (Recommended for Big Data) #

This approach avoids collecting data to the driver node (safer for large datasets) using Spark's window functions:

// Create a global window (applies aggregations across all rows)
val globalWindow = Window.partitionBy()

// Calculate pairwise co-occurrence totals and probabilities
val resultDF = df
  // Calculate global total sum (denominator for joint probability)
  .withColumn("global_total", sum("total").over(globalWindow))
  // Calculate A+B co-occurrence total and joint probability
  .withColumn("ab_cooccur_total", sum(when(col("A") === 1 && col("B") === 1, col("total")).otherwise(0)).over(globalWindow))
  .withColumn("prob_ab_joint", col("ab_cooccur_total") / col("global_total"))
  // Calculate A+C co-occurrence total and joint probability
  .withColumn("ac_cooccur_total", sum(when(col("A") === 1 && col("C") === 1, col("total")).otherwise(0)).over(globalWindow))
  .withColumn("prob_ac_joint", col("ac_cooccur_total") / col("global_total"))
  // Calculate B+C co-occurrence total and joint probability
  .withColumn("bc_cooccur_total", sum(when(col("B") === 1 && col("C") === 1, col("total")).otherwise(0)).over(globalWindow))
  .withColumn("prob_bc_joint", col("bc_cooccur_total") / col("global_total"))
  // Optional: Drop intermediate calculation columns
  .drop("global_total", "ab_cooccur_total", "ac_cooccur_total", "bc_cooccur_total")

If you need conditional probabilities instead, add these steps:

// Calculate single-channel totals for conditional probability denominators
val conditionalResultDF = resultDF
  .withColumn("total_a", sum(when(col("A") === 1, col("total")).otherwise(0)).over(globalWindow))
  .withColumn("total_b", sum(when(col("B") === 1, col("total")).otherwise(0)).over(globalWindow))
  .withColumn("total_c", sum(when(col("C") === 1, col("total")).otherwise(0)).over(globalWindow))
  // Add conditional probability columns
  .withColumn("prob_a_given_b", col("ab_cooccur_total") / col("total_b"))
  .withColumn("prob_a_given_c", col("ac_cooccur_total") / col("total_c"))
  .withColumn("prob_b_given_c", col("bc_cooccur_total") / col("total_c"))
  .drop("total_a", "total_b", "total_c", "ab_cooccur_total", "ac_cooccur_total", "bc_cooccur_total")

Step 4: Alternative Implementation (Small Datasets) #

If your dataset is small, you can collect aggregation results to the driver and broadcast them as literal values:

// Calculate all required aggregates
val aggMetrics = df.agg(
  sum("total").alias("global_total"),
  sum(when(col("A") === 1 && col("B") === 1, col("total")).otherwise(0)).alias("ab_total"),
  sum(when(col("A") === 1 && col("C") === 1, col("total")).otherwise(0)).alias("ac_total"),
  sum(when(col("B") === 1 && col("C") === 1, col("total")).otherwise(0)).alias("bc_total")
).collect()(0)

// Extract values
val globalTotal = aggMetrics.getAs[Long]("global_total").toDouble
val abTotal = aggMetrics.getAs[Long]("ab_total").toDouble
val acTotal = aggMetrics.getAs[Long]("ac_total").toDouble
val bcTotal = aggMetrics.getAs[Long]("bc_total").toDouble

// Add probability columns
val smallDatasetResult = df
  .withColumn("prob_ab_joint", lit(abTotal / globalTotal))
  .withColumn("prob_ac_joint", lit(acTotal / globalTotal))
  .withColumn("prob_bc_joint", lit(bcTotal / globalTotal))

Step 5: View the Result #

Run resultDF.show() to see your original DataFrame with the new probability columns appended.

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