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

{- | Numeric split candidates: per-column Double expressions, arithmetic
expansion, and threshold conditions. Nullable columns also contribute an
@isNothing@ candidate, so missingness is splittable directly rather than
only via the null-routing of value splits. 'numericCondVecs' materializes
the pool with one interpret per distinct expression.
-}
module DataFrame.DecisionTree.Numeric (
    NumExpr (..),
    numExprCols,
    numExprEq,
    combineNumExprs,
    numericConditions,
    generateNumericConds,
    missingnessConditions,
    percentilesOf,
    numericCondVecs,
    numericExprsWithTerms,
    numericCols,
) where

import DataFrame.DecisionTree.CondVec (CondVec (..), materializeCondVec)
import DataFrame.DecisionTree.Types (SynthConfig (..), TreeConfig (..))
import qualified DataFrame.Functions as F
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (DataFrame, columnNames, unsafeGetColumn)
import DataFrame.Internal.Expression (Expr (..), eqExpr, getColumns, normalize)
import DataFrame.Internal.Interpreter (interpret)
import DataFrame.Internal.Types
import DataFrame.Operators

import Data.List (sort)
import Data.Maybe (catMaybes, 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)

-- | A numeric feature expression, non-nullable or nullable.
data NumExpr
    = NDouble !(Expr Double)
    | NMaybeDouble !(Expr (Maybe Double))

numExprCols :: NumExpr -> [T.Text]
numExprCols :: NumExpr -> [Text]
numExprCols (NDouble Expr Double
e) = Expr Double -> [Text]
forall a. Expr a -> [Text]
getColumns Expr Double
e
numExprCols (NMaybeDouble Expr (Maybe Double)
e) = Expr (Maybe Double) -> [Text]
forall a. Expr a -> [Text]
getColumns Expr (Maybe Double)
e

numExprEq :: NumExpr -> NumExpr -> Bool
numExprEq :: NumExpr -> NumExpr -> Bool
numExprEq (NDouble Expr Double
e1) (NDouble Expr Double
e2) = Expr Double -> Expr Double -> Bool
forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr Double
e1 Expr Double
e2
numExprEq (NMaybeDouble Expr (Maybe Double)
e1) (NMaybeDouble Expr (Maybe Double)
e2) = Expr (Maybe Double) -> Expr (Maybe Double) -> Bool
forall a. Columnable a => Expr a -> Expr a -> Bool
eqExpr Expr (Maybe Double)
e1 Expr (Maybe Double)
e2
numExprEq NumExpr
_ NumExpr
_ = Bool
False

-- | Safe division: @0@ (or @Nothing@) where the divisor is zero.
safeDivD :: Expr Double -> Expr Double -> Expr Double
safeDivD :: Expr Double -> Expr Double -> Expr Double
safeDivD Expr Double
a Expr Double
b = Expr Bool -> Expr Double -> Expr Double -> Expr Double
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
F.ifThenElse (Expr Double
b Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Eq (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
./= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
0 :: Double)) (Expr Double
a Expr Double -> Expr Double -> Expr (WidenResultDiv Double Double)
forall a b.
(DivWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b),
 Fractional (PromoteDiv (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResultDiv a b)
./ Expr Double
b) (Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
0 :: Double))

safeDivMaybe :: Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double)
safeDivMaybe :: Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double)
safeDivMaybe Expr Bool
nonZero Expr (Maybe Double)
q = Expr Bool
-> Expr (Maybe Double)
-> Expr (Maybe Double)
-> Expr (Maybe Double)
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
F.ifThenElse Expr Bool
nonZero Expr (Maybe Double)
q (Maybe Double -> Expr (Maybe Double)
forall a. Columnable a => a -> Expr a
F.lit (Maybe Double
forall a. Maybe a
Nothing :: Maybe Double))

-- | Arithmetic combinations (@+@, @-@, @*@, safe @/@) of two numeric exprs.
combineNumExprs :: NumExpr -> NumExpr -> [NumExpr]
combineNumExprs :: NumExpr -> NumExpr -> [NumExpr]
combineNumExprs (NDouble Expr Double
e1) (NDouble Expr Double
e2) =
    (Expr Double -> NumExpr) -> [Expr Double] -> [NumExpr]
