Adding NaN handeling to bolt

bolt/spark/array.py

@@ -334,6 +334,21 @@ def mean(self, axis=None, keepdims=False):
         """
         return self._stat(axis, name='mean', keepdims=keepdims)
 
+    def nanmean(self, axis=None, keepdims=False):
+        """
+        Return the nanmean of the array over the given axis.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nanmean', keepdims=keepdims)
+
     def var(self, axis=None, keepdims=False):
         """
         Return the variance of the array over the given axis.
@@ -347,7 +362,22 @@ def var(self, axis=None, keepdims=False):
         keepdims : boolean, optional, default=False
             Keep axis remaining after operation with size 1.
         """
-        return self._stat(axis, name='variance', keepdims=keepdims)
+        return self._stat(axis, name='var', keepdims=keepdims)
+
+    def nanvar(self, axis=None, keepdims=False):
+        """
+        Return the variance of the array over the given axis ignoring NaNs.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nanvar', keepdims=keepdims)
 
     def std(self, axis=None, keepdims=False):
         """
@@ -362,7 +392,22 @@ def std(self, axis=None, keepdims=False):
         keepdims : boolean, optional, default=False
             Keep axis remaining after operation with size 1.
         """
-        return self._stat(axis, name='stdev', keepdims=keepdims)
+        return self._stat(axis, name='std', keepdims=keepdims)
+
+    def nanstd(self, axis=None, keepdims=False):
+        """
+        Return the standard deviation of the array over the given axis ignoring NaNs.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nanstd', keepdims=keepdims)
 
     def sum(self, axis=None, keepdims=False):
         """
@@ -380,6 +425,21 @@ def sum(self, axis=None, keepdims=False):
         from operator import add
         return self._stat(axis, func=add, keepdims=keepdims)
 
+    def nansum(self, axis=None, keepdims=False):
+        """
+        Return the sum of the array over the given axis ignoring NaNs.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nansum', keepdims=keepdims)
+
     def max(self, axis=None, keepdims=False):
         """
         Return the maximum of the array over the given axis.
@@ -396,6 +456,21 @@ def max(self, axis=None, keepdims=False):
         from numpy import maximum
         return self._stat(axis, func=maximum, keepdims=keepdims)
 
+    def nanmax(self, axis=None, keepdims=False):
+        """
+        Return the maximum of the array over the given axis ignoring NaNs.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nanmax', keepdims=keepdims)
+
     def min(self, axis=None, keepdims=False):
         """
         Return the minimum of the array over the given axis.
@@ -412,6 +487,36 @@ def min(self, axis=None, keepdims=False):
         from numpy import minimum
         return self._stat(axis, func=minimum, keepdims=keepdims)
 
+    def nanmin(self, axis=None, keepdims=False):
+        """
+        Return the minimum of the array over the given axis ignoring NaNs.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nanmin', keepdims=keepdims)
+
+    def nancount(self, axis=None, keepdims=False):
+        """
+        Return the count of non NaN values.
+
+        Parameters
+        ----------
+        axis : tuple or int, optional, default=None
+            Axis to compute statistic over, if None
+            will compute over all axes
+
+        keepdims : boolean, optional, default=False
+            Keep axis remaining after operation with size 1.
+        """
+        return self._stat(axis, name='nancount', keepdims=keepdims)
+
     def concatenate(self, arry, axis=0):
         """
         Join this array with another array.