如何使用PySpark将tick逐笔交易数据转换为1分钟OHLCV DataFrame

需求说明 #

需要将逐笔tick交易数据聚合为1分钟粒度的OHLCV（Open、High、Low、Close、Volume）格式PySpark DataFrame。

现有信息 #

原始输入数据样例 #

+---------+-------+--------+----------+-------------+
|       id|  price|  volume|   capital|    timestamp|
+---------+-------+--------+----------+-------------+
|237367441|9351.71|0.043982| 411.30692|1580515200518|
|237367442|9351.71|0.001413| 13.213966|1580515200580|
|237367443|9352.86|0.001236| 11.560135|1580515200690|
|237367444|9352.85|0.001976| 18.481232|1580515201148|
|237367445| 9352.0|1.214703| 11359.902|1580515201176|
|...|...|...|...|...|
+---------+-------+--------+----------+-------------+
仅展示前20行

目标输出样例 #

+-------------------+---------+---------+---------+---------+------------+
|           datetime|     open|     high|      low|    close|      volume|
+-------------------+---------+---------+---------+---------+------------+
|2020-02-01 00:00:00| 9351.710| 9375.000| 9351.710| 9358.590|  532596.780|
|2020-02-01 00:01:00| 9359.880| 9366.730| 9356.750| 9359.960|  127918.275|
|2020-02-01 00:02:00| 9359.400| 9361.170| 9356.750| 9356.800|   93449.282|
|...|...|...|...|...|...|
+-------------------+---------+---------+---------+---------+------------+

已有代码 #

时间转换工具类

import datetime
class Datetime_manager():
    def timestamp_to_datetime_utc(x):
        return datetime.datetime.utcfromtimestamp(x/ 1000.0).strftime("%Y-%m-%d %H:%M:%S")

主逻辑代码

from Helpers.Datetime_manager import Datetime_manager
import pyspark.sql.functions as func
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType, FloatType, LongType

# 初始化SparkSession
spark = SparkSession.builder.appName('tick_to_ohlcv').getOrCreate()

# 定义schema
schema = StructType([
    StructField("id", IntegerType()), 
    StructField("price", FloatType()),
    StructField("volume", FloatType()),
    StructField("capital", FloatType()),
    StructField("timestamp", LongType()),
])

# 读取数据
df = spark.read.csv('file.csv', sep=",", schema=schema)

# 转换时间格式
get_datetime = func.udf(lambda x : Datetime_manager.timestamp_to_datetime_utc(x))
df_ = df.withColumn("datetime", get_datetime(df.timestamp))
deleted_column = df_.drop('timestamp')
deleted_column.show()

聚合实现代码 #

方案1：兼容现有代码（使用自定义UDF） #

在已有代码基础上添加以下逻辑：

from pyspark.sql import Window

# 1. 生成分钟粒度的时间分组字段
df_with_minute = deleted_column.withColumn(
    "minute_time",
    func.date_trunc("minute", func.col("datetime").cast("timestamp"))
)

# 2. 定义窗口：按分钟分组，按原始交易id排序（id自增可保证交易顺序）
window_spec = Window.partitionBy("minute_time").orderBy("id")

# 3. 计算每个分组内的指标
df_with_rn = df_with_minute \
    .withColumn("rn", func.row_number().over(window_spec)) \
    .withColumn("max_rn", func.max("rn").over(Window.partitionBy("minute_time"))) \
    .withColumn("open", func.first("price").over(window_spec)) \
    .withColumn("high", func.max("price").over(Window.partitionBy("minute_time"))) \
    .withColumn("low", func.min("price").over(Window.partitionBy("minute_time"))) \
    .withColumn("close", func.last("price").over(window_spec)) \
    .withColumn("volume", func.sum("volume").over(Window.partitionBy("minute_time")))

# 4. 去重，每个分钟只保留一行结果
ohlcv_df = df_with_rn \
    .filter(func.col("rn") == func.col("max_rn")) \
    .select(
        func.col("minute_time").alias("datetime"),
        func.round("open", 3).alias("open"),
        func.round("high", 3).alias("high"),
        func.round("low", 3).alias("low"),
        func.round("close", 3).alias("close"),
        func.round("volume", 3).alias("volume")
    ) \
    .orderBy("datetime")

# 输出结果
ohlcv_df.show()

方案2：性能优化版（替换自定义UDF为PySpark内置函数） #

自定义UDF在大数据量下性能较差，可直接使用PySpark内置时间函数处理时间戳，无需引入外部工具类：

import pyspark.sql.functions as func
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType, FloatType, LongType

# 初始化SparkSession
spark = SparkSession.builder.appName('tick_to_ohlcv').getOrCreate()

# 定义schema
schema = StructType([
    StructField("id", IntegerType()), 
    StructField("price", FloatType()),
    StructField("volume", FloatType()),
    StructField("capital", FloatType()),
    StructField("timestamp", LongType()),
])

# 读取数据
df = spark.read.csv('file.csv', sep=",", schema=schema)

# 直接用内置函数转换时间戳并截断到分钟
df_with_minute = df.withColumn(
    "minute_time",
    func.date_trunc("minute", func.to_timestamp(func.col("timestamp")/1000))
)

# 聚合逻辑和方案1一致
window_spec = Window.partitionBy("minute_time").orderBy("id")

ohlcv_df = df_with_minute \
    .withColumn("rn", func.row_number().over(window_spec)) \
    .withColumn("max_rn", func.max("rn").over(Window.partitionBy("minute_time"))) \
    .withColumn("open", func.first("price").over(window_spec)) \
    .withColumn("high", func.max("price").over(Window.partitionBy("minute_time"))) \
    .withColumn("low", func.min("price").over(Window.partitionBy("minute_time"))) \
    .withColumn("close", func.last("price").over(window_spec)) \
    .withColumn("volume", func.sum("volume").over(Window.partitionBy("minute_time"))) \
    .filter(func.col("rn") == func.col("max_rn")) \
    .select(
        func.col("minute_time").alias("datetime"),
        func.round("open", 3).alias("open"),
        func.round("high", 3).alias("high"),
        func.round("low", 3).alias("low"),
        func.round("close", 3).alias("close"),
        func.round("volume", 3).alias("volume")
    ) \
    .orderBy("datetime")

ohlcv_df.show()

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