{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- | Categorical split candidates: Breiman prefixes for binary targets,
subset/singleton enumeration otherwise, and cross-column equality. Each
value-list becomes an OR-of-equalities condition.
-}
module DataFrame.DecisionTree.Categorical (
    TargetInfo (..),
    mkTargetInfo,
    distinctValuesUpTo,
    validBoxedValues,
    orEqs,
    subsetSplits,
    subsetLists,
    singletonSplits,
    singletonLists,
    breimanPrefixSplits,
    breimanPrefixLists,
    catValueLists,
    membershipVec,
    crossColumnConds,
    discreteConditions,
    discreteCondVecs,
) where

import DataFrame.DecisionTree.CondVec (CondVec (..), materializeCondVec)
import DataFrame.DecisionTree.Types (
    ColumnOrdering,
    SynthConfig (..),
    TreeConfig (..),
    withOrdFrom,
 )
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn)
import DataFrame.Internal.Expression (Expr (..))
import DataFrame.Internal.Interpreter (interpret)
import DataFrame.Internal.Types
import DataFrame.Operators

import Data.Either (fromRight)
import Data.Function (on)
import Data.List (inits, sort, sortBy, subsequences)
import qualified Data.Map.Strict as M
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Type.Equality (testEquality, (:~:) (..))
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import Type.Reflection (typeRep)

-- | Valid-slot view of a nullable boxed column (null slots hold crash-thunks).
validBoxedValues :: Bitmap -> V.Vector a -> V.Vector a
validBoxedValues :: forall a. Bitmap -> Vector a -> Vector a
validBoxedValues Bitmap
bm = (Int -> a -> Bool) -> Vector a -> Vector a
forall a. (Int -> a -> Bool) -> Vector a -> Vector a
V.ifilter (\Int
i a
_ -> Bitmap -> Int -> Bool
bitmapTestBit Bitmap
bm Int
i)

{- | Target-column summary driving the categorical generator: binary vs
multi-class, the deterministic positive class, and the raw label vector.
-}
data TargetInfo target = TargetInfo
    { forall target. TargetInfo target -> Bool
tiIsBinary :: !Bool
    , forall target. TargetInfo target -> Maybe target
tiPositiveClass :: !(Maybe target)
    , forall target. TargetInfo target -> Vector target
tiValues :: !(V.Vector target)
    }

{- | Compute 'TargetInfo' once per fit. The positive class for binary targets
is the lexicographically-first distinct value, for deterministic pools.
-}
mkTargetInfo ::
    forall target.
    (Columnable target, Ord target) =>
    T.Text -> DataFrame -> Maybe (TargetInfo target)
mkTargetInfo :: forall target.
(Columnable target, Ord target) =>
Text -> DataFrame -> Maybe (TargetInfo target)
mkTargetInfo Text
target DataFrame
df = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @target DataFrame
df (Text -> Expr target
forall a. Columnable a => Text -> Expr a
Col Text
target) of
    Right (TColumn Column
column) ->
        (DataFrameException -> Maybe (TargetInfo target))
-> (Vector target -> Maybe (TargetInfo target))
-> Either DataFrameException (Vector target)
-> Maybe (TargetInfo target)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (TargetInfo target)
-> DataFrameException -> Maybe (TargetInfo target)
forall a b. a -> b -> a
const Maybe (TargetInfo target)
forall a. Maybe a
Nothing) (TargetInfo target -> Maybe (TargetInfo target)
forall a. a -> Maybe a
Just (TargetInfo target -> Maybe (TargetInfo target))
-> (Vector target -> TargetInfo target)
-> Vector target
-> Maybe (TargetInfo target)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector target -> TargetInfo target
forall target. Ord target => Vector target -> TargetInfo target
targetInfoFromValues) (forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @target Column
column)
    Either DataFrameException (TypedColumn target)
_ -> Maybe (TargetInfo target)
forall a. Maybe a
Nothing

targetInfoFromValues :: (Ord target) => V.Vector target -> TargetInfo target
targetInfoFromValues :: forall target. Ord target => Vector target -> TargetInfo target
targetInfoFromValues Vector target
vals = Bool -> Maybe target -> Vector target -> TargetInfo target
forall target.
Bool -> Maybe target -> Vector target -> TargetInfo target
TargetInfo Bool
isBinary Maybe target
posClass Vector target
vals
  where
    distinct :: [target]
distinct = Set target -> [target]
forall a. Set a -> [a]
Set.toAscList ([target] -> Set target
forall a. Ord a => [a] -> Set a
Set.fromList (Vector target -> [target]
forall a. Vector a -> [a]
V.toList Vector target
vals))
    isBinary :: Bool
isBinary = [target] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [target]
distinct Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2
    posClass :: Maybe target
posClass = case [target]
distinct of
        (target
p : [target]
_) | Bool
isBinary -> target -> Maybe target
forall a. a -> Maybe a
Just target
p
        [target]
