Palantir Foundry中Spark如何用partitionByRange写增量数据集

问题背景 #

• 在Palantir Foundry中启用Spark分区裁剪能力，需调用transforms.api.IncrementalTransformOutput.write_dataframe()方法，传入partitionBy=['col1', 'col2', 'col_N']参数即可实现。
• 实际测试中，在分区列数据分布完全均匀的增量数据集上使用上述配置时，输出文件大小差异极大，覆盖6MB到128MB区间；当前分区列仅存在24种不同取值组合，初步推测文件大小不均与默认分区逻辑的哈希冲突有关。
• 由于使用的两个分区列均为可排序值类型，希望改用partitionByRange()范围分区逻辑生成大小更均匀的文件。
• 已知平台限制：Foundry不允许对增量数据集使用分桶（bucketing）能力，无法通过分桶方案优化数据分布。
• 核心诉求：确认Palantir Foundry是否支持范围分区写入操作；由于增量数据集规模会随时间持续增长，需要将文件大小稳定维持在128MB左右的合理区间，获得可预测的文件大小输出，支撑后续参数调优。

解决方案 #

Palantir Foundry原生支持对增量数据集使用范围分区写入，无需依赖分桶能力，具体实现逻辑如下：
首先需要纠正一个认知偏差：观察到的文件大小不均并非哈希冲突导致。write_dataframe()方法传入的partitionBy参数对应的是Hive风格目录级分区，作用是按分区列取值切分存储子目录，是分区裁剪能力生效的基础，和Spark RDD层的哈希分区器、范围分区器不属于同一逻辑层级。单目录下文件大小在6MB到128MB区间波动，核心原因是写入前数据在Spark task间分布不均，默认写入逻辑没有做单分区内的文件大小对齐。

注意：目录级分区的partitionBy参数不可省略或替换，否则将导致分区裁剪能力完全失效，下游查询性能会出现大幅下降。

要实现范围分区、稳定输出128MB左右的目标大小文件，只需要在写入前对DataFrame做范围重分区处理，保留原有写入时的partitionBy配置即可，参考实现代码如下：

from transforms.api import transform, incremental, Input, Output

@incremental(semantic_version=1)
@transform(
    output=Output("/path/to/target/output_dataset"),
    source_df=Input("/path/to/source/input_dataset")
)
def range_partition_write(source_df, output):
    df = source_df.dataframe()
    # 定义分区列，需和分区裁剪依赖的列保持一致
    partition_columns = ["col1", "col2"]
    
    # 按128MB单文件目标大小计算所需分区数，可根据实际使用的压缩格式调整系数
    target_single_file_size = 128 * 1024 * 1024
    # 估算当前待写入数据总大小
    total_data_size = df.rdd.map(lambda row: len(str(row))).sum()
    # 分区数不低于分区列的取值组合数，避免出现部分分区空写
    required_partition_num = max(24, int(total_data_size / target_single_file_size) + 1)
    
    # 写入前按分区列做范围重分区，替代默认的哈希分发逻辑
    df = df.repartitionByRange(required_partition_num, *partition_columns)
    
    # 写入时保留partitionBy参数，保证目录分区结构正确，分区裁剪正常生效
    output.write_dataframe(
        df,
        partitionBy=partition_columns
    )

调优说明 #

• 上述实现完全适配Foundry增量数据集的写入规则，不会触发平台的分桶能力相关校验报错。
• 范围重分区是在DataFrame计算层完成的，不会破坏最终的目录分区结构，同时可以让每个Spark task处理的数据量更均匀，输出文件大小会稳定在128MB上下，不会出现数倍的大小差异。
• 后续如果需要微调文件大小，只需要调整target_single_file_size的取值或者required_partition_num的计算逻辑即可，不需要改动核心写入框架。

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