forall a b. (a -> b) -> [a] -> [b]
map Expr Double -> NumExpr
NDouble [Expr Double
e1 Expr Double -> Expr Double -> Expr (WidenResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.+ Expr Double
e2, Expr Double
e1 Expr Double -> Expr Double -> Expr (WidenResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.- Expr Double
e2, Expr Double
e1 Expr Double -> Expr Double -> Expr (WidenResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.* Expr Double
e2, Expr Double -> Expr Double -> Expr Double
safeDivD Expr Double
e1 Expr Double
e2]
combineNumExprs (NDouble Expr Double
e1) (NMaybeDouble Expr (Maybe Double)
e2) =
    (Expr (Maybe Double) -> NumExpr)
-> [Expr (Maybe Double)] -> [NumExpr]
forall a b. (a -> b) -> [a] -> [b]
map
        Expr (Maybe Double) -> NumExpr
NMaybeDouble
        [ Expr Double
e1 Expr Double
-> Expr (Maybe Double) -> Expr (WidenResult Double (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.+ Expr (Maybe Double)
e2
        , Expr Double
e1 Expr Double
-> Expr (Maybe Double) -> Expr (WidenResult Double (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.- Expr (Maybe Double)
e2
        , Expr Double
e1 Expr Double
-> Expr (Maybe Double) -> Expr (WidenResult Double (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.* Expr (Maybe Double)
e2
        , Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double)
safeDivMaybe (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e2 Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Eq (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
./= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
0 :: Double))) (Expr Double
e1 Expr Double
-> Expr (Maybe Double)
-> Expr (WidenResultDiv Double (Maybe Double))
forall a b.
(DivWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b),
 Fractional (PromoteDiv (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResultDiv a b)
./ Expr (Maybe Double)
e2)
        ]
combineNumExprs (NMaybeDouble Expr (Maybe Double)
e1) (NDouble Expr Double
e2) =
    (Expr (Maybe Double) -> NumExpr)
-> [Expr (Maybe Double)] -> [NumExpr]
forall a b. (a -> b) -> [a] -> [b]
map
        Expr (Maybe Double) -> NumExpr
NMaybeDouble
        [ Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr Double -> Expr (WidenResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.+ Expr Double
e2
        , Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr Double -> Expr (WidenResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.- Expr Double
e2
        , Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr Double -> Expr (WidenResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.* Expr Double
e2
        , Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double)
safeDivMaybe (Expr Double
e2 Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Eq (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
./= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
0 :: Double)) (Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr Double -> Expr (WidenResultDiv (Maybe Double) Double)
forall a b.
(DivWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b),
 Fractional (PromoteDiv (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResultDiv a b)
./ Expr Double
e2)
        ]
combineNumExprs (NMaybeDouble Expr (Maybe Double)
e1) (NMaybeDouble Expr (Maybe Double)
e2) =
    (Expr (Maybe Double) -> NumExpr)
-> [Expr (Maybe Double)] -> [NumExpr]
forall a b. (a -> b) -> [a] -> [b]
map
        Expr (Maybe Double) -> NumExpr
NMaybeDouble
        [ Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr (Maybe Double)
-> Expr (WidenResult (Maybe Double) (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.+ Expr (Maybe Double)
e2
        , Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr (Maybe Double)
-> Expr (WidenResult (Maybe Double) (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.- Expr (Maybe Double)
e2
        , Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr (Maybe Double)
-> Expr (WidenResult (Maybe Double) (Maybe Double))
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (Promote (BaseType a) (BaseType b)) (WidenResult a b),
 Num (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResult a b)
.* Expr (Maybe Double)
e2
        , Expr Bool -> Expr (Maybe Double) -> Expr (Maybe Double)
safeDivMaybe (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e2 Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Eq (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
./= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
0 :: Double))) (Expr (Maybe Double)
e1 Expr (Maybe Double)
-> Expr (Maybe Double)
-> Expr (WidenResultDiv (Maybe Double) (Maybe Double))
forall a b.
(DivWidenOp (BaseType a) (BaseType b),
 NullLift2Op
   a b (PromoteDiv (BaseType a) (BaseType b)) (WidenResultDiv a b),
 Fractional (PromoteDiv (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (WidenResultDiv a b)
./ Expr (Maybe Double)
e2)
        ]

numericConditions :: TreeConfig -> DataFrame -> [Expr Bool]
numericConditions :: TreeConfig -> DataFrame -> [Expr Bool]
numericConditions = TreeConfig -> DataFrame -> [Expr Bool]
generateNumericConds

generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool]
generateNumericConds :: TreeConfig -> DataFrame -> [Expr Bool]
generateNumericConds TreeConfig
cfg DataFrame
df = [Expr Bool]
thresholdConds [Expr Bool] -> [Expr Bool] -> [Expr Bool]
forall a. [a] -> [a] -> [a]
++ DataFrame -> [Expr Bool]
missingnessConditions DataFrame
df
  where
    thresholdConds :: [Expr Bool]
thresholdConds = do
        NumExpr
expr <- SynthConfig -> DataFrame -> [NumExpr]
numericExprsWithTerms (TreeConfig -> SynthConfig
synthConfig TreeConfig
cfg) DataFrame
df
        Double
threshold <- TreeConfig -> DataFrame -> NumExpr -> [Double]
numericThresholds TreeConfig
cfg DataFrame
df NumExpr
expr
        NumExpr -> Double -> [Expr Bool]
condsFromExpr NumExpr
expr Double
threshold

missingnessConditions :: DataFrame -> [Expr Bool]
missingnessConditions :: DataFrame -> [Expr Bool]
missingnessConditions DataFrame
df = (Text -> [Expr Bool]) -> [Text] -> [Expr Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Text -> [Expr Bool]
missingCond (DataFrame -> [Text]
columnNames DataFrame
df)
  where
    missingCond :: Text -> [Expr Bool]
missingCond Text
name = case Text -> DataFrame -> Column
unsafeGetColumn Text
name DataFrame
df of
        BoxedColumn (Just Bitmap
_) (Vector a
_ :: V.Vector b) -> [Expr (Maybe a) -> Expr Bool
forall a. Columnable a => Expr (Maybe a) -> Expr Bool
F.isNothing (forall a. Columnable a => Text -> Expr a
Col @(Maybe b) Text
name)]
        UnboxedColumn (Just Bitmap
_) (Vector a
_ :: VU.Vector b) -> [Expr (Maybe a) -> Expr Bool
forall a. Columnable a => Expr (Maybe a) -> Expr Bool
F.isNothing (forall a. Columnable a => Text -> Expr a
Col @(Maybe b) Text
name)]
        PackedText (Just Bitmap
_) PackedTextData
_ -> [Expr (Maybe Text) -> Expr Bool
forall a. Columnable a => Expr (Maybe a) -> Expr Bool
F.isNothing (forall a. Columnable a => Text -> Expr a
Col @(Maybe T.Text) Text
name)]
        Column
_ -> []

-- | Thresholds for nullable expressions come from the observed values only.
numericThresholds :: TreeConfig -> DataFrame -> NumExpr -> [Double]
numericThresholds :: TreeConfig -> DataFrame -> NumExpr -> [Double]
numericThresholds TreeConfig
cfg DataFrame
df (NDouble Expr Double
e) = TreeConfig -> DataFrame -> Expr Double -> [Double]
thresholdsForExpr TreeConfig
cfg DataFrame
df Expr Double
e
numericThresholds TreeConfig
cfg DataFrame
df (NMaybeDouble Expr (Maybe Double)
e) =
    [Double]
-> (Vector (Maybe Double) -> [Double])
-> Maybe (Vector (Maybe Double))
-> [Double]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
        []
        ([Int] -> [Double] -> [Double]
percentilesOf (TreeConfig -> [Int]
percentiles TreeConfig
cfg) ([Double] -> [Double])
-> (Vector (Maybe Double) -> [Double])
-> Vector (Maybe Double)
-> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe Double] -> [Double]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Double] -> [Double])
-> (Vector (Maybe Double) -> [Maybe Double])
-> Vector (Maybe Double)
-> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Maybe Double) -> [Maybe Double]
forall a. Vector a -> [a]
V.toList)
        (DataFrame -> Expr (Maybe Double) -> Maybe (Vector (Maybe Double))
interpretMaybeDoubleCol DataFrame
df Expr (Maybe Double)
e)

thresholdsForExpr :: TreeConfig -> DataFrame -> Expr Double -> [Double]
thresholdsForExpr :: TreeConfig -> DataFrame -> Expr Double -> [Double]
thresholdsForExpr TreeConfig
cfg DataFrame
df Expr Double
e =
    [Double]
-> (Vector Double -> [Double]) -> Maybe (Vector Double) -> [Double]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ([Int] -> [Double] -> [Double]
percentilesOf (TreeConfig -> [Int]
percentiles TreeConfig
cfg) ([Double] -> [Double])
-> (Vector Double -> [Double]) -> Vector Double -> [Double]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Double -> [Double]
forall a. Vector a -> [a]
V.toList) (DataFrame -> Expr Double -> Maybe (Vector Double)
interpretDoubleCol DataFrame
df Expr Double
e)

condsFromExpr :: NumExpr -> Double -> [Expr Bool]
condsFromExpr :: NumExpr -> Double -> [Expr Bool]
condsFromExpr (NDouble Expr Double
e) Double
t = [Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.<= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.>= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.< Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.> Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t]
condsFromExpr (NMaybeDouble Expr (Maybe Double)
e) Double
t =
    (Expr (Maybe Bool) -> Expr Bool)
-> [Expr (Maybe Bool)] -> [Expr Bool]
forall a b. (a -> b) -> [a] -> [b]
map
        (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False)
        [Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.<= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.>= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.< Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t, Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.> Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t]

{- | Percentile thresholds for a value list: sort once, index each percentile.
Shared by 'generateNumericConds' and 'numericCondVecs' for identical results.
-}
percentilesOf :: [Int] -> [Double] -> [Double]
percentilesOf :: [Int] -> [Double] -> [Double]
percentilesOf [Int]
ps [Double]
valsList
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = []
    | Bool
otherwise = (Int -> Double) -> [Int] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
p -> Vector Double
sortedV Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
p Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
100))) [Int]
ps
  where
    !sortedV :: Vector Double
sortedV = [Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList ([Double] -> [Double]
forall a. Ord a => [a] -> [a]
sort [Double]
valsList)
    !n :: Int
n = Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
sortedV

interpretDoubleCol :: DataFrame -> Expr Double -> Maybe (V.Vector Double)
interpretDoubleCol :: DataFrame -> Expr Double -> Maybe (Vector Double)
interpretDoubleCol DataFrame
df Expr Double
e = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @Double DataFrame
df Expr Double
e of
    Right (TColumn Column
column) -> (DataFrameException -> Maybe (Vector Double))
-> (Vector Double -> Maybe (Vector Double))
-> Either DataFrameException (Vector Double)
-> Maybe (Vector Double)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (Vector Double)
-> DataFrameException -> Maybe (Vector Double)
forall a b. a -> b -> a
const Maybe (Vector Double)
forall a. Maybe a
Nothing) Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just (forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @Double Column
column)
    Either DataFrameException (TypedColumn Double)
_ -> Maybe (Vector Double)
forall a. Maybe a
Nothing

interpretMaybeDoubleCol ::
    DataFrame -> Expr (Maybe Double) -> Maybe (V.Vector (Maybe Double))
interpretMaybeDoubleCol :: DataFrame -> Expr (Maybe Double) -> Maybe (Vector (Maybe Double))
interpretMaybeDoubleCol DataFrame
df Expr (Maybe Double)
e = case forall a.
Columnable a =>
DataFrame -> Expr a -> Either DataFrameException (TypedColumn a)
interpret @(Maybe Double) DataFrame
df Expr (Maybe Double)
e of
    Right (TColumn Column
column) -> (DataFrameException -> Maybe (Vector (Maybe Double)))
-> (Vector (Maybe Double) -> Maybe (Vector (Maybe Double)))
-> Either DataFrameException (Vector (Maybe Double))
-> Maybe (Vector (Maybe Double))
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe (Vector (Maybe Double))
-> DataFrameException -> Maybe (Vector (Maybe Double))
forall a b. a -> b -> a
const Maybe (Vector (Maybe Double))
forall a. Maybe a
Nothing) Vector (Maybe Double) -> Maybe (Vector (Maybe Double))
forall a. a -> Maybe a
Just (forall a (v :: * -> *).
(Vector v a, Columnable a) =>
Column -> Either DataFrameException (v a)
toVector @(Maybe Double) Column
column)
    Either DataFrameException (TypedColumn (Maybe Double))
_ -> Maybe (Vector (Maybe Double))
forall a. Maybe a
Nothing

{- | Materialize the numeric pool with one interpret per distinct expression,
deriving each threshold/operator truth vector by direct comparison.
Byte-identical to materializing 'numericConditions' one at a time.
-}
numericCondVecs :: TreeConfig -> DataFrame -> DataFrame -> [CondVec]
numericCondVecs :: TreeConfig -> DataFrame -> DataFrame -> [CondVec]
numericCondVecs TreeConfig
cfg DataFrame
dfGen DataFrame
df =
    (NumExpr -> [CondVec]) -> [NumExpr] -> [CondVec]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap NumExpr -> [CondVec]
forExpr (SynthConfig -> DataFrame -> [NumExpr]
numericExprsWithTerms (TreeConfig -> SynthConfig
synthConfig TreeConfig
cfg) DataFrame
dfGen)
        [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) (DataFrame -> [Expr Bool]
missingnessConditions DataFrame
dfGen)
  where
    forExpr :: NumExpr -> [CondVec]
forExpr (NDouble Expr Double
e) = [CondVec]
-> (Vector Double -> [CondVec])
-> Maybe (Vector Double)
-> [CondVec]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (TreeConfig -> Expr Double -> Vector Double -> [CondVec]
condsForDouble TreeConfig
cfg Expr Double
e) (DataFrame -> Expr Double -> Maybe (Vector Double)
interpretDoubleCol DataFrame
df Expr Double
e)
    forExpr (NMaybeDouble Expr (Maybe Double)
e) = [CondVec]
-> (Vector (Maybe Double) -> [CondVec])
-> Maybe (Vector (Maybe Double))
-> [CondVec]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (TreeConfig
-> Expr (Maybe Double) -> Vector (Maybe Double) -> [CondVec]
condsForMaybe TreeConfig
cfg Expr (Maybe Double)
e) (DataFrame -> Expr (Maybe Double) -> Maybe (Vector (Maybe Double))
interpretMaybeDoubleCol DataFrame
df Expr (Maybe Double)
e)

condsForDouble :: TreeConfig -> Expr Double -> V.Vector Double -> [CondVec]
condsForDouble :: TreeConfig -> Expr Double -> Vector Double -> [CondVec]
condsForDouble TreeConfig
cfg Expr Double
e Vector Double
vals = (Double -> [CondVec]) -> [Double] -> [CondVec]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Expr Double -> Vector Double -> Int -> Double -> [CondVec]
doubleCondsAt Expr Double
e Vector Double
vals (Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
vals)) [Double]
ts
  where
    ts :: [Double]
ts = [Int] -> [Double] -> [Double]
percentilesOf (TreeConfig -> [Int]
percentiles TreeConfig
cfg) (Vector Double -> [Double]
forall a. Vector a -> [a]
V.toList Vector Double
vals)

doubleCondsAt :: Expr Double -> V.Vector Double -> Int -> Double -> [CondVec]
doubleCondsAt :: Expr Double -> Vector Double -> Int -> Double -> [CondVec]
doubleCondsAt Expr Double
e Vector Double
vals Int
n Double
t =
    [ Expr Bool -> Vector Bool -> CondVec
CondVec (Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.<= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.>= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.< Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Expr Double
e Expr Double -> Expr Double -> Expr (NullCmpResult Double Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.> Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
t))
    ]
  where
    gen :: (Double -> Bool) -> Vector Bool
gen Double -> Bool
p = Int -> (Int -> Bool) -> Vector Bool
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n (\Int
i -> Double -> Bool
p (Vector Double
vals Vector Double -> Int -> Double
forall a. Vector a -> Int -> a
V.! Int
i))

condsForMaybe ::
    TreeConfig -> Expr (Maybe Double) -> V.Vector (Maybe Double) -> [CondVec]
condsForMaybe :: TreeConfig
-> Expr (Maybe Double) -> Vector (Maybe Double) -> [CondVec]
condsForMaybe TreeConfig
cfg Expr (Maybe Double)
e Vector (Maybe Double)
mvals = (Double -> [CondVec]) -> [Double] -> [CondVec]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Expr (Maybe Double)
-> Vector (Maybe Double) -> Int -> Double -> [CondVec]
maybeCondsAt Expr (Maybe Double)
e Vector (Maybe Double)
mvals (Vector (Maybe Double) -> Int
forall a. Vector a -> Int
V.length Vector (Maybe Double)
mvals)) [Double]
ts
  where
    ts :: [Double]
ts = [Int] -> [Double] -> [Double]
percentilesOf (TreeConfig -> [Int]
percentiles TreeConfig
cfg) ([Maybe Double] -> [Double]
forall a. [Maybe a] -> [a]
catMaybes (Vector (Maybe Double) -> [Maybe Double]
forall a. Vector a -> [a]
V.toList Vector (Maybe Double)
mvals))

maybeCondsAt ::
    Expr (Maybe Double) -> V.Vector (Maybe Double) -> Int -> Double -> [CondVec]
maybeCondsAt :: Expr (Maybe Double)
-> Vector (Maybe Double) -> Int -> Double -> [CondVec]
maybeCondsAt Expr (Maybe Double)
e Vector (Maybe Double)
mvals Int
n Double
t =
    [ Expr Bool -> Vector Bool -> CondVec
CondVec (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.<= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t)) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.>= Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t)) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.< Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t)) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
t))
    , Expr Bool -> Vector Bool -> CondVec
CondVec (Bool -> Expr (Maybe Bool) -> Expr Bool
forall a. Columnable a => a -> Expr (Maybe a) -> Expr a
F.fromMaybe Bool
False (Expr (Maybe Double)
e Expr (Maybe Double)
-> Expr Double -> Expr (NullCmpResult (Maybe Double) Double)
forall a b.
(NumericWidenOp (BaseType a) (BaseType b),
 NullLift2Op a b Bool (NullCmpResult a b),
 Ord (Promote (BaseType a) (BaseType b))) =>
Expr a -> Expr b -> Expr (NullCmpResult a b)
.> Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
t)) ((Double -> Bool) -> Vector Bool
gen (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
t))
    ]
  where
    gen :: (Double -> Bool) -> Vector Bool