_ -> Maybe target
forall a. Maybe a
Nothing

{- | Distinct values, capped: @Right vs@ (sorted) under the cap, else @Left@
the count-so-far so the caller routes to the high-cardinality path.
-}
distinctValuesUpTo :: (Ord a) => Int -> V.Vector a -> Either Int [a]
distinctValuesUpTo :: forall a. Ord a => Int -> Vector a -> Either Int [a]
distinctValuesUpTo Int
cap Vector a
values = Set a -> Int -> Either Int [a]
go Set a
forall a. Set a
Set.empty Int
0
  where
    n :: Int
n = Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
values
    go :: Set a -> Int -> Either Int [a]
go !Set a
s !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = [a] -> Either Int [a]
forall a b. b -> Either a b
Right (Set a -> [a]
forall a. Set a -> [a]
Set.toAscList Set a
s)
        | Set a -> Int
forall a. Set a -> Int
Set.size Set a
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
cap = Int -> Either Int [a]
forall a b. a -> Either a b
Left (Set a -> Int
forall a. Set a -> Int
Set.size Set a
s)
        | Bool
otherwise = Set a -> Int -> Either Int [a]
go (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert (Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.unsafeIndex Vector a
values Int
i) Set a
s) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{- | OR-of-equalities for a value-list, shared by the expression and
truth-vector discrete paths so they stay byte-identical.
-}
orEqs :: (a -> Expr Bool) -> [a] -> Expr Bool
orEqs :: forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs a -> Expr Bool
eqLit = (Expr Bool -> Expr Bool -> Expr Bool) -> [Expr Bool] -> Expr Bool
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Expr Bool -> Expr Bool -> Expr Bool
(.||.) ([Expr Bool] -> Expr Bool)
-> ([a] -> [Expr Bool]) -> [a] -> Expr Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Expr Bool) -> [a] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map a -> Expr Bool
eqLit

subsetSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool]
subsetSplits :: forall a. (a -> Expr Bool) -> [a] -> [Expr Bool]
subsetSplits a -> Expr Bool
eqLit = ([a] -> Expr Bool) -> [[a]] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Expr Bool) -> [a] -> Expr Bool
forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs a -> Expr Bool
eqLit) ([[a]] -> [Expr Bool]) -> ([a] -> [[a]]) -> [a] -> [Expr Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. [a] -> [[a]]
subsetLists

-- | Proper non-empty, non-full subsets of the values.
subsetLists :: [a] -> [[a]]
subsetLists :: forall a. [a] -> [[a]]
subsetLists [a]
vs = Int -> [[a]] -> [[a]]
forall a. Int -> [a] -> [a]
drop Int
1 ([[a]] -> [[a]]
forall a. HasCallStack => [a] -> [a]
init ([a] -> [[a]]
forall a. [a] -> [[a]]
subsequences [a]
vs))

singletonSplits :: (a -> Expr Bool) -> [a] -> [Expr Bool]
singletonSplits :: forall a. (a -> Expr Bool) -> [a] -> [Expr Bool]
singletonSplits = (a -> Expr Bool) -> [a] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map

singletonLists :: [a] -> [[a]]
singletonLists :: forall a. [a] -> [[a]]
singletonLists = (a -> [a]) -> [a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [])

breimanPrefixSplits ::
    (Ord a, Ord target) =>
    target ->
    V.Vector a ->
    V.Vector target ->
    [a] ->
    (a -> Expr Bool) ->
    [Expr Bool]
breimanPrefixSplits :: forall a target.
(Ord a, Ord target) =>
target
-> Vector a
-> Vector target
-> [a]
-> (a -> Expr Bool)
-> [Expr Bool]
breimanPrefixSplits target
pc Vector a
values Vector target
targetVals [a]
distinctVals a -> Expr Bool
eqLit =
    ([a] -> Expr Bool) -> [[a]] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Expr Bool) -> [a] -> Expr Bool
forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs a -> Expr Bool
eqLit) (target -> Vector a -> Vector target -> [a] -> [[a]]
forall a target.
(Ord a, Ord target) =>
target -> Vector a -> Vector target -> [a] -> [[a]]
breimanPrefixLists target
pc Vector a
values Vector target
targetVals [a]
distinctVals)

{- | Breiman's binary-target split set: sort levels by Laplace-smoothed
positive rate, then take every contiguous non-trivial prefix.
-}
breimanPrefixLists ::
    (Ord a, Ord target) => target -> V.Vector a -> V.Vector target -> [a] -> [[a]]
breimanPrefixLists :: forall a target.
(Ord a, Ord target) =>
target -> Vector a -> Vector target -> [a] -> [[a]]
breimanPrefixLists target
pc Vector a
values Vector target
targetVals [a]
distinctVals =
    [a] -> [[a]]
