HiveServer2Hook: Support kwargs in get_pandas_df

apache · Aug 5, 2020 · fa7192e · fa7192e
1 parent 090f71a
commit fa7192e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/airflow/providers/apache/hive/hooks/hive.py b/airflow/providers/apache/hive/hooks/hive.py
@@ -1044,9 +1044,10 @@ def get_records(self, hql: Union[str, Text],
         """
         return self.get_results(hql, schema=schema, hive_conf=hive_conf)['data']
 
-    def get_pandas_df(self, hql: Union[str, Text],
+    def get_pandas_df(self, hql: Union[str, Text],  # type: ignore
                       schema: str = 'default',
-                      hive_conf: Optional[Dict[Any, Any]] = None
+                      hive_conf: Optional[Dict[Any, Any]] = None,
+                      **kwargs
                       ) -> pandas.DataFrame:
         """
         Get a pandas dataframe from a Hive query
@@ -1069,6 +1070,6 @@ def get_pandas_df(self, hql: Union[str, Text],
         :return: pandas.DateFrame
         """
         res = self.get_results(hql, schema=schema, hive_conf=hive_conf)
-        df = pandas.DataFrame(res['data'])
+        df = pandas.DataFrame(res['data'], **kwargs)
         df.columns = [c[0] for c in res['header']]
         return df