Fixes #34

src/tech/v3/datatype/argops.clj

@@ -49,11 +49,11 @@
        (= argtype :scalar)
        (const-reader/const-reader item n-const-elems)
        (= argtype :iterable)
-       (-> (dtype-cmc/make-container :list (dtype-base/elemwise-datatype item) {}
+       (-> (dtype-cmc/make-container :list (dtype-base/operational-elemwise-datatype item) {}
                                      item)
            (dtype-base/->reader))
        :else
-       (dtype-base/->reader item))))
+       (dtype-base/->reader item (dtype-base/operational-elemwise-datatype item)))))
   (^Buffer [item]
    (ensure-reader item Long/MAX_VALUE)))
 
@@ -271,8 +271,8 @@
   "Given a reader of values an a source comparator, return either an
   IntComparator or a LongComparator depending on the number of indexes
   in the reader that compares the values using the passed in comparator."
-  [src-comparator values]
-  (let [src-dtype (dtype-base/elemwise-datatype values)
+  [src-comparator nan-strategy values]
+  (let [src-dtype (dtype-base/operational-elemwise-datatype values)
         values (ensure-reader values)
         n-values (.lsize values)
         src-comparator (if-let [bin-pred (:binary-predicate (meta src-comparator))]
@@ -292,8 +292,18 @@
           (reify Comparators$IntComp
             (compareInts [this lhs rhs]
               (let [lhs-value (.readDouble values lhs)
-                    rhs-value (.readDouble values rhs)]
-                (.compare comp lhs-value rhs-value)))))
+                    rhs-value (.readDouble values rhs)
+                    lhs-nan? (Double/isNaN lhs-value)
+                    rhs-nan? (Double/isNaN rhs-value)]
+                (if (or lhs-nan? rhs-nan?)
+                  (if (identical? nan-strategy :exception)
+                    (throw (Exception. "##NaN value detected"))
+                    (if (and lhs-nan? rhs-nan?)
+                      0
+                      (if (identical? nan-strategy :first)
+                        (long (if lhs-nan? -1 1))
+                        (long (if lhs-nan? 1 -1)))))
+                  (.compare comp lhs-value rhs-value))))))
         :else
         (let [^Comparator comp (->comparator src-comparator)]
           (reify Comparators$IntComp
@@ -337,13 +347,15 @@
 
 (defn argsort
   "Sort values in index space.  By default uses a parallelized quicksort algorithm."
-  ([comparator {:keys [parallel?]
-           :or {parallel? true}}
+  ([comparator {:keys [parallel?
+                       nan-strategy]
+                :or {parallel? true
+                     nan-strategy :last}}
     values]
    (let [n-elems (dtype-base/ecount values)
-         val-dtype (dtype-base/elemwise-datatype values)
+         val-dtype (dtype-base/operational-elemwise-datatype values)
          comparator (-> (find-base-comparator comparator val-dtype)
-                        (index-comparator values))]
+                        (index-comparator nan-strategy values))]
      (cond
        (== n-elems 0)
        (int-array 0)

src/tech/v3/datatype/base.clj

@@ -38,6 +38,22 @@
     (dtype-proto/elemwise-datatype item)))
 
 
+(defn operational-elemwise-datatype
+  "Return the datatype one would expect when iterating through a container
+  of this type and performing a numeric operation on it.  For integer types, when a column
+  has missing values they get promoted to floating point numbers and ##NaN is substituted
+  in.  For scalars, return your elemental datatype."
+  [item]
+  ;;false has a datatype in this world.
+  (cond
+    (nil? item) :object
+    (instance? Double item) :float64
+    (instance? Integer item) :int32
+    (instance? Long item) :int64
+    :else
+    (dtype-proto/operational-elemwise-datatype item)))
+
+
 (defn datatype
   "Return this object's actual datatype.
   **This is not the same as the DEPRECATED get-datatype.  That function maps

test/tech/v3/datatype_test.clj

@@ -471,6 +471,15 @@
            (vec new-data)))))
 
 
+(deftest argsort-nan
+  (let [data (dtype/make-container :float64 [##NaN 2 5 ##NaN 1 6 3 4 ##NaN])
+        idx-first (argops/argsort nil {:nan-strategy :first} data)
+        idx-last (argops/argsort nil {:nan-strategy :last} data)]
+    (is (= [0 3 8] (vec (take 3 idx-first))))
+    (is (= [0 3 8] (vec (take-last 3 idx-last))))
+    (is (thrown? Exception (argops/argsort nil {:nan-strategy :exception} data)))))
+
+
 (deftest reader-as-persistent-vector-test
   (let [src-data (range 20)
         ldata (long-array src-data)