gen Double -> Bool
p = Int -> (Int -> Bool) -> Vector Bool
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
n (\Int
i -> Bool -> (Double -> Bool) -> Maybe Double -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False Double -> Bool
p (Vector (Maybe Double)
mvals Vector (Maybe Double) -> Int -> Maybe Double
forall a. Vector a -> Int -> a
V.! Int
i))

{- | Arithmetic candidate expansion, generated already-deduped: each round
combines @frontier × base@ and admits only normalized-novel candidates.
Produces @base@ plus @maxExprDepth-1@ combination rounds.
-}
numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr]
numericExprsWithTerms :: SynthConfig -> DataFrame -> [NumExpr]
numericExprsWithTerms SynthConfig
cfg DataFrame
df
    | Bool -> Bool
not (SynthConfig -> Bool
enableArithOps SynthConfig
cfg) = [NumExpr]
base
    | Bool
otherwise =
        [NumExpr]
base [NumExpr] -> [NumExpr] -> [NumExpr]
forall a. [a] -> [a] -> [a]
++ SynthConfig
-> [NumExpr] -> Int -> [NumExpr] -> Set String -> [NumExpr]
expandRounds SynthConfig
cfg [NumExpr]
base (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (SynthConfig -> Int
maxExprDepth SynthConfig
cfg Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) [NumExpr]
base Set String
seen0
  where
    base :: [NumExpr]
base = DataFrame -> [NumExpr]
numericCols DataFrame
df
    seen0 :: Set String
seen0 = [String] -> Set String
forall a. Ord a => [a] -> Set a
Set.fromList ((NumExpr -> String) -> [NumExpr] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map NumExpr -> String
keyNum [NumExpr]
base)

keyNum :: NumExpr -> String
keyNum :: NumExpr -> String
keyNum (NDouble Expr Double
e) = Expr Double -> String
forall a. Show a => a -> String
show (Expr Double -> Expr Double
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr Double
e)
keyNum (NMaybeDouble Expr (Maybe Double)
e) = Expr (Maybe Double) -> String
forall a. Show a => a -> String
show (Expr (Maybe Double) -> Expr (Maybe Double)
forall a. (Show a, Typeable a) => Expr a -> Expr a
normalize Expr (Maybe Double)
e)

isDisallowed :: SynthConfig -> NumExpr -> NumExpr -> Bool
isDisallowed :: SynthConfig -> NumExpr -> NumExpr -> Bool
isDisallowed SynthConfig
cfg NumExpr
e1 NumExpr
e2 =
    ((Text, Text) -> Bool) -> [(Text, Text)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Text
l, Text
r) -> Text
l Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
cols Bool -> Bool -> Bool
&& Text
r Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
cols) (SynthConfig -> [(Text, Text)]
disallowedCombinations SynthConfig
cfg)
  where
    cols :: [Text]
