From b7e20da0b50f47a097a6f5a045be5bd0977e1e81 Mon Sep 17 00:00:00 2001
From: fserucas <97905656+fserucas@users.noreply.github.com>
Date: Wed, 13 Aug 2025 10:12:05 +0100
Subject: [PATCH] Add imputeExpr function

This change adds a `imputeExpr` function which behaves like `impute`
but accepts Col Expr.
---
 src/DataFrame/Operations/Transformations.hs | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs
index 6b3d2047..b88f6319 100644
--- a/src/DataFrame/Operations/Transformations.hs
+++ b/src/DataFrame/Operations/Transformations.hs
@@ -164,3 +164,14 @@ impute columnName value df = case getColumn columnName df of
     Left exception -> throw exception
     Right res      -> res
   _ -> error "Cannot impute to a non-Empty column"
+
+imputeExpr :: forall a. (Columnable a) => T.Text -> Expr a -> DataFrame -> DataFrame
+--imputeExpr = undefined
+imputeExpr columnName expr df =
+  case getColumn columnName df of
+    Nothing -> throw $ ColumnNotFoundException columnName "impute" (map fst $ M.toList $ columnIndices df)
+    Just (OptionalColumn _) ->
+      case expr of
+        Col a -> impute columnName a df
+        _ -> error "Cannot impute a non Col Expression"
+    _ -> error "Cannot impute to a non-Empty column"