forall a. [a] -> [[a]]
nonTrivialPrefixes (Map a (Int, Int) -> [a] -> [a]
forall a. Ord a => Map a (Int, Int) -> [a] -> [a]
sortByRate (target -> Vector a -> Vector target -> Map a (Int, Int)
forall a target.
(Ord a, Eq target) =>
target -> Vector a -> Vector target -> Map a (Int, Int)
levelCounts target
pc Vector a
values Vector target
targetVals) [a]
distinctVals)

levelCounts ::
    (Ord a, Eq target) =>
    target -> V.Vector a -> V.Vector target -> M.Map a (Int, Int)
levelCounts :: forall a target.
(Ord a, Eq target) =>
target -> Vector a -> Vector target -> Map a (Int, Int)
levelCounts target
pc Vector a
values Vector target
targetVals = (Map a (Int, Int) -> Int -> a -> Map a (Int, Int))
-> Map a (Int, Int) -> Vector a -> Map a (Int, Int)
forall a b. (a -> Int -> b -> a) -> a -> Vector b -> a
V.ifoldl' Map a (Int, Int) -> Int -> a -> Map a (Int, Int)
add Map a (Int, Int)
forall k a. Map k a
M.empty Vector a
values
  where
    add :: Map a (Int, Int) -> Int -> a -> Map a (Int, Int)
add Map a (Int, Int)
acc Int
i a
v = ((Int, Int) -> (Int, Int) -> (Int, Int))
-> a -> (Int, Int) -> Map a (Int, Int) -> Map a (Int, Int)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith (Int, Int) -> (Int, Int) -> (Int, Int)
forall {a} {b}. (Num a, Num b) => (a, b) -> (a, b) -> (a, b)
plus a
v (Bool -> Int
forall {a}. Num a => Bool -> a
indicator (Vector target -> Int -> target
forall a. Vector a -> Int -> a
V.unsafeIndex Vector target
targetVals Int
i target -> target -> Bool
forall a. Eq a => a -> a -> Bool
== target
pc), Int
1) Map a (Int, Int)
acc
    plus :: (a, b) -> (a, b) -> (a, b)
plus (a
p1, b
n1) (a
p2, b
n2) = (a
p1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
p2, b
n1 b -> b -> b
forall a. Num a => a -> a -> a
+ b
n2)
    indicator :: Bool -> a
indicator Bool
b = if Bool
b then a
1 else a
0

laplaceRate :: (Ord a) => M.Map a (Int, Int) -> a -> Double
laplaceRate :: forall a. Ord a => Map a (Int, Int) -> a -> Double
laplaceRate Map a (Int, Int)
counts a
v = case a -> Map a (Int, Int) -> Maybe (Int, Int)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
v Map a (Int, Int)
counts of
    Maybe (Int, Int)
Nothing -> Double
0.5
    Just (Int
pos, Int
n) -> (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
pos Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
1) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
2)

