allow column filter function as column-selector in replace-missing (#183

) * replace-missing with fn as columns-selector * update docs of replace-missing
techascent · Dec 10, 2020 · fc20dfb · fc20dfb
1 parent 309dc98
commit fc20dfb
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/src/tech/v3/dataset/missing.clj b/src/tech/v3/dataset/missing.clj
@@ -176,8 +176,13 @@
 
 (defn replace-missing
   "Replace missing values in some columns with a given strategy.
-  The columns selector may be any legal argument to select-columns.
+  The columns selector may be:
+
+  - seq of any legal column names
+  - or a column filter function, such as `numeric` and `categorical`
+
   Strategies may be:
+
   - `:down` - take value from previous non-missing row if possible else use next
     non-missing row.
   - `:up` - take value from next non-missing row if possible else use previous
@@ -186,17 +191,21 @@
      rows.
   - `:lerp` - Linearly interpolate values between previous and next nonmissing rows.
   - `:value` - Value will be provided - see below.
-  value may be provided which will then be used.  Value may be a function in which
-  case it will be called on the column with missing values elided and the return will
-  be used to as the filler."
+
+      value may be provided which will then be used.  Value may be a function in which
+      case it will be called on the column with missing values elided and the return will
+      be used to as the filler."
   ([ds] (replace-missing ds :mid))
   ([ds strategy] (replace-missing ds :all strategy))
   ([ds columns-selector strategy]
    (replace-missing ds columns-selector strategy nil))
   ([ds columns-selector strategy value]
    (let [strategy (or strategy :mid)
-         row-cnt (ds-base/row-count ds)]
-     (->> (ds-base/select-columns ds columns-selector)
+         row-cnt (ds-base/row-count ds)
+         selected (if (fn? columns-selector)
+                    (columns-selector ds)
+                    (ds-base/select-columns ds columns-selector))]
+     (->> selected
           (ds-base/columns)
           (reduce (fn [ds col]
                     (let [^RoaringBitmap missing (col/missing col)]

diff --git a/test/tech/v3/dataset_test.clj b/test/tech/v3/dataset_test.clj
@@ -8,6 +8,7 @@
             [tech.v3.dataset.string-table :as str-table]
             [tech.v3.dataset.join :as ds-join]
             [tech.v3.dataset.test-utils :as test-utils]
+            [tech.v3.dataset.column-filters :as cf]
             ;;Loading multimethods required to load the files
             [tech.v3.libs.poi]
             [tech.v3.libs.fastexcel]
@@ -739,6 +740,17 @@
                  (ds/missing)
                  (dtype/ecount))))))
 
+(deftest replace-missing-selector-fn
+  (let [ds (ds/->dataset {:a [nil nil 2 4]
+                          :b [nil nil 4 6]
+                          :c [nil nil "A" "B"]})
+        ds-replaced (-> ds
+                        (ds/replace-missing cf/numeric :value dfn/mean)
+                        (ds/replace-missing cf/categorical :value "C"))]
+    (is (= [3 3 2 4] (vec (ds-replaced :a))))
+    (is (= [5 5 4 6] (vec (ds-replaced :b))))
+    (is (= ["C" "C" "A" "B"] (vec (ds-replaced :c))))))
+
 
 (deftest replace-missing-ldt
   (let [dtds (ds/->dataset {:dt [(java.time.LocalDateTime/of 2020 1 1 1 1 1)