PySpark列正则提取转Double后仍按字典序而非数值序排序

数据示例 #

+-----------------+
|       columnName|
+-----------------+
|            1 (1)|
|             null|
|          11 (10)|
|            2 (3)|
+-----------------+

注：数据中的数字可以是浮点数或整数，示例为简化使用了整数。

需求 #

• 排序时将列值为null的行始终移至末尾
• 按括号内的数字进行排序

问题现象 #

上述示例中，预期括号内数字的排序顺序为1, 3, 10，但实际排序结果为1, 10, 3，显然是按字典序排序导致（"10"的首字符"1"被优先匹配，而非按完整数字值排序）。尝试将提取结果转为Double类型但未解决问题，需排查原因。

现有代码 #

from pyspark.sql import functions as F
from pyspark.sql.types import Column

def apply_regex(column: Column, pattern: str, group: int):
    extracted = F.regexp_extract(column, pattern, group)
    return F.when(extracted == "", column).otherwise(extracted.cast("double"))

def sort(dataframe, column_name: str, direction: str):
  col = F.col(column_name)
  # Group 3 is the number inside the parentheses
  col = apply_regex(col, r"(\d+(\.\d+)?) \((\d+(\.\d+)?)\)", 3)
  col = F.when(col.isNull(), F.lit(None)).otherwise(col)
  sort_exprs = [col.asc() if direction == "asc" else col.desc()]
  return dataframe.orderBy(*sort_exprs)

问题原因 #

核心问题出在apply_regex函数的类型不一致：

• 当正则匹配成功时，返回的是extracted.cast("double")（数值类型）；
• 当匹配失败（提取结果为空字符串）时，返回的是原始column（字符串类型）。

Spark在处理when表达式时，会自动将两个分支的结果转换为相同类型——这里会把数值转成字符串，最终整个列还是字符串类型，排序自然按字典序进行，而非数值大小。

另外，原排序逻辑没有单独处理null值的位置，默认升序时null会排在最前面，不符合需求。

修复方案 #

修改apply_regex函数确保返回统一的数值类型，同时调整排序逻辑让null始终排在末尾：

from pyspark.sql import functions as F
from pyspark.sql.types import Column

def apply_regex(column: Column, pattern: str, group: int):
    extracted = F.regexp_extract(column, pattern, group)
    # 匹配失败时返回Double类型的null，而非原始字符串列，保证类型统一
    return F.when(extracted != "", extracted.cast("double")).otherwise(F.lit(None).cast("double"))

def sort(dataframe, column_name: str, direction: str):
    col = F.col(column_name)
    # 提取括号内的数字，得到Double类型列
    sorted_col = apply_regex(col, r"(\d+(\.\d+)?) \((\d+(\.\d+)?)\)", 3)
    # 实现null值移至末尾：先按是否为null排序（非null在前），再按数值排序
    if direction == "asc":
        sort_expr = [sorted_col.isNotNull().desc(), sorted_col.asc()]
    else:
        sort_expr = [sorted_col.isNotNull().desc(), sorted_col.desc()]
    return dataframe.orderBy(*sort_expr)

修复说明 #

1. 统一返回类型：修改后apply_regex函数始终返回Double类型，要么是提取到的数值，要么是Double类型的null，避免了Spark的隐式类型转换，确保排序按数值大小进行。
2. 强制null值在末尾：通过sorted_col.isNotNull().desc()作为第一排序条件，非null行会被优先排序，null值自动落在最后，无论升序还是降序都能满足需求。

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