From 0413cf1eb0c66c5f300960b3a167b149170b09f7 Mon Sep 17 00:00:00 2001 From: Randy Zwitch Date: Fri, 22 Mar 2019 08:43:30 -0400 Subject: [PATCH] Fix array handling (#183) Provides ability to return list columns within sql_execute. Closes #68 --- pymapd/_parsers.py | 59 +++++++++++++++++------ tests/test_integration.py | 98 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 143 insertions(+), 14 deletions(-) diff --git a/pymapd/_parsers.py b/pymapd/_parsers.py index cc89a69..dc2f133 100644 --- a/pymapd/_parsers.py +++ b/pymapd/_parsers.py @@ -48,25 +48,56 @@ _thrift_values_to_encodings = T.TEncodingType._VALUES_TO_NAMES +def _format_result_timestamp(desc, arr): + + return [None if v is None else + datetime_in_precisions(v, desc.col_type.precision) + for v in arr] + + +def _format_result_date(arr): + + base = datetime.datetime(1970, 1, 1) + return [None if v is None else + (base + datetime.timedelta(seconds=v)).date() + for v in arr] + + +def _format_result_time(arr): + + return [None if v is None else seconds_to_time(v) for v in arr] + + def _extract_col_vals(desc, val): # type: (T.TColumnType, T.TColumn) -> Any + typename = T.TDatumType._VALUES_TO_NAMES[desc.col_type.type] nulls = val.nulls - vals = getattr(val.data, _typeattr[typename] + '_col') - vals = [None if null else v - for null, v in zip(nulls, vals)] - - if typename == 'TIMESTAMP': - vals = [None if v is None else - datetime_in_precisions(v, desc.col_type.precision) - for v in vals] - elif typename == 'DATE': - base = datetime.datetime(1970, 1, 1) - vals = [None if v is None else - (base + datetime.timedelta(seconds=v)).date() for v in vals] - elif typename == 'TIME': - vals = [None if v is None else seconds_to_time(v) for v in vals] + # arr_col has multiple levels to parse, not accounted for in original code + # https://github.com/omnisci/pymapd/issues/68 + if hasattr(val.data, 'arr_col') and val.data.arr_col: + vals = [None if null else getattr(v.data, _typeattr[typename] + '_col') + for null, v in zip(nulls, val.data.arr_col)] + + if typename == 'TIMESTAMP': + vals = [_format_result_timestamp(desc, v) for v in vals] + elif typename == 'DATE': + vals = [_format_result_date(v) for v in vals] + elif typename == 'TIME': + vals = [_format_result_time(v) for v in vals] + + # else clause original code path + else: + vals = getattr(val.data, _typeattr[typename] + '_col') + vals = [None if null else v for null, v in zip(nulls, vals)] + + if typename == 'TIMESTAMP': + vals = _format_result_timestamp(desc, vals) + elif typename == 'DATE': + vals = _format_result_date(vals) + elif typename == 'TIME': + vals = _format_result_time(vals) return vals diff --git a/tests/test_integration.py b/tests/test_integration.py index d0bf0fd..9bbe1d8 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -408,3 +408,101 @@ def test_load_table_creates(self, con, not_a_table): 'double_', 'varchar_', 'text_', 'time_', 'timestamp_', 'date_']) con.load_table(not_a_table, data, create=True) + + def test_array_in_result_set(self, con): + + # text + con.execute("DROP TABLE IF EXISTS test_lists;") + con.execute("CREATE TABLE IF NOT EXISTS test_lists \ + (col1 TEXT, col2 TEXT[]);") + + row = [("row1", "{hello,goodbye,aloha}"), + ("row2", "{hello2,goodbye2,aloha2}")] + + con.load_table_rowwise("test_lists", row) + ans = con.execute("select * from test_lists").fetchall() + + expected = [('row1', ['hello', 'goodbye', 'aloha']), + ('row2', ['hello2', 'goodbye2', 'aloha2'])] + + assert ans == expected + + # int + con.execute("DROP TABLE IF EXISTS test_lists;") + con.execute("CREATE TABLE IF NOT EXISTS test_lists \ + (col1 TEXT, col2 INT[]);") + + row = [("row1", "{10,20,30}"), ("row2", "{40,50,60}")] + + con.load_table_rowwise("test_lists", row) + ans = con.execute("select * from test_lists").fetchall() + + expected = [('row1', [10, 20, 30]), ('row2', [40, 50, 60])] + + assert ans == expected + + # timestamp + con.execute("DROP TABLE IF EXISTS test_lists;") + con.execute("CREATE TABLE IF NOT EXISTS test_lists \ + (col1 TEXT, col2 TIMESTAMP[]);") + + row = [("row1", + "{2019-03-02 00:00:00,2019-03-02 00:00:00,2019-03-02 00:00:00}"), # noqa + ("row2", + "{2019-03-02 00:00:00,2019-03-02 00:00:00,2019-03-02 00:00:00}")] # noqa + + con.load_table_rowwise("test_lists", row) + ans = con.execute("select * from test_lists").fetchall() + + expected = [('row1', + [datetime.datetime(2019, 3, 2, 0, 0), + datetime.datetime(2019, 3, 2, 0, 0), + datetime.datetime(2019, 3, 2, 0, 0)]), + ('row2', + [datetime.datetime(2019, 3, 2, 0, 0), + datetime.datetime(2019, 3, 2, 0, 0), + datetime.datetime(2019, 3, 2, 0, 0)])] + + assert ans == expected + + # date + con.execute("DROP TABLE IF EXISTS test_lists;") + con.execute("CREATE TABLE IF NOT EXISTS test_lists \ + (col1 TEXT, col2 DATE[]);") + + row = [("row1", "{2019-03-02,2019-03-02,2019-03-02}"), + ("row2", "{2019-03-02,2019-03-02,2019-03-02}")] + + con.load_table_rowwise("test_lists", row) + ans = con.execute("select * from test_lists").fetchall() + + expected = [('row1', + [datetime.date(2019, 3, 2), + datetime.date(2019, 3, 2), + datetime.date(2019, 3, 2)]), + ('row2', + [datetime.date(2019, 3, 2), + datetime.date(2019, 3, 2), + datetime.date(2019, 3, 2)])] + + assert ans == expected + + # time + con.execute("DROP TABLE IF EXISTS test_lists;") + con.execute("CREATE TABLE IF NOT EXISTS test_lists \ + (col1 TEXT, col2 TIME[]);") + + row = [("row1", "{23:59,23:59,23:59}"), + ("row2", "{23:59,23:59,23:59}")] + + con.load_table_rowwise("test_lists", row) + ans = con.execute("select * from test_lists").fetchall() + + expected = [('row1', + [datetime.time(23, 59), datetime.time(23, 59), + datetime.time(23, 59)]), + ('row2', + [datetime.time(23, 59), datetime.time(23, 59), + datetime.time(23, 59)])] + + assert ans == expected