cols = NumExpr -> [Text]
numExprCols NumExpr
e1 [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> NumExpr -> [Text]
numExprCols NumExpr
e2

roundProducts :: SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr]
roundProducts :: SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr]
roundProducts SynthConfig
cfg [NumExpr]
frontier [NumExpr]
base =
    [ NumExpr
c
    | NumExpr
e1 <- [NumExpr]
frontier
    , NumExpr
e2 <- [NumExpr]
base
    , Bool -> Bool
not (NumExpr -> NumExpr -> Bool
numExprEq NumExpr
e1 NumExpr
e2)
    , Bool -> Bool
not (SynthConfig -> NumExpr -> NumExpr -> Bool
isDisallowed SynthConfig
cfg NumExpr
e1 NumExpr
e2)
    , NumExpr
c <- NumExpr -> NumExpr -> [NumExpr]
combineNumExprs NumExpr
e1 NumExpr
e2
    ]

expandRounds ::
    SynthConfig -> [NumExpr] -> Int -> [NumExpr] -> Set.Set String -> [NumExpr]
expandRounds :: SynthConfig
-> [NumExpr] -> Int -> [NumExpr] -> Set String -> [NumExpr]
expandRounds SynthConfig
_ [NumExpr]
_ Int
0 [NumExpr]
_ Set String
_ = []
expandRounds SynthConfig
cfg [NumExpr]
base Int
d [NumExpr]
frontier Set String
seen
    | [NumExpr] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [NumExpr]
