Merge pull request #2 from airbnb/hive_stats_resolution

Bring in more resolution to hivestats
apache · Apr 20, 2015 · 765446b · 765446b
2 parents 7eb484a + 1954bea
commit 765446b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/airflow/hooks/mysql_hook.py b/airflow/hooks/mysql_hook.py
@@ -56,11 +56,14 @@ def run(self, sql):
         cur.close()
         conn.close()
 
-    def insert_rows(self, table, rows):
+    def insert_rows(self, table, rows, target_fields=None):
         """
         A generic way to insert a set of tuples into a table,
         the whole set of inserts is treated as one transaction
         """
+        if target_fields:
+            target_fields = ", ".join(target_fields)
+            target_fields = "({})".format(target_fields)
         conn = self.get_conn()
         cur = conn.cursor()
         for row in rows:
@@ -73,8 +76,10 @@ def insert_rows(self, table, rows):
                 else:
                     l.append(str(cell))
             values = tuple(l)
-            sql = "INSERT INTO {0} VALUES ({1});".format(
-                table, ",".join(values))
+            sql = "INSERT INTO {0} {1} VALUES ({2});".format(
+                table,
+                target_fields,
+                ",".join(values))
             cur.execute(sql)
         conn.commit()
         cur.close()

diff --git a/airflow/operators/hive_stats_operator.py b/airflow/operators/hive_stats_operator.py
@@ -7,7 +7,6 @@
 from airflow.utils import apply_defaults
 
 
-
 class HiveStatsCollectionOperator(BaseOperator):
     """
     Gathers partition statistics using a dynmically generated Presto
@@ -35,7 +34,9 @@ class HiveStatsCollectionOperator(BaseOperator):
     :type col_blacklist: list
     :param assignment_func: a function that receives a column name and
         a type, and returns a dict of metric names and an Presto expressions.
-        If None is returned, the global defaults are applied.
+        If None is returned, the global defaults are applied. If an
+        empty dictionary is returned, no stats are computed for that
+        column.
     :type assignment_func: function
     """
 
@@ -68,6 +69,7 @@ def __init__(
         self.mysql_conn_id = mysql_conn_id
         self.assignment_func = assignment_func
         self.ds = '{{ ds }}'
+        self.dttm = '{{ execution_date.isoformat() }}'
 
     def get_default_exprs(self, col, col_type):
         if col in self.col_blacklist:
@@ -100,7 +102,9 @@ def execute(self, context=None):
             d = {}
             if self.assignment_func:
                 d = self.assignment_func(col, col_type)
-            if not d:
+                if d is None:
+                    d = self.get_default_exprs(col, col_type)
+            else:
                 d = self.get_default_exprs(col, col_type)
             exprs.update(d)
         exprs.update(self.extra_exprs)
@@ -142,6 +146,19 @@ def execute(self, context=None):
 
         logging.info("Pivoting and loading cells into the Airflow db")
         rows = [
-            (self.ds, self.table, part_json) + (r[0][0], r[0][1], r[1])
+            (self.ds, self.dttm, self.table, part_json) +
+            (r[0][0], r[0][1], r[1])
             for r in zip(exprs, row)]
-        mysql.insert_rows(table='hive_stats', rows=rows)
+        mysql.insert_rows(
+            table='hive_stats',
+            rows=rows,
+            target_fields=[
+                'ds',
+                'dttm',
+                'table_name',
+                'partition_repr',
+                'col',
+                'metric',
+                'value',
+            ]
+        )