如何在SQL/PySpark中基于FIFO逻辑映射交易以计算资本利得税

步骤1：构造示例交易数据 #

先创建筛选后的单只股票交易DataFrame，模拟你提到的4笔买入和2笔卖出场景：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

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

# 示例交易数据
transactions_data = [
    ("2020-06-18", "Buy", 100, 50.0),
    ("2020-06-29", "Buy", 100, 55.0),
    ("2020-07-15", "Buy", 200, 60.0),
    ("2020-08-02", "Buy", 150, 65.0),
    ("2020-09-10", "Sell", 150, 70.0),
    ("2020-10-05", "Sell", 50, 75.0)
]

transactions_df = spark.createDataFrame(
    transactions_data,
    schema=["trade_date", "trade_type", "quantity", "price"]
).withColumn("trade_date", F.to_date("trade_date"))

步骤2：拆分买入和卖出交易并计算累计量 #

处理买入交易 #

按日期排序买入交易，计算累计持有数量，用于后续区间匹配：

# 筛选买入交易并计算累计数量
buy_df = transactions_df.filter(F.col("trade_type") == "Buy") \
    .withColumn("buy_id", F.row_number().over(Window.orderBy("trade_date"))) \
    .withColumn("cumulative_buy_qty", F.sum("quantity").over(Window.orderBy("trade_date").rowsBetween(Window.unboundedPreceding, 0))) \
    .withColumn("prev_cumulative_buy_qty", F.lag("cumulative_buy_qty", 1, 0).over(Window.orderBy("trade_date")))

处理卖出交易 #

同样按日期排序卖出交易，计算累计卖出数量：

# 筛选卖出交易并计算累计数量
sell_df = transactions_df.filter(F.col("trade_type") == "Sell") \
    .withColumn("sell_id", F.row_number().over(Window.orderBy("trade_date"))) \
    .withColumn("cumulative_sell_qty", F.sum("quantity").over(Window.orderBy("trade_date").rowsBetween(Window.unboundedPreceding, 0))) \
    .withColumn("prev_cumulative_sell_qty", F.lag("cumulative_sell_qty", 1, 0).over(Window.orderBy("trade_date")))

步骤3：FIFO匹配买入与卖出交易 #

通过交叉连接筛选累计量的重叠区间，匹配对应的买入份额，计算实际匹配数量：

# 交叉连接买入和卖出，筛选符合FIFO的匹配区间
matched_df = buy_df.crossJoin(sell_df) \
    .filter(
        (F.col("cumulative_sell_qty") > F.col("prev_cumulative_buy_qty")) &
        (F.col("prev_cumulative_sell_qty") < F.col("cumulative_buy_qty"))
    ) \
    .withColumn(
        "matched_qty",
        F.least(F.col("cumulative_buy_qty"), F.col("cumulative_sell_qty")) -
        F.greatest(F.col("prev_cumulative_buy_qty"), F.col("prev_cumulative_sell_qty"))
    ) \
    .select(
        "sell_id", "trade_date", "quantity", "price",
        "buy_id", F.col("trade_date").alias("buy_date"),
        F.col("quantity").alias("buy_quantity"), F.col("price").alias("buy_price"),
        "matched_qty"
    )

步骤4：生成最终资本利得计算表 #

整理结果并计算每笔匹配的成本、收益和资本利得：

final_result = matched_df \
    .withColumn("cost_basis", F.col("matched_qty") * F.col("buy_price")) \
    .withColumn("proceeds", F.col("matched_qty") * F.col("price")) \
    .withColumn("capital_gain", F.col("proceeds") - F.col("cost_basis")) \
    .select(
        "sell_id", "trade_date", "quantity", "price",
        "buy_date", "buy_price", "matched_qty",
        "cost_basis", "proceeds", "capital_gain"
    )

final_result.show()

执行后会得到符合预期的FIFO匹配结果：

• 第一笔卖出150股匹配2020-06-18的100股，以及2020-06-29的50股
• 第二笔卖出50股匹配2020-06-29剩余的50股

Spark SQL实现方案 #

若偏好SQL语法，可将DataFrame注册为临时视图后执行查询：

transactions_df.createOrReplaceTempView("stock_transactions")

-- 计算买入累计量
CREATE OR REPLACE TEMP VIEW buy_transactions AS
SELECT 
    trade_date,
    trade_type,
    quantity,
    price,
    ROW_NUMBER() OVER (ORDER BY trade_date) AS buy_id,
    SUM(quantity) OVER (ORDER BY trade_date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS cumulative_buy_qty,
    COALESCE(LAG(SUM(quantity) OVER (ORDER BY trade_date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), 1), 0) AS prev_cumulative_buy_qty
FROM stock_transactions
WHERE trade_type = 'Buy';

-- 计算卖出累计量
CREATE OR REPLACE TEMP VIEW sell_transactions AS
SELECT 
    trade_date,
    trade_type,
    quantity,
    price,
    ROW_NUMBER() OVER (ORDER BY trade_date) AS sell_id,
    SUM(quantity) OVER (ORDER BY trade_date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS cumulative_sell_qty,
    COALESCE(LAG(SUM(quantity) OVER (ORDER BY trade_date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW), 1), 0) AS prev_cumulative_sell_qty
FROM stock_transactions
WHERE trade_type = 'Sell';

-- FIFO匹配并计算资本利得
SELECT 
    s.sell_id,
    s.trade_date AS sell_date,
    s.quantity AS sell_quantity,
    s.price AS sell_price,
    b.trade_date AS buy_date,
    b.price AS buy_price,
    LEAST(b.cumulative_buy_qty, s.cumulative_sell_qty) - GREATEST(b.prev_cumulative_buy_qty, s.prev_cumulative_sell_qty) AS matched_qty,
    (LEAST(b.cumulative_buy_qty, s.cumulative_sell_qty) - GREATEST(b.prev_cumulative_buy_qty, s.prev_cumulative_sell_qty)) * b.price AS cost_basis,
    (LEAST(b.cumulative_buy_qty, s.cumulative_sell_qty) - GREATEST(b.prev_cumulative_buy_qty, s.prev_cumulative_sell_qty)) * s.price AS proceeds,
    ((LEAST(b.cumulative_buy_qty, s.cumulative_sell_qty) - GREATEST(b.prev_cumulative_buy_qty, s.prev_cumulative_sell_qty)) * s.price) - 
    ((LEAST(b.cumulative_buy_qty, s.cumulative_sell_qty) - GREATEST(b.prev_cumulative_buy_qty, s.prev_cumulative_sell_qty)) * b.price) AS capital_gain
FROM buy_transactions b
CROSS JOIN sell_transactions s
WHERE s.cumulative_sell_qty > b.prev_cumulative_buy_qty
AND s.prev_cumulative_sell_qty < b.cumulative_buy_qty
ORDER BY s.sell_id, b.buy_id;

内容的提问来源于stack exchange，提问作者Łukasz Kastelik