admitted = []
    | Bool
otherwise = [NumExpr]
admitted [NumExpr] -> [NumExpr] -> [NumExpr]
forall a. [a] -> [a] -> [a]
++ SynthConfig
-> [NumExpr] -> Int -> [NumExpr] -> Set String -> [NumExpr]
expandRounds SynthConfig
cfg [NumExpr]
base (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [NumExpr]
admitted Set String
seen'
  where
    ([NumExpr]
admitted, Set String
seen') = Set String -> [NumExpr] -> ([NumExpr], Set String)
admitNovel Set String
seen (SynthConfig -> [NumExpr] -> [NumExpr] -> [NumExpr]
roundProducts SynthConfig
cfg [NumExpr]
frontier [NumExpr]
base)

admitNovel :: Set.Set String -> [NumExpr] -> ([NumExpr], Set.Set String)
admitNovel :: Set String -> [NumExpr] -> ([NumExpr], Set String)
admitNovel Set String
seen0 = Set String -> [NumExpr] -> [NumExpr] -> ([NumExpr], Set String)
go Set String
seen0 []
  where
    go :: Set String -> [NumExpr] -> [NumExpr] -> ([NumExpr], Set String)
go Set String
seen [NumExpr]
acc [] = ([NumExpr] -> [NumExpr]
forall a. [a] -> [a]
reverse [NumExpr]
acc, Set String
seen)
    go Set String
seen [NumExpr]
acc (NumExpr
c : [NumExpr]
cs)
        | NumExpr -> String
keyNum NumExpr
c String -> Set String -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set String
seen = Set String -> [NumExpr] -> [NumExpr] -> ([NumExpr], Set String)
go Set String
seen [NumExpr]
acc [NumExpr]
cs
        | Bool
otherwise = Set String -> [NumExpr] -> [NumExpr] -> ([NumExpr], Set String)
go (String -> Set String -> Set String
forall a. Ord a => a -> Set a -> Set a
Set.insert (NumExpr -> String
keyNum NumExpr
c) Set String
seen) (NumExpr
c NumExpr -> [NumExpr] -> [NumExpr]
forall a. a -> [a] -> [a]
: [NumExpr]
acc) [NumExpr]
cs

numericCols :: DataFrame -> [NumExpr]
numericCols :: DataFrame -> [NumExpr]
numericCols DataFrame
df = (Text -> [NumExpr]) -> [Text] -> [NumExpr]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (DataFrame -> Text -> [NumExpr]
numExprsOfColumn DataFrame
df) (DataFrame -> [Text]
columnNames DataFrame
df)

numExprsOfColumn :: DataFrame -> T.Text -> [NumExpr]
numExprsOfColumn :: DataFrame -> Text -> [NumExpr]
numExprsOfColumn DataFrame
df Text
colName = case Text -> DataFrame -> Column
unsafeGetColumn Text
colName DataFrame
df of
    UnboxedColumn Maybe Bitmap
Nothing (Vector a
_ :: VU.Vector b) -> forall b. Columnable b => Text -> [NumExpr]
strictNumeric @b Text
colName
    BoxedColumn (Just Bitmap
_) (Vector a
_ :: V.Vector b) -> forall b. Columnable b => Text -> [NumExpr]
nullableNumeric @b Text
colName
    UnboxedColumn (Just Bitmap
_) (Vector a
_ :: VU.Vector b) -> forall b. Columnable b => Text -> [NumExpr]
nullableNumeric @b Text
colName
    Column
_ -> []

strictNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr]
strictNumeric :: forall b. Columnable b => Text -> [NumExpr]
strictNumeric Text
c = case TypeRep b -> TypeRep Double -> Maybe (b :~: Double)
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) of
    Just b :~: Double
