Optimizer.hs

{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}

module DataFrame.Lazy.Internal.Optimizer (optimize) where

import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.Text as T
import qualified DataFrame.Internal.Expression as E
import DataFrame.Lazy.Internal.LogicalPlan
import DataFrame.Lazy.Internal.PhysicalPlan
import DataFrame.Schema (Schema (..), elements)

{- | Optimise a logical plan and lower it to a physical plan: fuse filters, push
predicates toward the scan, drop dead derived columns, then let @toPhysical@
select concrete operators.
-}
optimize :: Int -> LogicalPlan -> PhysicalPlan
optimize batchSz =
    toPhysical batchSz
        . eliminateDeadColumns
        . pushPredicates
        . fuseFilters

-- ---------------------------------------------------------------------------
-- Rule 1: Filter fusion
-- ---------------------------------------------------------------------------

-- | Merge @Filter p1 (Filter p2 child)@ into @Filter (p1 && p2) child@.
fuseFilters :: LogicalPlan -> LogicalPlan
fuseFilters (Filter p1 (Filter p2 child)) =
    fuseFilters (Filter (andExpr p1 p2) (fuseFilters child))
fuseFilters (Filter p child) = Filter p (fuseFilters child)
fuseFilters (Project cols child) = Project cols (fuseFilters child)
fuseFilters (Derive name expr child) = Derive name expr (fuseFilters child)
fuseFilters (Join jt l r left right) =
    Join jt l r (fuseFilters left) (fuseFilters right)
fuseFilters (Aggregate keys aggs child) =
    Aggregate keys aggs (fuseFilters child)
fuseFilters (Sort cols child) = Sort cols (fuseFilters child)
fuseFilters (Limit n child) = Limit n (fuseFilters child)
fuseFilters leaf = leaf

-- | Logical AND of two @Bool@ expressions.
andExpr :: E.Expr Bool -> E.Expr Bool -> E.Expr Bool
andExpr =
    E.Binary
        ( E.MkBinaryOp
            { E.binaryFn = (&&)
            , E.binaryName = "and"
            , E.binarySymbol = Just "&&"
            , E.binaryCommutative = True
            , E.binaryPrecedence = 3
            }
        )

-- ---------------------------------------------------------------------------
-- Rule 2: Predicate pushdown
-- ---------------------------------------------------------------------------

{- | Push Filter nodes as close to the Scan as possible: past a @Derive@ it
doesn't reference, past a @Project@ that keeps all its columns, and into the
@Scan@'s pushdown predicate.
-}
pushPredicates :: LogicalPlan -> LogicalPlan
pushPredicates (Filter p (Derive name expr child))
    | name `notElem` E.getColumns p =
        Derive name expr (pushPredicates (Filter p child))
    | otherwise =
        Filter p (Derive name expr (pushPredicates child))
pushPredicates (Filter p (Project cols child))
    | all (`elem` cols) (E.getColumns p) =
        Project cols (pushPredicates (Filter p child))
    | otherwise =
        Filter p (Project cols (pushPredicates child))
pushPredicates (Filter p child) = Filter p (pushPredicates child)
pushPredicates (Project cols child) = Project cols (pushPredicates child)
pushPredicates (Derive name expr child) = Derive name expr (pushPredicates child)
pushPredicates (Join jt l r left right) =
    Join jt l r (pushPredicates left) (pushPredicates right)
pushPredicates (Aggregate keys aggs child) =
    Aggregate keys aggs (pushPredicates child)
pushPredicates (Sort cols child) = Sort cols (pushPredicates child)
pushPredicates (Limit n child) = Limit n (pushPredicates child)
pushPredicates leaf = leaf

-- ---------------------------------------------------------------------------
-- Rule 3: Dead column elimination
-- ---------------------------------------------------------------------------

{- | Collect every column name referenced anywhere in the plan (filters, sorts,
aggregate/join keys, projections, derived expressions). 'Nothing' means all
columns are needed (no Project restricts the output).
-}
referencedCols :: LogicalPlan -> Maybe (S.Set T.Text)
referencedCols (Scan _ schema) = Just (S.fromList (M.keys (elements schema)))
referencedCols (Project cols _) = Just (S.fromList cols)
referencedCols (Filter p child) =
    fmap (S.union (S.fromList (E.getColumns p))) (referencedCols child)
referencedCols (Derive _ expr child) =
    fmap (S.union (S.fromList (uExprCols expr))) (referencedCols child)
referencedCols (Join _ l r left right) =
    let keySet = S.fromList [l, r]
        lRef = fmap (S.union keySet) (referencedCols left)
        rRef = fmap (S.union keySet) (referencedCols right)
     in liftMaybe2 S.union lRef rRef
referencedCols (Aggregate keys aggs child) =
    let aggCols = S.fromList (keys <> concatMap (uExprCols . snd) aggs)
     in fmap (S.union aggCols) (referencedCols child)
