Add positional renaming of columns #262

src/tech/v3/dataset/base.clj

@@ -342,16 +342,34 @@
   (select-by-index dataset col-index :all))
 
 (defn rename-columns
-  "Rename columns using a map.  Does not reorder columns."
-  [dataset colname-map]
-  (->> (ds-proto/columns dataset)
-       (map (fn [col]
-              (let [old-name (ds-col/column-name col)]
-                (if (contains? colname-map old-name)
-                  (ds-col/set-name col (get colname-map old-name))
-                  col))))
-       (ds-impl/new-dataset (dataset-name dataset)
-                            (meta dataset))))
+  "Rename columns using a map or vector of column names.
+
+  Does not reorder columns; rename is in-place for maps and
+  positional for vectors."
+  [dataset colnames]
+  (let [col-map? (map? colnames)
+        col-vector? (vector? colnames)
+        existing-cols (ds-proto/columns dataset)]
+    (cond (not (or col-map? col-vector?))
+          (throw (ex-info "column names must be a map or vector"
+                          {:column-names-type (type colnames)}))
+          (and col-vector? (not= (count existing-cols)
+                                 (count colnames)))
+          (throw (ex-info "rename column vector must be of equal size to current column count"
+                          {:current-column-count (count existing-cols)
+                           :target-column-count (count colnames)})))
+    (->> (map
+          (fn [col new-name]
+            (let [old-name (ds-col/column-name col)]
+              (if col-map?
+                (if (contains? colnames old-name)
+                  (ds-col/set-name col (get colnames old-name))
+                  col)
+                (ds-col/set-name col new-name))))
+          existing-cols
+          colnames)
+         (ds-impl/new-dataset (dataset-name dataset)
+                              (meta dataset)))))
 
 
 (defn select-rows

test/tech/v3/dataset_test.clj

@@ -17,7 +17,7 @@
             [taoensso.nippy :as nippy]
             [clojure.test :refer [deftest is]])
   (:import [java.util List HashSet UUID]
-           [java.io File]))
+           [java.io File ByteArrayInputStream]))
 
 (deftest datatype-parser
   (let [ds (ds/->dataset "test/data/datatype_parser.csv")]
@@ -678,6 +678,22 @@
                (ds/column false)
                vec)))))
 
+(deftest positional-column-rename
+  (let [DS (ds/->dataset
+            (-> "id,a,a\n0,aa,bb\n1,cc,dd"
+                .getBytes
+                ByteArrayInputStream.)
+            {:file-type :csv})
+        new-cols-incorrect [:a1 :a2]
+        new-cols-correct [:id :a1 :a2]]
+    (is (= new-cols-correct
+           (-> DS
+               (ds/rename-columns new-cols-correct)
+               ds/column-names)))
+    (is (thrown? Throwable
+                 (ds/rename-columns DS new-cols-incorrect)))
+    (is (thrown? Throwable
+                 (ds/rename-columns DS (set new-cols-correct))))))
 
 (deftest column-sequences-use-nil-missing
   (let [ds (ds/->dataset [{:a 1} {:b 2}])]