From a5d888e49b4f36b2d7db1bdf00b958be639a497c Mon Sep 17 00:00:00 2001
From: Peter Schafhalter <pschafhalter@berkeley.edu>
Date: Mon, 4 Jun 2018 10:50:13 -0700
Subject: [PATCH] [DataFrames] More dtypes optimizations (#2124)

* Pass dtypes for some DataFrame constructors

* More optimizations with dtypes_cache

* Optimizations
---
 python/ray/dataframe/dataframe.py | 46 +++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 12 deletions(-)

diff --git a/python/ray/dataframe/dataframe.py b/python/ray/dataframe/dataframe.py
index bdcb54cf1..84303dea8 100644
--- a/python/ray/dataframe/dataframe.py
+++ b/python/ray/dataframe/dataframe.py
@@ -584,7 +584,8 @@ class DataFrame(object):
         return DataFrame(block_partitions=self._block_partitions,
                          columns=new_cols,
                          col_metadata=self._col_metadata,
-                         row_metadata=self._row_metadata)
+                         row_metadata=self._row_metadata,
+                         dtypes_cache=self._dtypes_cache)
 
     def add_suffix(self, suffix):
         """Add a suffix to each of the column names.
@@ -596,7 +597,8 @@ class DataFrame(object):
         return DataFrame(block_partitions=self._block_partitions,
                          columns=new_cols,
                          col_metadata=self._col_metadata,
-                         row_metadata=self._row_metadata)
+                         row_metadata=self._row_metadata,
+                         dtypes_cache=self._dtypes_cache)
 
     def applymap(self, func):
         """Apply a function to a DataFrame elementwise.
@@ -625,7 +627,7 @@ class DataFrame(object):
         return DataFrame(block_partitions=self._block_partitions,
                          columns=self.columns,
                          index=self.index,
-                         dtypes_cache=self.dtypes)
+                         dtypes_cache=self._dtypes_cache)
 
     def groupby(self, by=None, axis=0, level=None, as_index=True, sort=True,
                 group_keys=True, squeeze=False, **kwargs):
@@ -682,7 +684,7 @@ class DataFrame(object):
         return self._arithmetic_helper(remote_func, axis, level)
 
     def abs(self):
-        """Apply an absolute value function to all numberic columns.
+        """Apply an absolute value function to all numeric columns.
 
         Returns:
             A new DataFrame with the applied absolute value.
@@ -698,7 +700,8 @@ class DataFrame(object):
 
         return DataFrame(block_partitions=new_block_partitions,
                          columns=self.columns,
-                         index=self.index)
+                         index=self.index,
+                         dtypes_cache=self._dtypes_cache)
 
     def isin(self, values):
         """Fill a DataFrame with booleans for cells contained in values.
@@ -732,9 +735,13 @@ class DataFrame(object):
         new_block_partitions = np.array([_map_partitions(
             lambda df: df.isna(), block) for block in self._block_partitions])
 
+        new_dtypes = pd.Series([np.dtype("bool")] * len(self.columns),
+                               index=self.columns)
+
         return DataFrame(block_partitions=new_block_partitions,
                          row_metadata=self._row_metadata,
-                         col_metadata=self._col_metadata)
+                         col_metadata=self._col_metadata,
+                         dtypes_cache=new_dtypes)
 
     def isnull(self):
         """Fill a DataFrame with booleans for cells containing a null value.
@@ -749,9 +756,13 @@ class DataFrame(object):
             lambda df: df.isnull(), block)
             for block in self._block_partitions])
 
+        new_dtypes = pd.Series([np.dtype("bool")] * len(self.columns),
+                               index=self.columns)
+
         return DataFrame(block_partitions=new_block_partitions,
                          row_metadata=self._row_metadata,
-                         col_metadata=self._col_metadata)
+                         col_metadata=self._col_metadata,
+                         dtypes_cache=new_dtypes)
 
     def keys(self):
         """Get the info axis for the DataFrame.
@@ -2172,7 +2183,8 @@ class DataFrame(object):
 
         return DataFrame(col_partitions=new_dfs,
                          col_metadata=self._col_metadata,
-                         index=index)
+                         index=index,
+                         dtypes_cache=self._dtypes_cache)
 
     def hist(self, data, column=None, by=None, grid=True, xlabelsize=None,
              xrot=None, ylabelsize=None, yrot=None, ax=None, sharex=False,
@@ -2982,9 +2994,13 @@ class DataFrame(object):
         new_block_partitions = np.array([_map_partitions(
             lambda df: df.notna(), block) for block in self._block_partitions])
 
+        new_dtypes = pd.Series([np.dtype("bool")] * len(self.columns),
+                               index=self.columns)
+
         return DataFrame(block_partitions=new_block_partitions,
                          row_metadata=self._row_metadata,
-                         col_metadata=self._col_metadata)
+                         col_metadata=self._col_metadata,
+                         dtypes_cache=new_dtypes)
 
     def notnull(self):
         """Perform notnull across the DataFrame.
@@ -3000,9 +3016,13 @@ class DataFrame(object):
             lambda df: df.notnull(), block)
             for block in self._block_partitions])
 
+        new_dtypes = pd.Series([np.dtype("bool")] * len(self.columns),
+                               index=self.columns)
+
         return DataFrame(block_partitions=new_block_partitions,
                          row_metadata=self._row_metadata,
-                         col_metadata=self._col_metadata)
+                         col_metadata=self._col_metadata,
+                         dtypes_cache=new_dtypes)
 
     def nsmallest(self, n, columns, keep='first'):
         raise NotImplementedError(
@@ -4146,7 +4166,8 @@ class DataFrame(object):
             return DataFrame(row_partitions=new_row_partitions,
                              col_partitions=new_column_partitions,
                              columns=new_columns,
-                             index=new_index)
+                             index=new_index,
+                             dtypes_cache=self._dtypes_cache)
 
     def sortlevel(self, level=0, axis=0, ascending=True, inplace=False,
                   sort_remaining=True):
@@ -4239,7 +4260,8 @@ class DataFrame(object):
         index = self._row_metadata.index[-n:]
         return DataFrame(col_partitions=new_dfs,
                          col_metadata=self._col_metadata,
-                         index=index)
+                         index=index,
+                         dtypes_cache=self._dtypes_cache)
 
     def take(self, indices, axis=0, convert=None, is_copy=True, **kwargs):
         raise NotImplementedError(