子查询去相关

chusql-core/engine/src/ChuSQL/Core/Engine/Algebra/Eval.hs

@@ -1,7 +1,7 @@
 module ChuSQL.Core.Engine.Algebra.Eval (evalRelOp, evalRelOpM, evalExprM, evalCondForRowM) where
 
 import ChuSQL.Core.Engine.Algebra.Cost (tableMetaOf)
-import ChuSQL.Core.Engine.Algebra.Expr (evalCondForRow, evalExpr, quantifiedCompare)
+import ChuSQL.Core.Engine.Algebra.Expr (evalCondForRow, evalExpr, hasSubquery, quantifiedCompare, subqueryRefsIn)
 import ChuSQL.Core.Engine.Algebra.Optimize (optimize)
 import ChuSQL.Core.Engine.Algebra.Op
 import ChuSQL.Core.Engine.Algebra.Planner (translate)
@@ -12,6 +12,7 @@ import ChuSQL.Core.Engine.Storage (IndexResult (..), MonadStorage (..))
 import ChuSQL.Core.Engine.Storage.Memory (runMemoryStorage)
 import ChuSQL.Core.Engine.Syntax.AST (Expr (..), JoinKind (..), Subquery (..))
 import Control.Monad (filterM)
+import Data.List (nub)
 import qualified Data.HashMap.Strict as HM
 
 -- 执行：按算子树算出结果行，外层行供相关子查询引用。
@@ -246,6 +247,53 @@ evalAggregate outer keys aggs rows = do
     let top = if null keys && null groups then [([], [])] else groups
     mapM (oneGroup outer keys aggs) top
 
+-- | 绑定键：子查询引用的外层列在这一行上的取值
+bindingKey :: [String] -> Row -> String
+bindingKey refs env = show [(c, lookup c env) | c <- refs]
+
+-- | 每个绑定留第一行当样本，去掉重复绑定
+bindingSamples :: [String] -> Row -> [Row] -> [(String, Row)]
+bindingSamples refs outer rows = go rows HM.empty []
+  where
+    -- | 逐行按绑定键去重
+    go [] _ acc = reverse acc
+    go (r : rs) seen acc =
+        let env = r ++ outer
+            k = bindingKey refs env
+        in case HM.lookup k seen of
+            Just _ -> go rs seen acc
+            Nothing -> go rs (HM.insert k () seen) ((k, env) : acc)
+
+-- | 每个绑定只换一次子查询，得到绑定到表达式的表
+bindingsTableM ::
+    (MonadStorage m) =>
+    Database ->
+    Row ->
+    [Expr] ->
+    [Row] ->
+    m (Either String (HM.HashMap String [Expr]))
+bindingsTableM db outer exprs rows = do
+    built <- mapM swap (bindingSamples refs outer rows)
+    pure (fmap HM.fromList (sequence built))
+  where
+    refs = nub (concatMap subqueryRefsIn exprs)
+    -- | 一个绑定样本上把子查询全换成字面量
+    swap (key, env) = do
+        es <- mapM (\e -> substSubqueries db e env) exprs
+        pure (fmap (\xs -> (key, xs)) (sequence es))
+
+-- | 取这一行绑定换好的表达式；没有这个绑定就是内部错误
+boundExprs :: [String] -> Row -> HM.HashMap String [Expr] -> Either String [Expr]
+boundExprs refs env table = case HM.lookup (bindingKey refs env) table of
+    Just es -> Right es
+    Nothing -> Left "subquery binding is missing"
+
+-- | 取这一行绑定换好的唯一表达式
+boundExpr :: [String] -> Row -> HM.HashMap String [Expr] -> Either String Expr
+boundExpr refs env table = case boundExprs refs env table of
+    Right [x] -> Right x
+    _ -> Left "subquery binding is missing"
+
 -- | 单子过滤：条件本身是单子求值的
 filterRowsM :: (Monad m) => (Row -> m (Either String Bool)) -> [Row] -> m (Either String [Row])
 filterRowsM p = go []
@@ -385,9 +433,19 @@ evalRelOpIn db outer (Compute items op) = do
     rows <- evalRelOpIn db outer op
     case rows of
         Left e -> pure (Left e)
-        Right rs -> do
-            computed <- mapM (computeM db outer items) rs
-            pure (sequence computed)
+        Right rs
+            | not (any (hasSubquery . snd) items) -> do
+                computed <- mapM (computeM db outer items) rs
+                pure (sequence computed)
+            | otherwise -> do
+                table <- bindingsTableM db outer (map snd items) rs
+                pure (table >>= \t -> traverse (computeBound t) rs)
+  where
+    refs = nub (concatMap (subqueryRefsIn . snd) items)
+    -- | 用这一行绑定换好的表达式算出一行
+    computeBound table r = do
+        es <- boundExprs refs (r ++ outer) table
+        mapM (\(label, e) -> (,) label <$> evalExpr e (r ++ outer)) (zip (map fst items) es)
 evalRelOpIn db outer (Aggregate keys aggs op) = do
     rows <- evalRelOpIn db outer op
     pure (rows >>= evalAggregate outer keys aggs)