sortByRate :: (Ord a) => M.Map a (Int, Int) -> [a] -> [a]
sortByRate :: forall a. Ord a => Map a (Int, Int) -> [a] -> [a]
sortByRate Map a (Int, Int)
counts = (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((Double, a) -> (Double, a) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ((Double, a) -> (Double, a) -> Ordering)
-> (a -> (Double, a)) -> a -> a -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (\a
v -> (Map a (Int, Int) -> a -> Double
forall a. Ord a => Map a (Int, Int) -> a -> Double
laplaceRate Map a (Int, Int)
counts a
v, a
v)))

nonTrivialPrefixes :: [a] -> [[a]]
nonTrivialPrefixes :: forall a. [a] -> [[a]]
nonTrivialPrefixes = Int -> [[a]] -> [[a]]
forall a. Int -> [a] -> [a]
drop Int
1 ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. HasCallStack => [a] -> [a]
init ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [[a]]
forall a. [a] -> [[a]]
inits

{- | Value-lists a categorical column contributes; shared by the expression and
truth-vector paths so both enumerate identical candidates in the same order.
-}
catValueLists ::
    (Ord a, Ord target) =>
    Bool -> Maybe target -> V.Vector target -> Int -> V.Vector a -> [[a]]
catValueLists :: forall a target.
(Ord a, Ord target) =>
Bool -> Maybe target -> Vector target -> Int -> Vector a -> [[a]]
catValueLists Bool
isBinary Maybe target
posClass Vector target
targetVals Int
subsetCap Vector a
values
    | Vector a -> Bool
forall a. Vector a -> Bool
V.null Vector a
values = []
    | Bool
isBinary, Just target
pc <- Maybe target
posClass = target -> Vector target -> Vector a -> [[a]]
forall a target.
(Ord a, Ord target) =>
target -> Vector target -> Vector a -> [[a]]
binaryLists target
pc Vector target
targetVals Vector a
values
    | Bool
otherwise = Int -> Vector a -> [[a]]
forall a. Ord a => Int -> Vector a -> [[a]]
multiclassLists Int
subsetCap Vector a
values

binaryLists ::
    (Ord a, Ord target) => target -> V.Vector target -> V.Vector a -> [[a]]
binaryLists :: forall a target.
(Ord a, Ord target) =>
target -> Vector target -> Vector a -> [[a]]
binaryLists target
pc Vector target
targetVals Vector a
values
    | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
distinct Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = []
    | Bool
otherwise = target -> Vector a -> Vector target -> [a] -> [[a]]
forall a target.
(Ord a, Ord target) =>
target -> Vector a -> Vector target -> [a] -> [[a]]
breimanPrefixLists target
pc Vector a
values Vector target
targetVals [a]
distinct
  where
    distinct :: [a]
distinct = [a] -> Either Int [a] -> [a]
forall b a. b -> Either a b -> b
fromRight (Vector a -> [a]
forall a. Ord a => Vector a -> [a]
ascDistinct Vector a
values) (Int -> Vector a -> Either Int [a]
forall a. Ord a => Int -> Vector a -> Either Int [a]
distinctValuesUpTo Int
64 Vector a
values)

multiclassLists :: (Ord a) => Int -> V.Vector a -> [[a]]
multiclassLists :: forall a. Ord a => Int -> Vector a -> [[a]]
multiclassLists Int
subsetCap Vector a
values = case Int -> Vector a -> Either Int [a]
forall a. Ord a => Int -> Vector a -> Either Int [a]
distinctValuesUpTo Int
subsetCap Vector a
values of
    Right [a]
vs | [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
vs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 -> [a] -> [[a]]
forall a. [a] -> [[a]]
subsetLists [a]
vs
    Right [a]
_ -> []
    Left Int
_ -> [a] -> [[a]]
forall a. [a] -> [[a]]
singletonLists (Vector a -> [a]
forall a. Ord a => Vector a -> [a]
ascDistinct Vector a
values)

ascDistinct :: (Ord a) => V.Vector a -> [a]
ascDistinct :: forall a. Ord a => Vector a -> [a]
ascDistinct = Set a -> [a]
forall a. Set a -> [a]
Set.toAscList (Set a -> [a]) -> (Vector a -> Set a) -> Vector a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> (Vector a -> [a]) -> Vector a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [a]
forall a. Vector a -> [a]
V.toList

{- | Truth vector of @col ∈ values@ read directly from the column; equal to
interpreting @orEqs (== v) values@ because the values are distinct.
-}
membershipVec :: (Ord a) => V.Vector a -> [a] -> VU.Vector Bool
membershipVec :: forall a. Ord a => Vector a -> [a] -> Vector Bool
membershipVec Vector a
colVals [a]
vs =
    let !s :: Set a
s = [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList [a]
vs
     in Int -> (Int -> Bool) -> Vector Bool
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate (Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
colVals) (\Int
i -> a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Vector a
colVals Vector a -> Int -> a
forall a. Vector a -> Int -> a
`V.unsafeIndex` Int
i) Set a
s)

{- | Per-fit categorical generation context bundling the target summary and
the column-ordering registry.
-}
data CatCtx target = CatCtx
    { forall target. CatCtx target -> Bool
ccBinary :: !Bool
    , forall target. CatCtx target -> Maybe target
ccPos :: !(Maybe target)
    , forall target. CatCtx target -> Vector target
ccTargets :: !(V.Vector target)
    , forall target. CatCtx target -> Int
ccSubsetCap :: !Int
    , forall target. CatCtx target -> ColumnOrdering
ccOrds :: !ColumnOrdering
    }

catCtx :: TargetInfo target -> TreeConfig -> CatCtx target
catCtx :: forall target. TargetInfo target -> TreeConfig -> CatCtx target
catCtx TargetInfo target
ti TreeConfig
cfg =
    Bool
-> Maybe target
-> Vector target
-> Int
-> ColumnOrdering
-> CatCtx target
forall target.
Bool
-> Maybe target
-> Vector target
-> Int
-> ColumnOrdering
-> CatCtx target
CatCtx
        (TargetInfo target -> Bool
forall target. TargetInfo target -> Bool
tiIsBinary TargetInfo target
ti)
        (TargetInfo target -> Maybe target
forall target. TargetInfo target -> Maybe target
tiPositiveClass TargetInfo target
ti)
        (TargetInfo target -> Vector target
forall target. TargetInfo target -> Vector target
tiValues TargetInfo target
ti)
        (SynthConfig -> Int
maxCategoricalSubsetCardinality (TreeConfig -> SynthConfig
synthConfig TreeConfig
cfg))
        (TreeConfig -> ColumnOrdering
columnOrdering TreeConfig
cfg)

catValueListsFor :: (Ord a, Ord target) => CatCtx target -> V.Vector a -> [[a]]
catValueListsFor :: forall a target.
(Ord a, Ord target) =>
CatCtx target -> Vector a -> [[a]]
catValueListsFor CatCtx target
ctx = Bool -> Maybe target -> Vector target -> Int -> Vector a -> [[a]]
forall a target.
(Ord a, Ord target) =>
Bool -> Maybe target -> Vector target -> Int -> Vector a -> [[a]]
catValueLists (CatCtx target -> Bool
forall target. CatCtx target -> Bool
ccBinary CatCtx target
ctx) (CatCtx target -> Maybe target
forall target. CatCtx target -> Maybe target
ccPos CatCtx target
ctx) (CatCtx target -> Vector target
forall target. CatCtx target -> Vector target
ccTargets CatCtx target
ctx) (CatCtx target -> Int
forall target. CatCtx target -> Int
ccSubsetCap CatCtx target
ctx)

-- | True for numeric columns (handled by the numeric pool, not here).
isNumericKind :: forall a. (Columnable a) => Bool
isNumericKind :: forall a. Columnable a => Bool
isNumericKind = case forall a. SBoolI (FloatingTypes a) => SBool (FloatingTypes a)
sFloating @a of
    SBool (FloatingTypes a)
STrue -> Bool
True
    SBool (FloatingTypes a)
SFalse -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @a of
        SBool (IntegralTypes a)
STrue -> Bool
True
        SBool (IntegralTypes a)
SFalse -> Bool
False

{- | All equality-based candidate splits from non-numeric columns: per-column
categorical conditions plus cross-column equality/order conditions.
-}
discreteConditions ::
    forall target.
    (Columnable target, Ord target) =>
    TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool]
discreteConditions :: forall target.
(Columnable target, Ord target) =>
TargetInfo target -> TreeConfig -> DataFrame -> [Expr Bool]
discreteConditions TargetInfo target
targetInfo TreeConfig
cfg DataFrame
df =
    (Text -> [Expr Bool]) -> [Text] -> [Expr Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (CatCtx target -> DataFrame -> Text -> [Expr Bool]
forall target.
(Columnable target, Ord target) =>
CatCtx target -> DataFrame -> Text -> [Expr Bool]
columnConds (TargetInfo target -> TreeConfig -> CatCtx target
forall target. TargetInfo target -> TreeConfig -> CatCtx target
catCtx TargetInfo target
targetInfo TreeConfig
cfg) DataFrame
df) (DataFrame -> [Text]
columnNames DataFrame
df)
        [Expr Bool] -> [Expr Bool] -> [Expr Bool]
forall a. [a] -> [a] -> [a]
++ TreeConfig -> DataFrame -> [Expr Bool]
crossColumnConds TreeConfig
cfg DataFrame
df

columnConds ::
    (Columnable target, Ord target) =>
    CatCtx target -> DataFrame -> T.Text -> [Expr Bool]
columnConds :: forall target.
(Columnable target, Ord target) =>
CatCtx target -> DataFrame -> Text -> [Expr Bool]
columnConds CatCtx target
ctx DataFrame
df Text
colName = case Text -> DataFrame -> Column
unsafeGetColumn Text
colName DataFrame
df of
    BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [Expr Bool]
nonNullColConds CatCtx target
ctx Text
colName Vector a
column
    BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column
    UnboxedColumn Maybe Bitmap
_ (Vector a
_ :: VU.Vector a) -> []
    pt :: Column
pt@(PackedText Maybe Bitmap
_ PackedTextData
_) -> case Column -> Column
materializePacked Column
pt of
        BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [Expr Bool]
nonNullColConds CatCtx target
ctx Text
colName Vector a
column
        BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column
        Column
_ -> []
    mc :: Column
mc@(MergedColumn Column
_ Column
_) -> case Column -> Column
materializeMerged Column
mc of
        BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [Expr Bool]
nonNullColConds CatCtx target
ctx Text
colName Vector a
column
        BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column
        Column
_ -> []

nonNullColConds ::
    forall a target.
    (Columnable a, Ord target) =>
    CatCtx target -> T.Text -> V.Vector a -> [Expr Bool]
nonNullColConds :: forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [Expr Bool]
nonNullColConds CatCtx target
ctx Text
colName Vector a
column =
    [Expr Bool] -> Maybe [Expr Bool] -> [Expr Bool]
forall a. a -> Maybe a -> a
fromMaybe
        []
        ( forall a r.
Columnable a =>
ColumnOrdering -> (Ord a => r) -> Maybe r
withOrdFrom @a
            (CatCtx target -> ColumnOrdering
forall target. CatCtx target -> ColumnOrdering
ccOrds CatCtx target
ctx)
            (([a] -> Expr Bool) -> [[a]] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Expr Bool) -> [a] -> Expr Bool
forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs (forall a. Columnable a => Text -> a -> Expr Bool
eqExprFor @a Text
colName)) (CatCtx target -> Vector a -> [[a]]
forall a target.
(Ord a, Ord target) =>
CatCtx target -> Vector a -> [[a]]
catValueListsFor CatCtx target
ctx Vector a
column))
        )

