PySpark中使用ntile函数排名时排除空值并保留所有行

直接通过条件判断结合窗口函数特性，就能实现需求，无需额外关联原数据。以下分两种场景给出方案：

场景1：Spark 3.0及以上版本（支持窗口filter子句） #

Spark 3.0+允许在窗口定义中使用filter，直接过滤掉null值后计算ntile，同时用when给null值返回null：

import numpy as np
import pandas as pd
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("NtileNonNull").getOrCreate()

data = pd.DataFrame({
    'ID':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15, 16, 18, 18, 19, 20],
    'group':['A']*10 + ['B']*10,
    'feature':[0.5, 3.4, 0.3, 0.1, 1.2, np.nan, 2, np.nan, 1.1, 2.2, np.nan, 1.32, 2.5, np.nan, 0.87, 1.56, 2.1, np.nan, 0.34, 7.43]
})

spark_df = spark.createDataFrame(data)

# 定义窗口：分组后按feature排序，仅包含非null值
window = Window.partitionBy('group').orderBy('feature').filter(F.col('feature').isNotNull())

# 仅对非null值计算ntile，null值返回null
spark_df = spark_df.withColumn(
    "feature_rank",
    F.when(F.col('feature').isNotNull(), F.ntile(5).over(window))
)

spark_df.display()

场景2：Spark 3.0以下版本（不支持窗口filter） #

通过计算分组内非null值总数、非null值的行号，手动推导ntile分桶结果：

import numpy as np
import pandas as pd
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("NtileNonNullLegacy").getOrCreate()

data = pd.DataFrame({
    'ID':[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15, 16, 18, 18, 19, 20],
    'group':['A']*10 + ['B']*10,
    'feature':[0.5, 3.4, 0.3, 0.1, 1.2, np.nan, 2, np.nan, 1.1, 2.2, np.nan, 1.32, 2.5, np.nan, 0.87, 1.56, 2.1, np.nan, 0.34, 7.43]
})

spark_df = spark.createDataFrame(data)

# 窗口1：分组统计非null值总数
window_group = Window.partitionBy('group')
spark_df = spark_df.withColumn('non_null_cnt', F.count('feature').over(window_group))

# 窗口2：分组按feature排序，给非null值分配行号
window_order = Window.partitionBy('group').orderBy('feature')
spark_df = spark_df.withColumn('rn', F.when(F.col('feature').isNotNull(), F.row_number().over(window_order)))

# 手动计算ntile：按行号和非null总数的比例分桶，ceil取整
spark_df = spark_df.withColumn(
    "feature_rank",
    F.when(
        F.col('feature').isNotNull(),
        F.ceil(F.col('rn') / (F.col('non_null_cnt') / 5)).cast('int')
    )
)

# 可选：删除中间临时列
spark_df = spark_df.drop('non_null_cnt', 'rn')

spark_df.display()

原理说明 #

• 两种方案都通过when条件判断，仅对feature非null的行计算排名，null行直接返回null。
• Spark 3.0+的方案利用窗口filter直接限定ntile的计算范围，更简洁；低版本则通过行号和总数的比例模拟ntile的分桶逻辑，结果与原生ntile一致。

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