Spark技术问询：基于多行文件按起始字符拆分生成RDD/SparkDataFrame

Great question! Let's break down how to split your input data into groups where each group starts with a string beginning with "00", and includes all subsequent lines until the next "00" entry. We'll cover both RDD and DataFrame approaches for Scala and PySpark.

RDD Approach #

Basic Partition-Wise Grouping (Simple but Partition-Dependent) #

If your data is small enough to fit in a single partition, or each partition starts with a "00" entry, this straightforward method works:

Scala

// Read input file into RDD
val inputRDD = spark.sparkContext.textFile("path/to/your/file.txt")

// Assign group IDs and group entries
val groupedRDD = inputRDD.mapPartitions(iter => {
  var groupId = 0
  var isFirstEntry = true
  iter.map { line =>
    // Increment group ID when we hit a new "00" entry (skip first one)
    if (line.startsWith("00")) {
      if (!isFirstEntry) groupId += 1
      isFirstEntry = false
    }
    (groupId, line)
  }
}).groupByKey().map(_._2.toArray) // Convert grouped values to array

// View results
groupedRDD.collect().foreach(group => println(group.mkString("[", ", ", "]")))

PySpark

# Read input file into RDD
input_rdd = spark.sparkContext.textFile("path/to/your/file.txt")

def assign_group_ids(iterator):
    group_id = 0
    is_first_entry = True
    for line in iterator:
        if line.startswith("00"):
            if not is_first_entry:
                group_id += 1
            is_first_entry = False
        yield (group_id, line)

# Assign groups and convert to arrays
grouped_rdd = input_rdd.mapPartitions(assign_group_ids).groupByKey().map(lambda x: list(x[1]))

# View results
for group in grouped_rdd.collect():
    print(f"[{', '.join(group)}]")

Global Grouping (Works Across Partitions) #

For distributed datasets where "00" entries might span partitions, we need a global approach using indices to track group boundaries:

Scala

val inputRDD = spark.sparkContext.textFile("path/to/your/file.txt")

// Assign global indices to each line
val indexedRDD = inputRDD.zipWithIndex().map(_.swap) // (index, line)

// Collect all indices where lines start with "00"
val startIndices = indexedRDD.filter(_._2.startsWith("00")).keys.collect().sorted

// Map each line to its corresponding group
val groupedRDD = indexedRDD.map { case (idx, line) =>
  // Find the latest start index that's <= current line's index
  val groupId = startIndices.lastIndexWhere(_ <= idx)
  (groupId, line)
}.groupByKey().map(_._2.toArray)

groupedRDD.collect().foreach(group => println(group.mkString("[", ", ", "]")))

PySpark

input_rdd = spark.sparkContext.textFile("path/to/your/file.txt")

# Assign global indices
indexed_rdd = input_rdd.zipWithIndex().map(lambda x: (x[1], x[0]))

# Collect indices of "00" starting lines
start_indices = sorted(indexed_rdd.filter(lambda x: x[1].startswith("00")).keys().collect())

def map_to_group(iterator):
    for idx, line in iterator:
        # Find the right group ID using binary search logic
        group_id = next(i for i, s_idx in enumerate(start_indices) if s_idx > idx) - 1
        yield (group_id, line)

grouped_rdd = indexed_rdd.mapPartitions(map_to_group).groupByKey().map(lambda x: list(x[1]))

for group in grouped_rdd.collect():
    print(f"[{', '.join(group)}]")

DataFrame Approach #

If you prefer working with DataFrames (recommended for most Spark use cases), we can use window functions to calculate group boundaries:

Scala

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

// Read input into DataFrame
val df = spark.read.text("path/to/your/file.txt").toDF("line")

// Define window (assumes input order is preserved; adjust if needed)
val windowSpec = Window.orderBy(lit(1))

// Mark group starts, calculate group IDs, then aggregate
val groupedDF = df
  .withColumn("is_group_start", when(col("line").startsWith("00"), 1).otherwise(0))
  .withColumn("group_id", sum("is_group_start").over(windowSpec) - 1) // Start group IDs at 0
  .groupBy("group_id")
  .agg(collect_list("line").alias("group"))
  .select("group")

// Show results
groupedDF.show(false)

PySpark

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Read input into DataFrame
df = spark.read.text("path/to/your/file.txt").toDF("line")

# Define window (preserves input order)
window_spec = Window.orderBy(F.lit(1))

# Calculate groups and aggregate
grouped_df = df \
    .withColumn("is_group_start", F.when(F.col("line").startswith("00"), 1).otherwise(0)) \
    .withColumn("group_id", F.sum("is_group_start").over(window_spec) - 1) \
    .groupBy("group_id") \
    .agg(F.collect_list("line").alias("group")) \
    .select("group")

# Show results
grouped_df.show(truncate=False)

All these methods will produce the exact output you're looking for:

• First group: [00000000000, 01000000000, 02000000000]
• Second group: [00000000000, 01000000000, 02000000000, 03000000000]

内容的提问来源于stack exchange，提问作者Pedro Kássio