PySpark 2.1.1中如何让inferSchema识别带前导0的数值列为字符串？

Spark 2.1.1原生的inferSchema=True逻辑确实没办法直接识别带前导0的数值并将其推断为字符串类型——它的推断逻辑是优先判断数据能否被解析为数值类型（整数、浮点数），一旦符合就会自动转换，完全忽略前导0这类格式信息。不过我们可以通过自定义Schema推断逻辑来实现你的需求，不用手动硬编码每一列的类型。

下面是具体的实现步骤和代码：

步骤1：先以全字符串格式读取原始数据 #

首先我们要确保原始数据的所有格式都被保留，所以先关闭自动推断，把所有列都读成字符串类型：

# 注意如果你的CSV有表头，记得加上header=True
df_raw = spark.read.csv('/some/file.csv', header=True, inferSchema=False)

步骤2：实现自定义Schema推断函数 #

我们可以编写一个函数，对每一列的样本数据进行分析：

• 如果某列所有非空值都是纯数字，但存在至少一个值以0开头且长度大于1 → 设为字符串类型（保留前导0）
• 如果所有非空值都是纯数字且没有长于1的前导0 → 设为整数类型
• 其他情况（包含非数字字符）→ 设为字符串类型

代码示例：

from pyspark.sql.types import StringType, IntegerType, StructType, StructField
from pyspark.sql.functions import col

def custom_infer_schema(df, sample_size=1000):
    # 抽取样本数据，避免全量扫描提升性能
    sample_df = df.limit(sample_size)
    schema_fields = []
    
    for column in df.columns:
        # 获取该列非空的样本值
        non_null_values = sample_df.select(col(column)) \
                                   .filter(col(column).isNotNull()) \
                                   .rdd.flatMap(lambda row: row) \
                                   .collect()
        
        if not non_null_values:
            # 空列默认设为字符串类型
            schema_fields.append(StructField(column, StringType(), nullable=True))
            continue
        
        # 检查所有值是否为纯数字
        all_numeric = all(val.isdigit() for val in non_null_values)
        if all_numeric:
            # 检查是否存在带前导0的长数值
            has_leading_zero = any(len(val) > 1 and val.startswith('0') for val in non_null_values)
            if has_leading_zero:
                schema_fields.append(StructField(column, StringType(), nullable=True))
            else:
                # 尝试转换为整数，验证是否符合整数范围
                try:
                    [int(val) for val in non_null_values]
                    schema_fields.append(StructField(column, IntegerType(), nullable=True))
                except ValueError:
                    # 如果转换失败（比如超大数）， fallback到字符串类型
                    schema_fields.append(StructField(column, StringType(), nullable=True))
        else:
            # 包含非数字字符，设为字符串类型
            schema_fields.append(StructField(column, StringType(), nullable=True))
    
    return StructType(schema_fields)

步骤3：应用自定义Schema转换数据 #

拿到推断好的Schema后，我们可以直接转换原始全字符串DataFrame的列类型：

# 生成自定义Schema
custom_schema = custom_infer_schema(df_raw)

# 转换列类型得到最终DataFrame
df_final = df_raw.select([
    col(column).cast(custom_schema[column].dataType).alias(column) 
    for column in df_raw.columns
])

注意事项 #

1. 样本大小调整：sample_size参数可以根据你的数据量调整，如果数据分布不均匀，建议用sample(withReplacement=False, fraction=0.1)抽样替代limit()，避免漏抽到带前导0的行导致误判。
2. 小数扩展：如果你的数据包含小数（比如001.23），上面的isdigit()判断会失效，你可以扩展逻辑，用正则表达式匹配有效的数值格式（比如^-?\d+(\.\d+)?$）来判断是否为数值类型。
3. 性能平衡：自定义推断需要扫描样本数据，对于超大规模数据集，建议合理控制样本比例，在准确性和性能之间做平衡。

另外，如果你的原始CSV文件中，带前导0的列是用引号包裹的（比如"001"），Spark默认会将其识别为字符串类型，但如果源头数据无法修改格式，上面的自定义推断方案就是最可靠的自动处理方式。

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