@@ -410,7 +468,17 @@ evalRelOpIn db outer (Filter e op) = do
     rows <- evalRelOpIn db outer op
     case rows of
         Left err -> pure (Left err)
-        Right rs -> filterRowsM (\r -> evalCondForRowM db e (r ++ outer)) rs
+        Right rs
+            | not (hasSubquery e) -> filterRowsM (\r -> evalCondForRowM db e (r ++ outer)) rs
+            | otherwise -> do
+                table <- bindingsTableM db outer [e] rs
+                pure (table >>= \t -> filterM (keepBound t) rs)
+  where
+    refs = nub (subqueryRefsIn e)
+    -- | 用这一行绑定换好的条件判定
+    keepBound table r = do
+        e' <- boundExpr refs (r ++ outer) table
+        evalCondForRow e' (r ++ outer)
 evalRelOpIn db outer (Project cols op) = do
     rows <- evalRelOpIn db outer op
     pure (fmap (map (project cols)) rows)

chusql-core/engine/src/ChuSQL/Core/Engine/Algebra/Expr.hs

@@ -1,4 +1,4 @@
-module ChuSQL.Core.Engine.Algebra.Expr (evalExpr, evalCondForRow, colsInExpr, aggregatesIn, hasAggregate, bareColumns, hasSubquery, hasDivision, inValues, threeValuedNot, quantifiedCompare) where
+module ChuSQL.Core.Engine.Algebra.Expr (evalExpr, evalCondForRow, colsInExpr, subqueryRefsIn, aggregatesIn, hasAggregate, bareColumns, hasSubquery, hasDivision, inValues, threeValuedNot, quantifiedCompare) where
 
 import ChuSQL.Core.Model
 import ChuSQL.Core.Engine.Syntax.AST
@@ -90,6 +90,39 @@ colsInExpr (ExistsSub sq _) = subqueryRefs sq
 colsInExpr (QuantCmp _ a sq _) = colsInExpr a ++ subqueryRefs sq
 colsInExpr _ = []
 
+-- | 表达式里子查询引用的外层列
+subqueryRefsIn :: Expr -> [String]
+subqueryRefsIn e = case e of
+    ScalarSub sq -> subqueryRefs sq
+    InSub a sq _ -> subqueryRefsIn a ++ subqueryRefs sq
+    InList a es _ -> concatMap subqueryRefsIn (a : es)
+    ExistsSub sq _ -> subqueryRefs sq
+    QuantCmp _ a sq _ -> subqueryRefsIn a ++ subqueryRefs sq
+    Add a b -> both a b
+    Sub a b -> both a b
+    Mul a b -> both a b
+    Div a b -> both a b
+    Neg a -> subqueryRefsIn a
+    Gt a b -> both a b
+    Lt a b -> both a b
+    Eq a b -> both a b
+    GtE a b -> both a b
+    LtE a b -> both a b
+    NotEq a b -> both a b
+    And a b -> both a b
+    Or a b -> both a b
+    IsNull a -> subqueryRefsIn a
+    IsNotNull a -> subqueryRefsIn a
+    CountOf a -> subqueryRefsIn a
+    SumOf a -> subqueryRefsIn a
+    AvgOf a -> subqueryRefsIn a
+    MinOf a -> subqueryRefsIn a
+    MaxOf a -> subqueryRefsIn a
+    _ -> []
+  where
+    -- | 两个子表达式的外层列拼起来
+    both a b = subqueryRefsIn a ++ subqueryRefsIn b
+
 -- | 表达式里有没有除法（除法可能报错，不能提前或延后求值）
 hasDivision :: Expr -> Bool
 hasDivision e = case e of

chusql-core/engine/test/Spec.hs

@@ -16,6 +16,7 @@ import ChuSQL.Core.Engine.Parallel (poolRun, workerPool)
 import ChuSQL.Core.Model
 import ChuSQL.Core.Engine.Semantic (prepare)
 import ChuSQL.Core.Engine.Storage (IndexResult (..), MonadStorage (..))
+import ChuSQL.Core.Engine.Storage.Memory (MemoryStorage (runMemoryStorage))
 import ChuSQL.Core.Engine.Storage.IPC (IPCStorage (..), closeConnection, defaultEnv, doListTables, envForDatabase, getDatabaseName, localStorageLink, runIPCStorage, setDatabaseName, setStorageLink)
 import ChuSQL.Core.Engine.Syntax.AST
 import ChuSQL.Core.Engine.Syntax.Parser
@@ -1400,6 +1401,24 @@ main = hspec $ do
             sameResultAsUnoptimized "SELECT name FROM users WHERE id IN (SELECT user_id FROM orders)"
             sameResultAsUnoptimized "SELECT u.name FROM users u WHERE EXISTS (SELECT * FROM orders o WHERE o.user_id = u.id)"
 
