ChuBlog/benchmark/chusql-engine/src/Main.hs

module Main where

import ChuSQL.Algebra.Eval (evalRelOp)
import ChuSQL.Algebra.Optimize (optimize)
import ChuSQL.Algebra.Planner (translate)
import ChuSQL.Model
import ChuSQL.Syntax.Parser (parseQuery)
import Control.Exception (evaluate)
import Data.IORef (IORef, newIORef, readIORef)
import System.CPUTime (getCPUTime)
import System.Environment (getArgs)
import Text.Printf (printf)

mkUsers :: Int -> Table
mkUsers n =
    Table
        { tableName = "users"
        , tableCols = [("id", TInt), ("name", TStr), ("age", TInt)]
        , tableRows =
            [ [("id", VInt i), ("name", VStr ("user" ++ show i)), ("age", VInt (i `mod` 100))]
            | i <- [1 .. n]
            ]
        }

mkOrders :: Int -> Int -> Table
mkOrders n m =
    Table
        { tableName = "orders"
        , tableCols = [("id", TInt), ("user_id", TInt), ("product", TStr)]
        , tableRows =
            [ [("id", VInt j), ("user_id", VInt (1 + (j - 1) `mod` n)), ("product", VStr ("item" ++ show (j `mod` 50)))]
            | j <- [1 .. m]
            ]
        }

checksum :: Either String [Row] -> Int
checksum (Left _) = -1
checksum (Right rows) = go 0 rows
  where
    go acc [] = acc
    go acc (r : rs) = go (acc + length r) rs

rowCount :: Either String [Row] -> Int
rowCount (Left _) = -1
rowCount (Right rows) = length rows

-- salt 是从 IORef 里读出来的：有了这个"IO 里来的"依赖，GHC 就没法把整段纯计算
-- 提到循环外只算一次 —— 否则批量计时会把 1 次的开销除以 k，测出 0.000 ms 这种假数据
runOnce :: IORef String -> Database -> String -> Bool -> IO Int
runOnce saltRef db sql useOpt = do
    salt <- readIORef saltRef
    let result = case parseQuery (sql ++ salt) >>= translate of
            Left e -> Left e
            Right plan -> evalRelOp db (if useOpt then optimize db plan else plan)
    forced <- evaluate (checksum result)
    return forced

resultRows :: IORef String -> Database -> String -> Bool -> IO Int
resultRows saltRef db sql useOpt = do
    salt <- readIORef saltRef
    let result = case parseQuery (sql ++ salt) >>= translate of
            Left e -> Left e
            Right plan -> evalRelOp db (if useOpt then optimize db plan else plan)
    _ <- evaluate (checksum result)
    return (rowCount result)

timed :: IO a -> IO Double
timed act = do
    t0 <- getCPUTime
    _ <- act
    t1 <- getCPUTime
    return (fromIntegral (t1 - t0) / 1e12)

benchOp :: Double -> IO Int -> IO (Double, Int)
benchOp target act = loop 1
  where
    loop k = do
        t <- timed (mapM_ (const act) [1 .. k])
        if t < target then loop (k * 2) else return (t / fromIntegral k, k)

report :: IORef String -> Database -> Double -> String -> IO ()
report saltRef db target sql = do
    nUn <- resultRows saltRef db sql False
    nOp <- resultRows saltRef db sql True
    (tUn, kUn) <- benchOp target (runOnce saltRef db sql False)
    (tOp, kOp) <- benchOp target (runOnce saltRef db sql True)
    let speedup = if tOp <= 0 then 0 else tUn / tOp
        mismatch = if nUn == nOp then "" else "   <<< 结果不一致！"
    printf "%s\n" sql
    printf
        "    结果行数 %-6d | 未优化 %9.3f ms | 优化后 %9.3f ms | 加速 %6.1fx%s\n"
        nOp
        (tUn * 1000)
        (tOp * 1000)
        speedup
        mismatch
    printf "    （实测批次：未优化 %d 次/批，优化后 %d 次/批）\n\n" kUn kOp

queries :: [String]
queries =
    [ "SELECT u.name FROM users u JOIN orders o ON u.id = o.user_id WHERE u.age > 90"
    , "SELECT u.name, o.product FROM users u JOIN orders o ON u.id = o.user_id"
    , "SELECT u.name FROM users u JOIN orders o ON u.id = o.user_id WHERE u.age > 90 ORDER BY o.product LIMIT 10"
    , "SELECT name FROM users WHERE age > 90"
    ]

warmupSql :: String
warmupSql = "SELECT name FROM users"

main :: IO ()
main = do
    args <- getArgs
    let pick i d = if length args > i then read (args !! i) else d
        n = pick 0 500
        m = pick 1 500
        target = pick 2 (0.3 :: Double)
        db = [("users", mkUsers n), ("orders", mkOrders n m)]
    printf "数据规模：users = %d 行，orders = %d 行\n" n m
    printf "计时方式：CPU 时间；每个用例自动加倍批量次数，直到一批累计耗时超过 %.2f 秒，再折算成单次耗时\n" target
    printf "（先跑一遍预热，避免把首次构造数据的开销算进去）\n\n"
    saltRef <- newIORef ""
    _ <- runOnce saltRef db warmupSql False
    _ <- runOnce saltRef db warmupSql True
    mapM_ (report saltRef db target) queries