Browse Source

HUE-3389 [editor] Provide select from partitions support in get_sample for Impala

Jenny Kim 9 years ago
parent
commit
abfab56f30

+ 7 - 9
apps/beeswax/src/beeswax/server/dbms.py

@@ -315,7 +315,11 @@ class HiveServer2Dbms(object):
     result = None
     hql = None
 
-    if self.server_name == 'impala':
+    # Filter on max # of partitions for partitioned tables
+    column = '`%s`' % column if column else '*'
+    if table.partition_keys:
+      hql = self._get_sample_partition_query(database, table, column, limit)
+    elif self.server_name == 'impala':
       if column or nested:
         from impala.dbms import ImpalaDbms
         select_clause, from_clause = ImpalaDbms.get_nested_select(database, table.name, column, nested)
@@ -323,14 +327,8 @@ class HiveServer2Dbms(object):
       else:
         hql = "SELECT * FROM `%s`.`%s` LIMIT %s" % (database, table.name, limit)
     else:
-      # Filter on max # of partitions for partitioned tables
-      # Impala's SHOW PARTITIONS is different from Hive, so we only support Hive for now
-      column = '`%s`' % column if column else '*'
-      if table.partition_keys:
-        hql = self._get_sample_partition_query(database, table, column, limit)
-      else:
-        hql = "SELECT %s FROM `%s`.`%s` LIMIT %s" % (column, database, table.name, limit)
-        # TODO: Add nested select support for HS2
+      hql = "SELECT %s FROM `%s`.`%s` LIMIT %s" % (column, database, table.name, limit)
+      # TODO: Add nested select support for HS2
 
     if hql:
       query = hql_query(hql)

+ 30 - 5
apps/beeswax/src/beeswax/server/hive_server2_lib.py

@@ -20,7 +20,7 @@ import itertools
 import json
 import re
 
-from itertools import imap
+from itertools import imap, izip
 from operator import itemgetter
 
 from django.utils.translation import ugettext as _
@@ -863,14 +863,39 @@ class HiveServerClient:
     table = self.get_table(database, table_name)
 
     query = 'SHOW PARTITIONS `%s`.`%s`' % (database, table_name)
-    if partition_spec:
+    if self.query_server['server_name'] == 'beeswax' and partition_spec:
       query += ' PARTITION(%s)' % partition_spec
 
     # We fetch N partitions then reverse the order later and get the max_parts. Use partition_spec to refine more the initial list.
     # Need to fetch more like this until SHOW PARTITIONS offers a LIMIT and ORDER BY
-    partition_table = self.execute_query_statement(query, max_rows=10000)
+    partition_table = self.execute_query_statement(query, max_rows=10000, orientation=TFetchOrientation.FETCH_NEXT)
 
-    partitions = [PartitionValueCompatible(partition, table) for partition in partition_table.rows()]
+    if self.query_server['server_name'] == 'impala':
+      try:
+        # Fetch all partition key names, which are listed before the #Rows column
+        cols = [col.name for col in partition_table.cols()]
+        stop = cols.index('#Rows')
+        partition_keys = cols[:stop]
+        num_parts = len(partition_keys)
+
+        # Get all partition values
+        rows = partition_table.rows()
+        partition_values = [partition[:num_parts] for partition in rows]
+
+        # Truncate last row which is the Total
+        partition_values = partition_values[:-1]
+        partitions_formatted = []
+
+        # Format partition key and values into Hive format: [key1=val1/key2=value2]
+        for values in partition_values:
+          zipped_parts = izip(partition_keys, values)
+          partitions_formatted.append(['/'.join(['%s=%s' % (part[0], part[1]) for part in zipped_parts])])
+
+        partitions = [PartitionValueCompatible(partition, table) for partition in partitions_formatted]
+      except Exception, e:
+        raise ValueError(_('Failed to determine partition keys for Impala table: `%s`.`%s`') % (database, table_name))
+    else:
+      partitions = [PartitionValueCompatible(partition, table) for partition in partition_table.rows()]
 
     if reverse_sort:
       partitions.reverse()
@@ -967,7 +992,7 @@ class PartitionValueCompatible:
     # Parses: ['datehour=2013022516'] or ['month=2011-07/dt=2011-07-01/hr=12']
     partition = partition_row[0]
     parts = partition.split('/')
-    self.partition_spec = ','.join(["%s='%s'" % (pv[0], pv[1]) for pv in [part.split('=') for part in parts]])
+    self.partition_spec = ','.join(["`%s`='%s'" % (pv[0], pv[1]) for pv in [part.split('=') for part in parts]])
     self.values = [pv[1] for pv in [part.split('=') for part in parts]]
     self.sd = type('Sd', (object,), properties,)
 

+ 2 - 2
apps/beeswax/src/beeswax/tests.py

@@ -1733,7 +1733,7 @@ for x in sys.stdin:
     finish = conf.QUERY_PARTITIONS_LIMIT.set_for_testing(1)
     try:
       table_name = 'test_partitions'
-      partition_spec = "(baz='baz_one' AND boom='boom_two')"
+      partition_spec = "(`baz`='baz_one' AND `boom`='boom_two')"
       table = self.db.get_table(database=self.db_name, table_name=table_name)
       hql = self.db._get_sample_partition_query(self.db_name, table, limit=10)
       assert_equal(hql, 'SELECT * FROM `%s`.`%s` WHERE %s LIMIT 10' % (self.db_name, table_name, partition_spec))
@@ -1744,7 +1744,7 @@ for x in sys.stdin:
     finish = conf.QUERY_PARTITIONS_LIMIT.set_for_testing(2)
     try:
       table_name = 'test_partitions'
-      partition_spec = "(baz='baz_one' AND boom='boom_two') OR (baz='baz_foo' AND boom='boom_bar')"
+      partition_spec = "(`baz`='baz_one' AND `boom`='boom_two') OR (`baz`='baz_foo' AND `boom`='boom_bar')"
       table = self.db.get_table(database=self.db_name, table_name=table_name)
       hql = self.db._get_sample_partition_query(self.db_name, table, limit=10)
       assert_equal(hql, 'SELECT * FROM `%s`.`%s` WHERE %s LIMIT 10' % (self.db_name, table_name, partition_spec))