You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2018/05/19 16:01:29 UTC
incubator-airflow git commit: [AIRFLOW-2448] Enhance
HiveCliHook.load_df to work with datetime
Repository: incubator-airflow
Updated Branches:
refs/heads/master e48b8e36a -> 67b351183
[AIRFLOW-2448] Enhance HiveCliHook.load_df to work with datetime
HiveCliHook.load_df can not handle DataFrame
which contains datetime for now.
This PR enhances it to work with datetime,
fixes some bug introduced by AIRFLOW-2441,
and addresses some flake8 issues.
Closes #3364 from sekikn/AIRFLOW-2448
Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/67b35118
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/67b35118
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/67b35118
Branch: refs/heads/master
Commit: 67b351183b0f85e9484f1f7f70e0b46300753b60
Parents: e48b8e3
Author: Kengo Seki <se...@apache.org>
Authored: Sat May 19 16:59:44 2018 +0100
Committer: Kaxil Naik <ka...@apache.org>
Committed: Sat May 19 17:00:47 2018 +0100
----------------------------------------------------------------------
airflow/hooks/hive_hooks.py | 30 +++++++++++++++----------
tests/hooks/test_hive_hook.py | 45 +++++++++++++++++++++++++++++++++++++-
2 files changed, 62 insertions(+), 13 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/67b35118/airflow/hooks/hive_hooks.py
----------------------------------------------------------------------
diff --git a/airflow/hooks/hive_hooks.py b/airflow/hooks/hive_hooks.py
index 92958c0..01a6a2a 100644
--- a/airflow/hooks/hive_hooks.py
+++ b/airflow/hooks/hive_hooks.py
@@ -20,11 +20,12 @@
from __future__ import print_function, unicode_literals
from six.moves import zip
-from past.builtins import basestring
+from past.builtins import basestring, unicode
import unicodecsv as csv
import itertools
import re
+import six
import subprocess
import time
from collections import OrderedDict
@@ -316,15 +317,16 @@ class HiveCliHook(BaseHook):
def _infer_field_types_from_df(df):
DTYPE_KIND_HIVE_TYPE = {
- 'b': 'BOOLEAN', # boolean
- 'i': 'BIGINT', # signed integer
- 'u': 'BIGINT', # unsigned integer
- 'f': 'DOUBLE', # floating-point
- 'c': 'STRING', # complex floating-point
- 'O': 'STRING', # object
- 'S': 'STRING', # (byte-)string
- 'U': 'STRING', # Unicode
- 'V': 'STRING' # void
+ 'b': 'BOOLEAN', # boolean
+ 'i': 'BIGINT', # signed integer
+ 'u': 'BIGINT', # unsigned integer
+ 'f': 'DOUBLE', # floating-point
+ 'c': 'STRING', # complex floating-point
+ 'M': 'TIMESTAMP', # datetime
+ 'O': 'STRING', # object
+ 'S': 'STRING', # (byte-)string
+ 'U': 'STRING', # Unicode
+ 'V': 'STRING' # void
}
d = OrderedDict()
@@ -336,15 +338,19 @@ class HiveCliHook(BaseHook):
pandas_kwargs = {}
with TemporaryDirectory(prefix='airflow_hiveop_') as tmp_dir:
- with NamedTemporaryFile(dir=tmp_dir) as f:
+ with NamedTemporaryFile(dir=tmp_dir, mode="w") as f:
if field_dict is None and (create or recreate):
field_dict = _infer_field_types_from_df(df)
df.to_csv(path_or_buf=f,
- sep=delimiter.encode(encoding),
+ sep=(delimiter.encode(encoding)
+ if six.PY2 and isinstance(delimiter, unicode)
+ else delimiter),
header=False,
index=False,
+ encoding=encoding,
+ date_format="%Y-%m-%d %H:%M:%S",
**pandas_kwargs)
f.flush()
http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/67b35118/tests/hooks/test_hive_hook.py
----------------------------------------------------------------------
diff --git a/tests/hooks/test_hive_hook.py b/tests/hooks/test_hive_hook.py
index 3c278e3..eadb0f1 100644
--- a/tests/hooks/test_hive_hook.py
+++ b/tests/hooks/test_hive_hook.py
@@ -21,6 +21,7 @@
import datetime
import pandas as pd
import random
+import re
import mock
import unittest
@@ -125,7 +126,7 @@ class TestHiveCliHook(unittest.TestCase):
kwargs = mock_to_csv.call_args[1]
self.assertEqual(kwargs["header"], False)
self.assertEqual(kwargs["index"], False)
- self.assertEqual(kwargs["sep"], delimiter.encode(encoding))
+ self.assertEqual(kwargs["sep"], delimiter)
mock_load_file.assert_called_once()
kwargs = mock_load_file.call_args[1]
@@ -134,6 +135,48 @@ class TestHiveCliHook(unittest.TestCase):
self.assertTrue(isinstance(kwargs["field_dict"], OrderedDict))
self.assertEqual(kwargs["table"], table)
+ @mock.patch('airflow.hooks.hive_hooks.HiveCliHook.run_cli')
+ def test_load_df_with_data_types(self, mock_run_cli):
+ d = OrderedDict()
+ d['b'] = [True]
+ d['i'] = [-1]
+ d['t'] = [1]
+ d['f'] = [0.0]
+ d['c'] = ['c']
+ d['M'] = [datetime.datetime(2018, 1, 1)]
+ d['O'] = [object()]
+ d['S'] = ['STRING'.encode('utf-8')]
+ d['U'] = ['STRING']
+ d['V'] = [None]
+ df = pd.DataFrame(d)
+
+ hook = HiveCliHook()
+ hook.load_df(df, 't')
+
+ query = """
+ CREATE TABLE IF NOT EXISTS t (
+ b BOOLEAN,
+ i BIGINT,
+ t BIGINT,
+ f DOUBLE,
+ c STRING,
+ M TIMESTAMP,
+ O STRING,
+ S STRING,
+ U STRING,
+ V STRING)
+ ROW FORMAT DELIMITED
+ FIELDS TERMINATED BY ','
+ STORED AS textfile
+ ;
+ """
+
+ def _trim(s):
+ return re.sub("\s+", " ", s.strip())
+
+ self.assertEqual(_trim(mock_run_cli.call_args_list[0][0][0]),
+ _trim(query))
+
class TestHiveMetastoreHook(HiveEnvironmentTest):
VALID_FILTER_MAP = {'key2': 'value2'}