nullableColConds ::
    forall a target.
    (Columnable a, Ord target) =>
    CatCtx target -> T.Text -> Bitmap -> V.Vector a -> [Expr Bool]
nullableColConds :: forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column
    | forall a. Columnable a => Bool
isNumericKind @a Bool -> Bool -> Bool
|| Vector a -> Bool
forall a. Vector a -> Bool
V.null Vector a
valid = []
    | Bool
otherwise =
        [Expr Bool] -> Maybe [Expr Bool] -> [Expr Bool]
forall a. a -> Maybe a -> a
fromMaybe
            []
            ( forall a r.
Columnable a =>
ColumnOrdering -> (Ord a => r) -> Maybe r
withOrdFrom @a
                (CatCtx target -> ColumnOrdering
forall target. CatCtx target -> ColumnOrdering
ccOrds CatCtx target
ctx)
                (([a] -> Expr Bool) -> [[a]] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Expr Bool) -> [a] -> Expr Bool
forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs (forall a. Columnable a => Text -> a -> Expr Bool
eqJustFor @a Text
colName)) (CatCtx target -> Vector a -> [[a]]
forall a target.
(Ord a, Ord target) =>
CatCtx target -> Vector a -> [[a]]
catValueListsFor CatCtx target
ctx Vector a
valid))
            )
  where
    valid :: Vector a
