feat: Add type level impute function.

Author: mchav

src/DataFrame/Typed.hs

@@ -130,6 +130,7 @@ module DataFrame.Typed (
 
     -- * Schema-modifying operations
     derive,
+    impute,
     select,
     exclude,
     rename,
@@ -174,6 +175,7 @@ module DataFrame.Typed (
     ExcludeSchema,
     RenameInSchema,
     RemoveColumn,
+    Impute,
     Append,
     Reverse,
     StripAllMaybe,
@@ -212,7 +214,6 @@ import DataFrame.Typed.Schema
 import DataFrame.Typed.TH (deriveSchema, deriveSchemaFromCsvFile)
 import DataFrame.Typed.Types (
     Column,
-    TExpr (..),
     TSortOrder (..),
     These (..),
     TypedDataFrame,

src/DataFrame/Typed/Aggregate.hs

@@ -28,7 +28,7 @@ import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)
 
 import DataFrame.Internal.Column (Columnable)
 import qualified DataFrame.Internal.DataFrame as D
-import DataFrame.Internal.Expression (Expr, NamedExpr, UExpr (..))
+import DataFrame.Internal.Expression (NamedExpr)
 import qualified DataFrame.Operations.Aggregation as DA
 
 import DataFrame.Typed.Freeze (unsafeFreeze)

src/DataFrame/Typed/Expr.hs

@@ -83,7 +83,6 @@ module DataFrame.Typed.Expr (
     desc,
 ) where
 
-import Data.Kind (Type)
 import Data.Proxy (Proxy (..))
 import Data.String (IsString (..))
 import qualified Data.Text as T
@@ -100,7 +99,7 @@ import DataFrame.Internal.Expression (
     UnaryOp (..),
  )
 import DataFrame.Typed.Schema (AssertPresent, Lookup)
-import DataFrame.Typed.Types (Column, TExpr (..), TSortOrder (..))
+import DataFrame.Typed.Types (TExpr (..), TSortOrder (..))
 
 -------------------------------------------------------------------------------
 -- Column reference — the core type-safe construction point

src/DataFrame/Typed/Freeze.hs

@@ -15,11 +15,8 @@ module DataFrame.Typed.Freeze (
     unsafeFreeze,
 ) where
 
-import Data.Kind (Type)
-import qualified Data.Map as M
-import Data.Proxy (Proxy (..))
 import qualified Data.Text as T
-import Type.Reflection (SomeTypeRep, someTypeRep)
+import Type.Reflection (SomeTypeRep)
 
 import qualified DataFrame.Internal.Column as C
 import qualified DataFrame.Internal.DataFrame as D

src/DataFrame/Typed/Join.hs

@@ -15,16 +15,11 @@ module DataFrame.Typed.Join (
     fullOuterJoin,
 ) where
 
-import Data.Proxy (Proxy (..))
-import qualified Data.Text as T
-import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)
+import GHC.TypeLits (Symbol)
 
-import DataFrame.Internal.Column (Columnable)
-import qualified DataFrame.Internal.DataFrame as D
-import qualified DataFrame.Operations.Core as D
 import qualified DataFrame.Operations.Join as DJ
 
-import DataFrame.Typed.Freeze (thaw, unsafeFreeze)
+import DataFrame.Typed.Freeze (unsafeFreeze)
 import DataFrame.Typed.Schema
 import DataFrame.Typed.Types (TypedDataFrame (..))
 

src/DataFrame/Typed/Operations.hs

@@ -30,6 +30,7 @@ module DataFrame.Typed.Operations (
 
     -- * Schema-modifying operations
     derive,
+    impute,
     select,
     exclude,
     rename,
@@ -54,7 +55,6 @@ module DataFrame.Typed.Operations (
     (|>),
 ) where
 
-import qualified Data.Foldable as F
 import Data.Function ((&))
 import Data.Proxy (Proxy (..))
 import qualified Data.Text as T
@@ -63,9 +63,9 @@ import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)
 import System.Random (RandomGen)
 import Prelude hiding (drop, filter, take)
 
+import qualified DataFrame.Functions as DF
 import DataFrame.Internal.Column (Columnable)
 import qualified DataFrame.Internal.Column as C
-import DataFrame.Internal.Expression (Expr (..))
 import qualified DataFrame.Operations.Aggregation as DA
 import qualified DataFrame.Operations.Core as D
 import DataFrame.Operations.Merge ()
@@ -75,7 +75,7 @@ import qualified DataFrame.Operations.Transformations as D
 
 -- Semigroup instance
 
-import DataFrame.Typed.Freeze (thaw, unsafeFreeze)
+import DataFrame.Typed.Freeze (unsafeFreeze)
 import DataFrame.Typed.Schema
 import DataFrame.Typed.Types (TExpr (..), TSortOrder (..), TypedDataFrame (..))
 import qualified DataFrame.Typed.Types as T
@@ -222,6 +222,20 @@ derive (TExpr expr) (TDF df) = unsafeFreeze (D.derive colName expr df)
   where
     colName = T.pack (symbolVal (Proxy @name))
 
+impute ::
+    forall name a cols.
+    ( KnownSymbol name
+    , Columnable a
+    ) =>
+    a ->
+    TypedDataFrame cols ->
+    TypedDataFrame (Impute name cols)
+impute value (TDF df) =
+    unsafeFreeze
+        (D.derive colName (DF.fromMaybe value (DF.col @(Maybe a) colName)) df)
+  where
+    colName = T.pack (symbolVal (Proxy @name))
+
 -- | Select a subset of columns by name.
 select ::
     forall (names :: [Symbol]) cols.

src/DataFrame/Typed/Schema.hs

@@ -18,6 +18,7 @@ module DataFrame.Typed.Schema (
     Lookup,
     HasName,
     RemoveColumn,
+    Impute,
     SubsetSchema,
     ExcludeSchema,
     RenameInSchema,
@@ -60,9 +61,10 @@ import Data.Proxy (Proxy (..))
 import qualified Data.Text as T
 import Data.These (These)
 import GHC.TypeLits
-import Type.Reflection (SomeTypeRep, Typeable, someTypeRep, typeRep)
+import Type.Reflection (SomeTypeRep, Typeable, someTypeRep)
 
 import DataFrame.Internal.Column (Columnable)
+import DataFrame.Internal.Types (If)
 import DataFrame.Typed.Types (Column)
 
 -------------------------------------------------------------------------------
@@ -77,6 +79,14 @@ type family Lookup (name :: Symbol) (cols :: [Type]) :: Type where
         TypeError
             ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' not found in schema")
 
+type family Impute (name :: Symbol) (cols :: [Type]) :: [Type] where
+    Impute name (Column name (Maybe a) ': rest) = Column name a ': rest
+    Impute name (Column name _ ': rest) =
+        TypeError
+            ('Text "Column '" ':<>: 'Text name ':<>: 'Text "' is not of kind Maybe *")
+    Impute name (col ': rest) = col ': Impute name rest
+    Impute name '[] = '[]
+
 -- | Add type to the end of a list.
 type family Snoc (xs :: [k]) (x :: k) :: [k] where
     Snoc '[] x = '[x]
@@ -108,11 +118,6 @@ type family ExcludeSchema (names :: [Symbol]) (cols :: [Type]) :: [Type] where
             (ExcludeSchema names rest)
             (Column n a ': ExcludeSchema names rest)
 
--- | Type-level if
-type family If (b :: Bool) (t :: k) (f :: k) :: k where
-    If 'True t _ = t
-    If 'False _ f = f
-
 -- | Type-level elem for Symbols
 type family IsElem (x :: Symbol) (xs :: [Symbol]) :: Bool where
     IsElem x '[] = 'False

src/DataFrame/Typed/TH.hs

@@ -20,7 +20,6 @@ import qualified Data.Map as M
 import qualified Data.Text as T
 
 import Language.Haskell.TH
-import Language.Haskell.TH.Syntax (Lift (..))
 
 import qualified DataFrame.IO.CSV as D
 import qualified DataFrame.Internal.Column as C