Fixes #85

techascent · May 30, 2020 · f389625 · f389625
1 parent d49608e
commit f389625
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/src/tech/ml/dataset/base.clj b/src/tech/ml/dataset/base.clj
@@ -201,7 +201,10 @@
   index-seq - either keyword :all or list of indexes.  May contain duplicates.
   "
   [dataset colname-seq index-seq]
-  (ds-proto/select dataset colname-seq index-seq))
+  (let [index-seq (if (number? index-seq)
+                    [index-seq]
+                    index-seq)]
+    (ds-proto/select dataset colname-seq index-seq)))
 
 
 (defn unordered-select

diff --git a/test/tech/ml/dataset_test.clj b/test/tech/ml/dataset_test.clj
@@ -554,6 +554,26 @@
            (vec (dtype/->reader (ds :a) :object))))))
 
 
+(deftest select-row
+  (let [ds (ds/name-values-seq->dataset {:V1 (take 9 (cycle [1 2]))
+                                         :V2 (range 1 10)
+                                         :V3 (take 9 (cycle [0.5 1.0 1.5]))
+                                         :V4 (take 9 (cycle [\A \B \C]))})]
+
+    (is (= [2 6 1.5 \C]
+           (-> (ds/select-rows ds 5)
+               (ds/value-reader)
+               (first)
+               (vec))))
+
+    (is (= [2 6 1.5 \C]
+           (-> (ds/select-rows ds [5])
+               (ds/value-reader)
+               (first)
+               (vec))))))
+
+
+
 (comment
 
   (def test-ds (ds/->dataset