valid = Bitmap -> Vector a -> Vector a
forall a. Bitmap -> Vector a -> Vector a
validBoxedValues Bitmap
bm Vector a
column

eqExprFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool
eqExprFor :: forall a. Columnable a => Text -> a -> Expr Bool
eqExprFor Text
colName a
v = forall a. Columnable a => Text -> Expr a
Col @a Text
colName Expr a -> Expr a -> Expr Bool
forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.==. a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
v

eqJustFor :: forall a. (Columnable a) => T.Text -> a -> Expr Bool
eqJustFor :: forall a. Columnable a => Text -> a -> Expr Bool
eqJustFor Text
colName a
v = forall a. Columnable a => Text -> Expr a
Col @(Maybe a) Text
colName Expr (Maybe a) -> Expr (Maybe a) -> Expr Bool
forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.==. Maybe a -> Expr (Maybe a)
forall a. Columnable a => a -> Expr a
Lit (a -> Maybe a
forall a. a -> Maybe a
Just a
v)

-- | Cross-column equality/order conditions over pairs of same-typed columns.
crossColumnConds :: TreeConfig -> DataFrame -> [Expr Bool]
crossColumnConds :: TreeConfig -> DataFrame -> [Expr Bool]
crossColumnConds TreeConfig
cfg DataFrame
df = ((Text, Text) -> [Expr Bool]) -> [(Text, Text)] -> [Expr Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (ColumnOrdering -> DataFrame -> (Text, Text) -> [Expr Bool]
pairConds (TreeConfig -> ColumnOrdering
columnOrdering TreeConfig
cfg) DataFrame
df) (TreeConfig -> DataFrame -> [(Text, Text)]
allowedPairs TreeConfig
cfg DataFrame
df)

allowedPairs :: TreeConfig -> DataFrame -> [(T.Text, T.Text)]
allowedPairs :: TreeConfig -> DataFrame -> [(Text, Text)]
allowedPairs TreeConfig
cfg DataFrame
df =
    [ (Text
l, Text
r)
    | Text
l <- DataFrame -> [Text]
columnNames DataFrame
df
    , Text
r <- DataFrame -> [Text]
columnNames DataFrame
df
    , Text
l Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
r
    , Bool -> Bool
not (TreeConfig -> Text -> Text -> Bool
isDisallowedPair TreeConfig
cfg Text
l Text
r)
    ]

isDisallowedPair :: TreeConfig -> T.Text -> T.Text -> Bool
isDisallowedPair :: TreeConfig -> Text -> Text -> Bool
isDisallowedPair TreeConfig
cfg Text
l Text
r =
    ((Text, Text) -> Bool) -> [(Text, Text)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
        (\(Text
l', Text
r') -> [Text] -> [Text]
forall a. Ord a => [a] -> [a]
sort [Text
l', Text
r'] [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text] -> [Text]
forall a. Ord a => [a] -> [a]
sort [Text
l, Text
r])
        (SynthConfig -> [(Text, Text)]
disallowedCombinations (TreeConfig -> SynthConfig
synthConfig TreeConfig
cfg))

pairConds :: ColumnOrdering -> DataFrame -> (T.Text, T.Text) -> [Expr Bool]
pairConds :: ColumnOrdering -> DataFrame -> (Text, Text) -> [Expr Bool]
pairConds ColumnOrdering
ords DataFrame
df (Text
l, Text
r) = case ( Column -> Column
materializePacked (Text -> DataFrame -> Column
unsafeGetColumn Text
l DataFrame
df)
                                , Column -> Column
materializePacked (Text -> DataFrame -> Column
unsafeGetColumn Text
r DataFrame
df)
                                ) of
    (BoxedColumn Maybe Bitmap
Nothing (Vector a
_ :: V.Vector a), BoxedColumn Maybe Bitmap
Nothing (Vector a
_ :: V.Vector b)) -> forall a b.
(Columnable a, Columnable b) =>
Text -> Text -> [Expr Bool]
strictPairConds @a @b Text
l Text
r
    (BoxedColumn (Just Bitmap
_) (Vector a
_ :: V.Vector a), BoxedColumn (Just Bitmap
_) (Vector a
_ :: V.Vector b)) -> forall a b.
(Columnable a, Columnable b) =>
ColumnOrdering -> Text -> Text -> [Expr Bool]
nullablePairConds @a @b ColumnOrdering
ords Text
l Text
r
    (Column, Column)
_ -> []

strictPairConds ::
    forall a b. (Columnable a, Columnable b) => T.Text -> T.Text -> [Expr Bool]
strictPairConds :: forall a b.
(Columnable a, Columnable b) =>
Text -> Text -> [Expr Bool]
strictPairConds Text
l Text
r = case TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
    Just a :~: b
Refl -> [forall a. Columnable a => Text -> Expr a
Col @a Text
l Expr a -> Expr a -> Expr Bool
forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.==. forall a. Columnable a => Text -> Expr a
Col @a Text
r]
    Maybe (a :~: b)
Nothing -> []

nullablePairConds ::
    forall a b.
    (Columnable a, Columnable b) =>
    ColumnOrdering -> T.Text -> T.Text -> [Expr Bool]
nullablePairConds :: forall a b.
(Columnable a, Columnable b) =>
ColumnOrdering -> Text -> Text -> [Expr Bool]
nullablePairConds ColumnOrdering
ords Text
l Text
r = case TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b) of
    Maybe (a :~: b)
Nothing -> []
    Just a :~: b
Refl -> forall a.
Columnable a =>
ColumnOrdering -> Text -> Text -> [Expr Bool]
nullableEqOrLe @a ColumnOrdering
ords Text
l Text
r

nullableEqOrLe ::
    forall a. (Columnable a) => ColumnOrdering -> T.Text -> T.Text -> [Expr Bool]
nullableEqOrLe :: forall a.
Columnable a =>
ColumnOrdering -> Text -> Text -> [Expr Bool]
nullableEqOrLe ColumnOrdering
ords Text
l Text
r
    | forall a. Columnable a => Bool
isTextType @a = [Expr Bool]
eqOnly
    | Bool
otherwise =
        [Expr Bool]
-> ([Expr Bool] -> [Expr Bool]) -> Maybe [Expr Bool] -> [Expr Bool]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
            [Expr Bool]
eqOnly
            ([Expr Bool] -> [Expr Bool] -> [Expr Bool]
forall a. [a] -> [a] -> [a]
++ [Expr Bool]
eqOnly)
            (forall a r.
Columnable a =>
ColumnOrdering -> (Ord a => r) -> Maybe r
withOrdFrom @a ColumnOrdering
ords [forall a. Columnable a => Text -> Expr a
Col @(Maybe a) Text
l Expr (Maybe a) -> Expr (Maybe a) -> Expr Bool
forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
.<=. forall a. Columnable a => Text -> Expr a
Col @(Maybe a) Text
r])
  where
    eqOnly :: [Expr Bool]
eqOnly = [forall a. Columnable a => Text -> Expr a
Col @(Maybe a) Text
l Expr (Maybe a) -> Expr (Maybe a) -> Expr Bool
forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.==. forall a. Columnable a => Text -> Expr a
Col @(Maybe a) Text
r]

isTextType :: forall a. (Columnable a) => Bool
isTextType :: forall a. Columnable a => Bool
isTextType = case TypeRep a -> TypeRep Text -> Maybe (a :~: Text)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @T.Text) of
    Just a :~: Text
Refl -> Bool
True
    Maybe (a :~: Text)
Nothing -> Bool
False

{- | 'discreteConditions' materialized with shared per-column reads: the
non-nullable categorical path builds truth vectors directly from one read
per column; nullable and cross-column fall back to interpret.
-}
discreteCondVecs ::
    forall target.
    (Columnable target, Ord target) =>
    TargetInfo target -> TreeConfig -> DataFrame -> [CondVec]
discreteCondVecs :: forall target.
(Columnable target, Ord target) =>
TargetInfo target -> TreeConfig -> DataFrame -> [CondVec]
discreteCondVecs TargetInfo target
targetInfo TreeConfig
cfg DataFrame
df =
    (Text -> [CondVec]) -> [Text] -> [CondVec]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (CatCtx target -> DataFrame -> Text -> [CondVec]
forall target.
(Columnable target, Ord target) =>
CatCtx target -> DataFrame -> Text -> [CondVec]
columnCondVecs (TargetInfo target -> TreeConfig -> CatCtx target
forall target. TargetInfo target -> TreeConfig -> CatCtx target
catCtx TargetInfo target
targetInfo TreeConfig
cfg) DataFrame
df) (DataFrame -> [Text]
columnNames DataFrame
df)
        [CondVec] -> [CondVec] -> [CondVec]
forall a. [a] -> [a] -> [a]
++ (Expr Bool -> Maybe CondVec) -> [Expr Bool] -> [CondVec]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DataFrame -> Expr Bool -> Maybe CondVec
materializeCondVec DataFrame
df) (TreeConfig -> DataFrame -> [Expr Bool]
crossColumnConds TreeConfig
cfg DataFrame
df)

columnCondVecs ::
    (Columnable target, Ord target) =>
    CatCtx target -> DataFrame -> T.Text -> [CondVec]
columnCondVecs :: forall target.
(Columnable target, Ord target) =>
CatCtx target -> DataFrame -> Text -> [CondVec]
columnCondVecs CatCtx target
ctx DataFrame
df Text
colName = case Text -> DataFrame -> Column
unsafeGetColumn Text
colName DataFrame
df of
    BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [CondVec]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [CondVec]
nonNullColCondVecs CatCtx target
ctx Text
colName Vector a
column
    BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> (Expr Bool -> Maybe CondVec) -> [Expr Bool] -> [CondVec]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DataFrame -> Expr Bool -> Maybe CondVec
materializeCondVec DataFrame
df) (CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column)
    UnboxedColumn Maybe Bitmap