+    describe "ChuSQL.Core.Engine (子查询去相关)" $ do
+        it "runs a non-correlated subquery only once" $ do
+            let (result, scans) = countedRun "SELECT name FROM users WHERE id IN (SELECT user_id FROM orders)"
+            fmap (map (map snd)) result `shouldBe` Right [[VStr "Alice"], [VStr "Bob"]]
+            scanTimes "orders" scans `shouldBe` 1
+
+        it "runs a correlated subquery once per distinct binding" $ do
+            let (result, scans) =
+                    countedRun "SELECT o.product FROM orders o WHERE EXISTS (SELECT * FROM users u WHERE u.id = o.user_id)"
+            fmap (map (map snd)) result `shouldBe` Right [[VStr "Book"], [VStr "Pen"], [VStr "Cup"]]
+            scanTimes "users" scans `shouldBe` 2
+            scanTimes "orders" scans `shouldBe` 1
+
+        it "runs a non-correlated subquery in a computed column only once" $ do
+            let (result, scans) = countedRun "SELECT name, (SELECT COUNT(*) FROM orders) FROM users"
+            length (either (const []) id result) `shouldBe` 3
+            scanTimes "orders" scans `shouldBe` 1
+
     describe "ChuSQL.Core.Engine (聚合与 GROUP BY)" $ do
         -- sql：在默认库上跑 SQL 取结果行
         let sql input = parseStatement input >>= rowsOf . runStatement testDB
@@ -2926,3 +2945,88 @@ leftSideFilters _ = Nothing
 planRoot :: RelOp -> RelOp
 planRoot (Project _ x) = planRoot x
 planRoot x = x
+
+-- | 记下扫过哪些表的存储：语义与内存实现一致，只是扫描多记一笔
+data Counting a = Counting {runCounting :: Database -> Either String (a, Database, [String])}
+
+instance Functor Counting where
+    fmap f (Counting m) = Counting $ \db -> fmap (\(a, db', log') -> (f a, db', log')) (m db)
+
+instance Applicative Counting where
+    pure a = Counting $ \db -> Right (a, db, [])
+    Counting mf <*> Counting ma = Counting $ \db -> do
+        (f, db1, l1) <- mf db
+        (a, db2, l2) <- ma db1
+        Right (f a, db2, l1 ++ l2)
+
+instance Monad Counting where
+    Counting m >>= k = Counting $ \db -> do
+        (a, db1, l1) <- m db
+        (b, db2, l2) <- runCounting (k a) db1
+        Right (b, db2, l1 ++ l2)
+
+-- | 内存实现的某个动作，转发给 Counting
+memory :: MemoryStorage a -> Counting a
+memory m = Counting $ \db -> case runMemoryStorage m db of
+    Left e -> Left e
+    Right (a, db') -> Right (a, db', [])
+
+-- | 内存实现的扫描
+memScan :: String -> MemoryStorage (Either String [Row])
+memScan = scan
+
+-- | 内存实现的插入
+memInsert :: String -> Row -> MemoryStorage (Either String ())
+memInsert = insert
+
+-- | 内存实现的整表替换
+memReplaceAll :: String -> [Row] -> MemoryStorage (Either String ())
+memReplaceAll = replaceAll
+
+-- | 内存实现的建表
+memCreateTable :: String -> [(String, Column)] -> MemoryStorage (Either String ())
+memCreateTable = createTable
+
+-- | 内存实现的删表
+memDropTable :: String -> MemoryStorage (Either String ())
+memDropTable = dropTable
+
+-- | 内存实现的删列
+memDropColumn :: String -> String -> MemoryStorage (Either String ())
+memDropColumn = dropColumn
+
+-- | 内存实现的换结构
+memReplaceSchema :: String -> [(String, Column)] -> [Row] -> MemoryStorage (Either String ())
+memReplaceSchema = replaceSchema
+
+instance MonadStorage Counting where
+    -- \| 扫描：记一笔表名，其余走内存实现
+    scan t = Counting $ \db -> case runMemoryStorage (memScan t) db of
+        Left e -> Left e
+        Right (r, db') -> Right (r, db', [t])
+    -- \| 插入：走内存实现
+    insert t r = memory (memInsert t r)
+    -- \| 整表替换：走内存实现
+    replaceAll t rs = memory (memReplaceAll t rs)
+    -- \| 建表：走内存实现
+    createTable name cols = memory (memCreateTable name cols)
+    -- \| 删表：走内存实现
+    dropTable name = memory (memDropTable name)
+    -- \| 删列：走内存实现
+    dropColumn name col = memory (memDropColumn name col)
+    -- \| 换结构：走内存实现
+    replaceSchema name cols rows = memory (memReplaceSchema name cols rows)
+    -- \| 快照：把当前库原样交出去
+    snapshot = Counting $ \db -> Right (db, db, [])
+
+-- | 解析一条 SQL 并在计数存储上跑，返回结果与扫描过的表
+countedRun :: String -> (Either String [Row], [String])
+countedRun input = case parseStatement input of
+    Left err -> (Left err, [])
+    Right stmt -> case runCounting (runStatementM stmt) testDB of
+        Left err -> (Left err, [])
+        Right (result, _, scans) -> (result, scans)
+
+-- | 结果里出现了几张表被扫过几次
+scanTimes :: String -> [String] -> Int
+scanTimes t scans = length [name | name <- scans, name == t]