referencedCols (Sort cols child) =
    fmap (S.union (S.fromList (fmap fst cols))) (referencedCols child)
referencedCols (Limit _ child) = referencedCols child
referencedCols (SourceDF _) = Nothing

liftMaybe2 :: (a -> b -> c) -> Maybe a -> Maybe b -> Maybe c
liftMaybe2 f (Just a) (Just b) = Just (f a b)
liftMaybe2 _ _ _ = Nothing

uExprCols :: E.UExpr -> [T.Text]
uExprCols (E.UExpr expr) = E.getColumns expr

-- | Drop @Derive@ nodes whose output column is never consumed downstream.
eliminateDeadColumns :: LogicalPlan -> LogicalPlan
eliminateDeadColumns plan = go (referencedCols plan) plan
  where
    go needed (Derive name expr child) =
        case needed of
            Nothing -> Derive name expr (go needed child)
            Just cols
                | name `S.notMember` cols -> go needed child
                | otherwise ->
                    Derive
                        name
                        expr
                        (go (Just (S.union cols (S.fromList (uExprCols expr)))) child)
    go needed (Filter p child) =
        Filter p (go (fmap (S.union (S.fromList (E.getColumns p))) needed) child)
    go _needed (Project cols child) =
        Project cols (go (Just (S.fromList cols)) child)
    go needed (Join jt l r left right) =
        let keySet = fmap (S.union (S.fromList [l, r])) needed
         in Join jt l r (go keySet left) (go keySet right)
    go needed (Aggregate keys aggs child) =
        let aggCols = fmap (S.union (S.fromList (keys <> concatMap (uExprCols . snd) aggs))) needed
         in Aggregate keys aggs (go aggCols child)
    go needed (Sort cols child) =
        Sort cols (go (fmap (S.union (S.fromList (fmap fst cols))) needed) child)
    go needed (Limit n child) = Limit n (go needed child)
    go needed (Scan ds schema) =
        case needed of
            Nothing -> Scan ds schema
            Just cols ->
                Scan ds (Schema (M.filterWithKey (\k _ -> k `S.member` cols) (elements schema)))
    go _ (SourceDF df) = SourceDF df

-- ---------------------------------------------------------------------------
-- Logical → Physical lowering
-- ---------------------------------------------------------------------------

{- | Lower the (already-optimised) logical plan to a physical plan. Joins always
lower to HashJoin; the executor may fall back to SortMerge at runtime.
-}
toPhysical :: Int -> LogicalPlan -> PhysicalPlan
-- A Filter directly on a Scan is folded into the scan's pushdown predicate.
toPhysical batchSz (Filter p (Scan (CsvSource path sep reader) schema)) =
    PhysicalScan
        (CsvSource path sep reader)
        (ScanConfig batchSz sep schema (Just p))
toPhysical batchSz (Scan (CsvSource path sep reader) schema) =
    PhysicalScan
        (CsvSource path sep reader)
        (ScanConfig batchSz sep schema Nothing)
toPhysical batchSz (Filter p (Scan (CsvSourceStreaming path sep reader) schema)) =
    PhysicalScan
        (CsvSourceStreaming path sep reader)
        (ScanConfig batchSz sep schema (Just p))
toPhysical batchSz (Scan (CsvSourceStreaming path sep reader) schema) =
    PhysicalScan
        (CsvSourceStreaming path sep reader)
        (ScanConfig batchSz sep schema Nothing)
toPhysical batchSz (Filter p (Scan (ParquetSource path) schema)) =
    PhysicalScan
        (ParquetSource path)
        (ScanConfig batchSz ',' schema (Just p))
toPhysical batchSz (Scan (ParquetSource path) schema) =
    PhysicalScan
        (ParquetSource path)
        (ScanConfig batchSz ',' schema Nothing)
toPhysical batchSz (Project cols child) =
    PhysicalProject cols (toPhysical batchSz child)
toPhysical batchSz (Filter p child) =
    PhysicalFilter p (toPhysical batchSz child)
toPhysical batchSz (Derive name expr child) =
    PhysicalDerive name expr (toPhysical batchSz child)
toPhysical batchSz (Join jt l r left right) =
    PhysicalHashJoin
        jt
        l
        r
        (toPhysical batchSz left)
        (toPhysical batchSz right)
toPhysical batchSz (Aggregate keys aggs child) =
    PhysicalHashAggregate keys aggs (toPhysical batchSz child)
toPhysical batchSz (Sort cols child) =
    PhysicalSort cols (toPhysical batchSz child)
toPhysical batchSz (Limit n child) =
    PhysicalLimit n (toPhysical batchSz child)
toPhysical batchSz (SourceDF df) = PhysicalSourceDF batchSz df