_ (Vector a
_ :: VU.Vector a) -> []
    pt :: Column
pt@(PackedText Maybe Bitmap
_ PackedTextData
_) -> case Column -> Column
materializePacked Column
pt of
        BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [CondVec]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [CondVec]
nonNullColCondVecs CatCtx target
ctx Text
colName Vector a
column
        BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> (Expr Bool -> Maybe CondVec) -> [Expr Bool] -> [CondVec]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DataFrame -> Expr Bool -> Maybe CondVec
materializeCondVec DataFrame
df) (CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column)
        Column
_ -> []
    mc :: Column
mc@(MergedColumn Column
_ Column
_) -> case Column -> Column
materializeMerged Column
mc of
        BoxedColumn Maybe Bitmap
Nothing (Vector a
column :: V.Vector a) -> CatCtx target -> Text -> Vector a -> [CondVec]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [CondVec]
nonNullColCondVecs CatCtx target
ctx Text
colName Vector a
column
        BoxedColumn (Just Bitmap
bm) (Vector a
column :: V.Vector a) -> (Expr Bool -> Maybe CondVec) -> [Expr Bool] -> [CondVec]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (DataFrame -> Expr Bool -> Maybe CondVec
materializeCondVec DataFrame
df) (CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Bitmap -> Vector a -> [Expr Bool]
nullableColConds CatCtx target
ctx Text
colName Bitmap
bm Vector a
column)
        Column
