Merge pull request #23 from deepfield/PB-266-create_with_partition

PB-266 Master CTAS with partition

ibis/impala/client.py

@@ -855,13 +855,6 @@ def create_table(self, table_name, obj=None, schema=None, database=None,
             ast = self._build_ast(to_insert)
             select = ast.queries[0]
 
-            if partition is not None:
-                # Fairly certain this is currently the case
-                raise ValueError('partition not supported with '
-                                 'create-table-as-select. Create an '
-                                 'empty partitioned table instead '
-                                 'and insert into those partitions.')
-
             statement = ddl.CTAS(table_name, select,
                                  database=database,
                                  can_exist=force,

ibis/impala/ddl.py

@@ -142,6 +142,13 @@ def _storage(self):
         }
         return storage_lines[self.format]
 
+    def _partition(self):
+        if self.partition:
+            return "\nPARTITIONED BY ({})".format(", ".join(self.partition))
+        else:
+            return ''
+
+
 
 class CTAS(CreateTable):
 
@@ -160,6 +167,7 @@ def __init__(self, table_name, select, database=None,
     def compile(self):
         buf = StringIO()
         buf.write(self._create_line())
+        buf.write(self._partition())
         buf.write(self._storage())
         buf.write(self._location())
 

ibis/impala/tests/test_ddl.py

@@ -291,6 +291,22 @@ def test_create_external_table_as(self):
 FROM test1""".format(path)
         assert result == expected
 
+    def test_create_table_as_with_partitions(self):
+        partition = ('year', 'month', 'day')
+        select = build_ast(self.con.table('test1')).queries[0]
+        statement = ddl.CTAS('db.test_table',
+                             select,
+                             partition=partition)
+        result = statement.compile()
+
+        expected = """\
+                CREATE TABLE db.test_table
+        PARTIONED BY (year, month, day) AS SELECT *
+        FROM test1"""
+        assert result == expected
+
+
+
     def test_create_table_with_location(self):
         path = '/path/to/table'
         schema = ibis.schema([('foo', 'string'),