Repository: ibis-project/ibis

Issue:
_fetch_from_cursor() ignores schema when constructing a pandas DataFrame
When executing a `TableExpr`, ibis returns a `DataFrame` with the results by calling the impyla cursor's `fetchall()` method.  This method gets back data from thrift already preconverted into Python types, in a way that mostly makes sense.  However, the standard python types are less rich than the Impala types, so an `int8` for example will be converted to the standard Python `int`, even though there is schema information available.  Since pandas supports the richer data types, it should recast the data to the original expected types, which are available from the cursor object.  There should be no risk of losing information in the downcasts, as the schemas should be compatible.



Discussion:

Resolving patch:
diff --git a/ibis/client.py b/ibis/client.py
--- a/ibis/client.py
+++ b/ibis/client.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from itertools import izip
+
 import hdfs
 
 from impala.error import Error as ImpylaError
@@ -120,11 +122,23 @@ def _build_ast_ensure_limit(self, expr, default_limit):
                     ast = sql.build_ast(expr)
         return ast, expr
 
+    def _db_type_to_dtype(self, db_type):
+        raise NotImplementedError
+
     def _fetch_from_cursor(self, cursor):
         import pandas as pd
         rows = cursor.fetchall()
+        # TODO(wesm): please evaluate/reimpl to optimize for perf/memory
+        dtypes = [self._db_type_to_dtype(x[1]) for x in cursor.description]
         names = [x[0] for x in cursor.description]
-        return pd.DataFrame.from_records(rows, columns=names)
+        cols = {}
+        for (col, name, dtype) in zip(izip(*rows), names, dtypes):
+            try:
+                cols[name] = pd.Series(col, dtype=dtype)
+            except TypeError:
+                # coercing to specified dtype failed, e.g. NULL vals in int col
+                cols[name] = pd.Series(col)
+        return pd.DataFrame(cols, columns=names)
 
 
 class ImpalaConnection(object):
@@ -214,6 +228,22 @@ class ImpalaClient(SQLClient):
     An Ibis client interface that uses Impala
     """
 
+    _HS2_TTypeId_to_dtype = {
+        'BOOLEAN': 'bool',
+        'TINYINT': 'int8',
+        'SMALLINT': 'int16',
+        'INT': 'int32',
+        'BIGINT': 'int64',
+        'TIMESTAMP': 'datetime64[ns]',
+        'FLOAT': 'float32',
+        'DOUBLE': 'float64',
+        'STRING': 'string',
+        'DECIMAL': 'object',
+        'BINARY': 'string',
+        'VARCHAR': 'string',
+        'CHAR': 'string'
+    }
+
     def __init__(self, con, hdfs_client=None, **params):
         self.con = con
 
@@ -258,6 +288,9 @@ def _fully_qualified_name(self, name, database):
         database = database or self.current_database
         return '{0}.`{1}`'.format(database, name)
 
+    def _db_type_to_dtype(self, db_type):
+        return self._HS2_TTypeId_to_dtype[db_type]
+
     def list_tables(self, like=None, database=None):
         """
         List tables in the current (or indicated) database. Like the SHOW


Test patch:
diff --git a/ibis/tests/test_impala_e2e.py b/ibis/tests/test_impala_e2e.py
--- a/ibis/tests/test_impala_e2e.py
+++ b/ibis/tests/test_impala_e2e.py
@@ -519,6 +519,16 @@ def test_int_builtins(self):
 
         self.assert_cases_equality(mod_cases + nullif_cases)
 
+    def test_column_types(self):
+        df = self.alltypes.execute()
+        assert df.tinyint_col.dtype.name == 'int8'
+        assert df.smallint_col.dtype.name == 'int16'
+        assert df.int_col.dtype.name == 'int32'
+        assert df.bigint_col.dtype.name == 'int64'
+        assert df.float_col.dtype.name == 'float32'
+        assert df.double_col.dtype.name == 'float64'
+        assert pd.core.common.is_datetime64_dtype(df.timestamp_col.dtype)
+
     def test_timestamp_builtins(self):
         i32 = ibis.literal(50000)
         i64 = ibis.literal(5 * 10 ** 8)