_ -> []

nonNullColCondVecs ::
    forall a target.
    (Columnable a, Ord target) => CatCtx target -> T.Text -> V.Vector a -> [CondVec]
nonNullColCondVecs :: forall a target.
(Columnable a, Ord target) =>
CatCtx target -> Text -> Vector a -> [CondVec]
nonNullColCondVecs CatCtx target
ctx Text
colName Vector a
column =
    [CondVec] -> Maybe [CondVec] -> [CondVec]
forall a. a -> Maybe a -> a
fromMaybe
        []
        ( forall a r.
Columnable a =>
ColumnOrdering -> (Ord a => r) -> Maybe r
withOrdFrom @a
            (CatCtx target -> ColumnOrdering
forall target. CatCtx target -> ColumnOrdering
ccOrds CatCtx target
ctx)
            (([a] -> CondVec) -> [[a]] -> [CondVec]
forall a b. (a -> b) -> [a] -> [b]
map (Text -> Vector a -> [a] -> CondVec
forall a.
(Columnable a, Ord a) =>
Text -> Vector a -> [a] -> CondVec
membershipCondVec Text
colName Vector a
column) (CatCtx target -> Vector a -> [[a]]
forall a target.
(Ord a, Ord target) =>
CatCtx target -> Vector a -> [[a]]
catValueListsFor CatCtx target
ctx Vector a
column))
        )

membershipCondVec ::
    forall a. (Columnable a, Ord a) => T.Text -> V.Vector a -> [a] -> CondVec
membershipCondVec :: forall a.
(Columnable a, Ord a) =>
Text -> Vector a -> [a] -> CondVec
membershipCondVec Text
colName Vector a
column [a]
vs = Expr Bool -> Vector Bool -> CondVec
CondVec ((a -> Expr Bool) -> [a] -> Expr Bool
forall a. (a -> Expr Bool) -> [a] -> Expr Bool
orEqs (forall a. Columnable a => Text -> a -> Expr Bool
eqExprFor @a Text
colName) [a]
vs) (Vector a -> [a] -> Vector Bool
forall a. Ord a => Vector a -> [a] -> Vector Bool
membershipVec Vector a
column [a]
vs)