PySpark中如何用字符串包裹saveAsTextFile的输出结果

Absolutely, you can pull this off directly with saveAsTextFile—no post-file-editing required! The trick is swapping out map for mapPartitions, since it lets you work with entire partitions (which correspond one-to-one with your output part-* files) and wrap each partition's full content with your <scrog> root tags.

Here's the adjusted code tailored to your needs:

import math, uuid

# Number of records per output file
perpage = 100
total_records = df.count()
num_partitions = math.ceil(total_records / perpage)

# Transform the RDD to wrap each partition with valid XML root tags
xml_rdd = (
    df.select('document')
    .rdd
    .repartition(num_partitions)
    .mapPartitions(lambda partition:
        # Collect and clean all XML docs in the current partition
        processed_docs = [
            row.document.replace('<?xml version="1.0" encoding="UTF-8"?>', '')
            for row in partition
        ]
        # Only add root tags if the partition has content
        if processed_docs:
            ["<scrog>"] + processed_docs + ["</scrog>"]
        else:
            []  # Skip empty partitions to avoid useless empty root tag files
    )
)

# Save the properly formatted XML partitions
output_path = f'file:///foo/bar/output/{str(uuid.uuid4())}'
xml_rdd.saveAsTextFile(output_path)

Breakdown of how this works: #

• mapPartitions operates on an iterator for each entire partition, so we can handle all rows in the partition in one go.
• We first gather and clean all XML documents in the partition (removing the redundant declaration as you originally did).
• If the partition has content, we prepend the opening <scrog> tag and append the closing </scrog> tag to the list of documents. This ensures every output file starts and ends with the root tag, making the XML fully valid.
• Empty partitions are skipped entirely, so you won't end up with files containing just <scrog></scrog> if a partition has no records.

For your 285-record example, this will generate 3 part-* files: the first two with 100 cleaned XML docs wrapped in <scrog> tags, and the third with 85 docs in the same valid structure.

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