Fix array handling (#183)

Provides ability to return list columns within sql_execute. Closes #68

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
 

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