Refl -> [Expr Double -> NumExpr
NDouble (Text -> Expr Double
forall a. Columnable a => Text -> Expr a
Col Text
c)]
    Maybe (b :~: Double)
Nothing -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @b of
        SBool (IntegralTypes b)
STrue -> [Expr Double -> NumExpr
NDouble (Expr b -> Expr Double
forall a. (Columnable a, Real a) => Expr a -> Expr Double
F.toDouble (forall a. Columnable a => Text -> Expr a
Col @b Text
c))]
        SBool (IntegralTypes b)
SFalse -> []

nullableNumeric :: forall b. (Columnable b) => T.Text -> [NumExpr]
nullableNumeric :: forall b. Columnable b => Text -> [NumExpr]
nullableNumeric Text
c = case TypeRep b -> TypeRep Double -> Maybe (b :~: Double)
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 @b) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @Double) of
    Just b :~: Double
Refl -> [Expr (Maybe Double) -> NumExpr
NMaybeDouble (forall a. Columnable a => Text -> Expr a
Col @(Maybe b) Text
c)]
    Maybe (b :~: Double)
Nothing -> case forall a. SBoolI (IntegralTypes a) => SBool (IntegralTypes a)
sIntegral @b of
        SBool (IntegralTypes b)
STrue -> [Expr (Maybe Double) -> NumExpr
NMaybeDouble ((b -> Double) -> Expr (Maybe b) -> Expr (Maybe Double)
forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr (Maybe a) -> Expr (Maybe b)
F.whenPresent (forall a b. (Real a, Fractional b) => a -> b
realToFrac @b @Double) (forall a. Columnable a => Text -> Expr a
Col @(Maybe b) Text
c))]
        SBool (IntegralTypes b)
SFalse -> []