Haskell Cassava处理CSV：保留任意额外列的实现方法

问题场景 #

现有两个CSV文件：

A.csv

A,B,C,D,E
1,2,3,4,5
5,4,3,2,1

B.csv

A,E,B,C,F
6,7,8,9,1
4,3,4,5,6

需要用Haskell的Cassava库实现以下需求：

• 对A、B、C三列执行严格解析（必须存在且为Int类型）
• 合并两个文件的行（可按需添加过滤逻辑）
• 最终输出的CSV要保留所有出现过的列（包括D、E、F这类未知的额外列），缺失值用占位符填充

期望输出的C.csv：

A,B,C,D,E,F
1,2,3,4,5,_
5,4,3,2,1,_
6,8,9,_,7,1
4,4,5,_,3,6

现有代码只能输出A、B、C三列，无法保留额外列，需要改造。

解决方案 #

完全可以用Cassava实现该需求，无需从头编写自定义解析器，只需扩展数据结构并调整解析/序列化逻辑即可。核心思路是将额外列存储为NamedRecord（Cassava内置的HashMap类型），在解析时保留这些列，序列化时合并固定列与额外列，最后统一填充缺失值。

关键改动说明 #

1. 扩展数据结构：在ABC中增加addCols字段，用于存储所有非A/B/C的列：

   data ABC = ABC {a :: Int, b :: Int, c :: Int, addCols :: C.NamedRecord} deriving Show
   

2. 解析逻辑调整：解析NamedRecord时，提取A/B/C列后，用HM.difference把剩余列存入addCols：

   instance C.FromNamedRecord ABC where
       parseNamedRecord m =
           pure ABC
           <*> m C..: "A"
           <*> m C..: "B"
           <*> m C..: "C"
           <*> pure (m `HM.difference` abcRefHashMap)
   

   其中abcRefHashMap是A/B/C列的集合，用于筛选出额外列。

3. 序列化逻辑调整：将固定列的NamedRecord与额外列的NamedRecord合并，确保所有列都被输出：

   instance C.ToNamedRecord ABC where
       toNamedRecord m =
           (addCols m) `HM.union` C.namedRecord ["A" C..= a m, "B" C..= b m, "C" C..= c m]
   

4. 生成完整表头：收集所有出现过的额外列，和A/B/C合并成完整表头：

   makeCompleteHeader :: [ABC] -> C.Header
   makeCompleteHeader ms = V.fromList $ abcDefinedCols ++ HM.keys (HM.unions (map addCols ms))
   

5. 填充缺失值：遍历所有行，为缺失的额外列添加占位符：

   combineABCs :: [ABC] -> [ABC] -> [ABC]
   combineABCs xs1 xs2 =
       let simpleSum = xs1 ++ xs2
           addColKeys = HM.keys (HM.unions (map addCols simpleSum))
           toAddHashMap = HM.fromList $ map (\k -> (k, "_")) addColKeys
       in map (\x -> x { addCols = fillAddCols (addCols x) toAddHashMap }) simpleSum
       where
           fillAddCols :: C.NamedRecord -> C.NamedRecord -> C.NamedRecord
           fillAddCols cur toAdd = HM.union cur (toAdd `HM.difference` cur)
   

完整代码 #

#!/usr/bin/env stack
-- stack --resolver lts-18.7 script --package cassava,bytestring,vector,unordered-containers

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

import qualified Data.ByteString.Lazy as B
import qualified Data.Csv as C
import qualified Data.HashMap.Strict as HM
import qualified Data.Vector as V

data ABC = ABC {a :: Int, b :: Int, c :: Int, addCols :: C.NamedRecord} deriving Show

abcDefinedCols = ["A", "B", "C"]
abcRefHashMap = HM.fromList $ map (\x -> (x, ())) abcDefinedCols

instance C.FromNamedRecord ABC where
    parseNamedRecord m =
        pure ABC
        <*> m C..: "A"
        <*> m C..: "B"
        <*> m C..: "C"
        <*> pure (m `HM.difference` abcRefHashMap)

instance C.ToNamedRecord ABC where
    toNamedRecord m =
        (addCols m) `HM.union` C.namedRecord ["A" C..= a m, "B" C..= b m, "C" C..= c m]

decodeABC :: B.ByteString -> [ABC]
decodeABC x =
    case C.decodeByName x of
        Left err -> error err
        Right (_,xs) -> V.toList xs

makeCompleteHeader :: [ABC] -> C.Header
makeCompleteHeader ms = V.fromList $ abcDefinedCols ++ HM.keys (HM.unions (map addCols ms))

combineABCs :: [ABC] -> [ABC] -> [ABC]
combineABCs xs1 xs2 =
    let simpleSum = xs1 ++ xs2
        addColKeys = HM.keys (HM.unions (map addCols simpleSum))
        toAddHashMap = HM.fromList $ map (\k -> (k, "_")) addColKeys
    in map (\x -> x { addCols = fillAddCols (addCols x) toAddHashMap }) simpleSum
    where
        fillAddCols :: C.NamedRecord -> C.NamedRecord -> C.NamedRecord
        fillAddCols cur toAdd = HM.union cur (toAdd `HM.difference` cur)

main :: IO ()
main = do
    fileA <- B.readFile "A.csv"
    fileB <- B.readFile "B.csv"
    let decodedA = decodeABC fileA
    let decodedB = decodeABC fileB
    putStrLn $ show decodedA
    putStrLn $ show decodedB
    let ab = combineABCs decodedA decodedB
    B.writeFile "C.csv" $ C.encodeByName (makeCompleteHeader ab) ab

补充说明 #

• 代码中使用"_"作为缺失值占位符，可根据需求替换为其他值（比如"n/a"）
• 若需要添加过滤逻辑，可在combineABCs函数中对simpleSum进行过滤处理
• 依赖的包包括cassava、bytestring、vector和unordered-containers，stack脚本已自动处理依赖

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