Fixes 242

techascent · May 10, 2021 · fd1211a · fd1211a
1 parent 1732b32
commit fd1211a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/tech/v3/dataset/base.clj b/src/tech/v3/dataset/base.clj
@@ -153,10 +153,10 @@
   [dataset colname-seq-or-fn]
   (let [colname-seq (set (if (fn? colname-seq-or-fn)
                            (column-names (colname-seq-or-fn dataset))
-                           colname-seq-or-fn))
-        all-cols (set (column-names dataset))
-        leftover (set/difference all-cols colname-seq)]
-    (select-columns dataset leftover)))
+                           colname-seq-or-fn))]
+    (select-columns dataset
+                    (->> (column-names dataset)
+                         (clojure.core/filter (complement colname-seq))))))
 
 
 (defn drop-columns

diff --git a/test/tech/v3/dataset_test.clj b/test/tech/v3/dataset_test.clj
@@ -1144,6 +1144,14 @@
                                 [:b])))))
 
 
+(deftest remove-columns-issue-242
+  (is (= [:a "c" :d :e]
+         (vec (-> (tech.v3.dataset/->dataset {:a [1] :b [2] "c" [3]
+                                              :d [4] :e [5]})
+                  (tech.v3.dataset/drop-columns [:b])
+                  (ds/column-names))))))
+
+
 (comment
 
   (def test-ds (ds/->dataset