Repository: ibis-project/ibis

Issue:
Add option to set Parquet compression codec
Will prioritize for 0.6



Discussion:

Resolving patch:
diff --git a/ibis/impala/client.py b/ibis/impala/client.py
--- a/ibis/impala/client.py
+++ b/ibis/impala/client.py
@@ -73,11 +73,12 @@ class ImpalaConnection(object):
 
     def __init__(self, pool_size=8, database='default', **params):
         self.params = params
-        self.codegen_disabled = False
         self.database = database
 
         self.lock = threading.Lock()
 
+        self.options = {}
+
         self.connection_pool = queue.Queue(pool_size)
         self.connection_pool_size = 0
         self.max_pool_size = pool_size
@@ -86,6 +87,9 @@ def __init__(self, pool_size=8, database='default', **params):
 
         self.ping()
 
+    def set_options(self, options):
+        self.options.update(options)
+
     def close(self):
         """
         Close all open Impyla sessions
@@ -97,7 +101,11 @@ def set_database(self, name):
         self.database = name
 
     def disable_codegen(self, disabled=True):
-        self.codegen_disabled = disabled
+        key = 'DISABLE_CODEGEN'
+        if disabled:
+            self.options[key] = '1'
+        elif key in self.options:
+            del self.options[key]
 
     def execute(self, query, async=False):
         if isinstance(query, DDL):
@@ -134,10 +142,10 @@ def fetchall(self, query):
     def _get_cursor(self):
         try:
             cur = self.connection_pool.get(False)
-            if cur.database != self.database:
+            if (cur.database != self.database or
+                    cur.options != self.options):
                 cur = self._new_cursor()
-            if cur.codegen_disabled != self.codegen_disabled:
-                cur.disable_codegen(self.codegen_disabled)
+
             return cur
         except queue.Empty:
             if self.connection_pool_size < self.max_pool_size:
@@ -157,26 +165,28 @@ def _new_cursor(self):
         cursor = con.cursor(convert_types=True)
         cursor.ping()
 
-        wrapper = ImpalaCursor(cursor, self, con, self.database)
-
-        if self.codegen_disabled:
-            wrapper.disable_codegen(self.codegen_disabled)
+        wrapper = ImpalaCursor(cursor, self, con, self.database,
+                               self.options.copy())
+        wrapper.set_options()
 
         return wrapper
 
     def ping(self):
         self._new_cursor()
 
+    def release(self, cur):
+        self.connection_pool.put(cur)
+
 
 class ImpalaCursor(object):
 
     def __init__(self, cursor, con, impyla_con, database,
-                 codegen_disabled=False):
+                 options):
         self.cursor = cursor
         self.con = con
         self.impyla_con = impyla_con
         self.database = database
-        self.codegen_disabled = codegen_disabled
+        self.options = options
 
     def __del__(self):
         self._close_cursor()
@@ -195,18 +205,17 @@ def __enter__(self):
     def __exit__(self, type, value, tb):
         self.release()
 
-    def disable_codegen(self, disabled=True):
-        self.codegen_disabled = disabled
-        query = ('SET disable_codegen={0}'
-                 .format('true' if disabled else 'false'))
-        self.cursor.execute(query)
+    def set_options(self):
+        for k, v in self.options.items():
+            query = 'SET {0}={1}'.format(k, v)
+            self.cursor.execute(query)
 
     @property
     def description(self):
         return self.cursor.description
 
     def release(self):
-        self.con.connection_pool.put(self)
+        self.con.release(self)
 
     def execute(self, stmt, async=False):
         self.cursor.execute_async(stmt)
@@ -684,6 +693,39 @@ def get_schema(self, table_name, database=None):
 
         return dt.Schema(names, ibis_types)
 
+    @property
+    def client_options(self):
+        return self.con.options
+
+    def get_options(self):
+        """
+        Return current query options for the Impala session
+        """
+        query = 'SET'
+        tuples = self.con.fetchall(query)
+        return dict(tuples)
+
+    def set_options(self, options):
+        self.con.set_options(options)
+
+    def reset_options(self):
+        # Must nuke all cursors
+        raise NotImplementedError
+
+    def set_compression_codec(self, codec):
+        """
+        Parameters
+        """
+        if codec is None:
+            codec = 'none'
+        else:
+            codec = codec.lower()
+
+        if codec not in ('none', 'gzip', 'snappy'):
+            raise ValueError('Unknown codec: {0}'.format(codec))
+
+        self.set_options({'COMPRESSION_CODEC': codec})
+
     def exists_table(self, name, database=None):
         """
         Determine if the indicated table or view exists


Test patch:
diff --git a/ibis/impala/tests/common.py b/ibis/impala/tests/common.py
--- a/ibis/impala/tests/common.py
+++ b/ibis/impala/tests/common.py
@@ -66,6 +66,9 @@ def __repr__(self):
         return 'IbisTestEnv(\n    {0})'.format(',\n    '.join(kvs))
 
 
+ENV = IbisTestEnv()
+
+
 def connect_test(env, with_hdfs=True):
     if with_hdfs:
         if env.auth_mechanism in ['GSSAPI', 'LDAP']:
@@ -93,13 +96,16 @@ class ImpalaE2E(object):
     def setUpClass(cls):
         ImpalaE2E.setup_e2e(cls)
 
+        # make sure this never gets messed up
+        opts = cls.con.get_options()
+        assert opts['DISABLE_CODEGEN'] == '1'
+
     @classmethod
     def tearDownClass(cls):
         ImpalaE2E.teardown_e2e(cls)
 
     @staticmethod
     def setup_e2e(cls):
-        ENV = IbisTestEnv()
         cls.con = connect_test(ENV)
         # Tests run generally faster without it
         if not ENV.use_codegen:
diff --git a/ibis/impala/tests/test_client.py b/ibis/impala/tests/test_client.py
--- a/ibis/impala/tests/test_client.py
+++ b/ibis/impala/tests/test_client.py
@@ -300,3 +300,37 @@ def test_query_cancel(self):
         assert elapsed < 5
 
         assert q.is_finished()
+
+    def test_set_compression_codec(self):
+        old_opts = self.con.get_options()
+        assert old_opts['COMPRESSION_CODEC'].upper() == 'NONE'
+
+        self.con.set_compression_codec('snappy')
+        opts = self.con.get_options()
+        assert opts['COMPRESSION_CODEC'].upper() == 'SNAPPY'
+
+        self.con.set_compression_codec(None)
+        opts = self.con.get_options()
+        assert opts['COMPRESSION_CODEC'].upper() == 'NONE'
+
+    def test_disable_codegen(self):
+        self.con.disable_codegen(False)
+        opts = self.con.get_options()
+        assert opts['DISABLE_CODEGEN'] == '0'
+
+        self.con.disable_codegen()
+        opts = self.con.get_options()
+        assert opts['DISABLE_CODEGEN'] == '1'
+
+        impala_con = self.con.con
+        cur1 = impala_con.execute('SET')
+        cur2 = impala_con.execute('SET')
+
+        opts1 = dict(cur1.fetchall())
+        cur1.release()
+
+        opts2 = dict(cur2.fetchall())
+        cur2.release()
+
+        assert opts1['DISABLE_CODEGEN'] == '1'
+        assert opts2['DISABLE_CODEGEN'] == '1'