Tests for Column of Arrays (heavydb-internal PR 6571).

rbc/extension_functions.thrift

@@ -53,6 +53,20 @@ enum TExtArgumentType {
   ColumnListTextEncodingDict,
   ColumnTimestamp,
   Timestamp,
+  ColumnArrayInt8,
+  ColumnArrayInt16,
+  ColumnArrayInt32,
+  ColumnArrayInt64,
+  ColumnArrayFloat,
+  ColumnArrayDouble,
+  ColumnArrayBool,
+  ColumnListArrayInt8,
+  ColumnListArrayInt16,
+  ColumnListArrayInt32,
+  ColumnListArrayInt64,
+  ColumnListArrayFloat,
+  ColumnListArrayDouble,
+  ColumnListArrayBool,
 }
 
 /* See QueryEngine/TableFunctions/TableFunctionsFactory.h for required

rbc/heavydb/remoteheavydb.py

@@ -569,7 +569,9 @@ def thrift_call(self, name, *args, **kwargs):
                 msg = (f"{msg.error_msg} Please use server options --enable-runtime-udf"
                        " and/or --enable-table-functions")
                 raise HeavyDBServerError(msg)
-
+            m = re.match(r'Error executing table function', msg.error_msg)
+            if m:
+                raise HeavyDBServerError(msg.error_msg)
             # TODO: catch more known server failures here.
             raise
 
@@ -859,13 +861,6 @@ def _get_ext_arguments_map(self):
             'Array<int64_t>': typemap['TExtArgumentType']['ArrayInt64'],
             'Array<float>': typemap['TExtArgumentType']['ArrayFloat'],
             'Array<double>': typemap['TExtArgumentType']['ArrayDouble'],
-            'Column<bool>': typemap['TExtArgumentType'].get('ColumnBool'),
-            'Column<int8_t>': typemap['TExtArgumentType'].get('ColumnInt8'),
-            'Column<int16_t>': typemap['TExtArgumentType'].get('ColumnInt16'),
-            'Column<int32_t>': typemap['TExtArgumentType'].get('ColumnInt32'),
-            'Column<int64_t>': typemap['TExtArgumentType'].get('ColumnInt64'),
-            'Column<float>': typemap['TExtArgumentType'].get('ColumnFloat'),
-            'Column<double>': typemap['TExtArgumentType'].get('ColumnDouble'),
             'Column<TextEncodingDict>': typemap['TExtArgumentType'].get(
                 'ColumnTextEncodingDict'),
             'Cursor': typemap['TExtArgumentType']['Cursor'],
@@ -877,13 +872,6 @@ def _get_ext_arguments_map(self):
                 'GeoMultiPolygon'),
             'TextEncodingNone': typemap['TExtArgumentType'].get('TextEncodingNone'),
             'TextEncodingDict': typemap['TExtArgumentType'].get('TextEncodingDict'),
-            'ColumnList<bool>': typemap['TExtArgumentType'].get('ColumnListBool'),
-            'ColumnList<int8_t>': typemap['TExtArgumentType'].get('ColumnListInt8'),
-            'ColumnList<int16_t>': typemap['TExtArgumentType'].get('ColumnListInt16'),
-            'ColumnList<int32_t>': typemap['TExtArgumentType'].get('ColumnListInt32'),
-            'ColumnList<int64_t>': typemap['TExtArgumentType'].get('ColumnListInt64'),
-            'ColumnList<float>': typemap['TExtArgumentType'].get('ColumnListFloat'),
-            'ColumnList<double>': typemap['TExtArgumentType'].get('ColumnListDouble'),
             'ColumnList<TextEncodingDict>': typemap['TExtArgumentType'].get(
                 'ColumnListTextEncodingDict'),
             'Timestamp': typemap['TExtArgumentType'].get('Timestamp'),
@@ -896,6 +884,22 @@ def _get_ext_arguments_map(self):
 
         ext_arguments_map['bool8'] = ext_arguments_map['bool']
 
+        for T, Tname in [
+                ('bool', 'Bool'),
+                ('int8_t', 'Int8'),
+                ('int16_t', 'Int16'),
+                ('int32_t', 'Int32'),
+                ('int64_t', 'Int64'),
+                ('float', 'Float'),
+                ('double', 'Double')]:
+            ext_arguments_map[f'Column<{T}>'] = typemap['TExtArgumentType'].get(f'Column{Tname}')
+            ext_arguments_map[f'ColumnList<{T}>'] = typemap['TExtArgumentType'].get(
+                f'ColumnList{Tname}')
+            ext_arguments_map[f'ColumnArray<{T}>'] = typemap['TExtArgumentType'].get(
+                f'ColumnArray{Tname}')
+            ext_arguments_map[f'ColumnListArray<{T}>'] = typemap['TExtArgumentType'].get(
+                f'ColumnListArray{Tname}')
+
         for ptr_type, T in [
                 ('bool', 'bool'),
                 ('bool8', 'bool'),
@@ -936,6 +940,9 @@ def type_to_extarg(self, t):
             s = t.tostring(use_annotation=False, use_name=False)
             extarg = self._get_ext_arguments_map().get(s)
             if extarg is None:
+                if s in self._get_ext_arguments_map():
+                    raise ValueError(f'cannot convert {t}(={s}) to ExtArgumentType:'
+                                     ' bug in _get_ext_arguments_map?')
                 raise ValueError(f'cannot convert {t}(={s}) to ExtArgumentType')
             return extarg
         elif isinstance(t, str):
@@ -971,7 +978,8 @@ def retrieve_targets(self):
                     messages.append(f'thrift type map: add new type {typ}')
                 elif member not in typemap[typ]:
                     messages.append(
-                        f'thrift type map: add new member {typ}.{member}')
+                        f'thrift type map: add new member {typ}.{member}'
+                        ' to rbc/extension_functions.thrift')
                 elif ivalue != typemap[typ][member]:
                     messages.append(
                         f'thrift type map: update {typ}.{member}'

rbc/tests/__init__.py

@@ -117,7 +117,7 @@ def require_version(version, message=None, label=None):
         if not available_version:
             pytest.skip(reason)
         # Requires update when heavydb-internal bumps up version number:
-        current_development_version = (6, 1, 0)
+        current_development_version = (6, 2, 0)
         if available_version[:3] > current_development_version:
             warnings.warn(f'{available_version}) is newer than development version'
                           f' ({current_development_version}), please update the latter!')

rbc/tests/heavydb/test_column_array.py

@@ -0,0 +1,55 @@
+import os
+import pytest
+from rbc.tests import heavydb_fixture
+
+
+@pytest.fixture(scope="module")
+def heavydb():
+
+    for o in heavydb_fixture(globals(), minimal_version=(6, 1)):
+        define(o)
+        o.base_name = os.path.splitext(os.path.basename(__file__))[0]
+        yield o
+
+
+def define(heavydb):
+    pass
+
+
+@pytest.mark.parametrize("colname", ['f8', 'f4', 'i1', 'i2', 'i4', 'i8', 'b'])
+def test_sum_along_row(heavydb, colname):
+
+    _, result = heavydb.sql_execute(f'SELECT {colname} FROM {heavydb.table_name}array')
+    mysum = any if colname == 'b' else sum
+    expected_result = [(mysum(item[0]),) for item in result]
+
+    _, result = heavydb.sql_execute('SELECT * FROM table(sum_along_row(CURSOR('
+                                    f'SELECT {colname} FROM {heavydb.table_name}array)))')
+    result = list(result)
+    assert result == expected_result
+
+
+@pytest.mark.parametrize("colname", ['f8', 'f4', 'i1', 'i2', 'i4', 'i8', 'b'])
+def test_array_copier(heavydb, colname):
+
+    _, result = heavydb.sql_execute(f'SELECT {colname} FROM {heavydb.table_name}array')
+    expected_result = list(result)
+
+    _, result = heavydb.sql_execute('SELECT * FROM table(array_copier(CURSOR('
+                                    f'SELECT {colname} FROM {heavydb.table_name}array)))')
+    result = list(result)
+    assert result == expected_result
+
+
+@pytest.mark.parametrize("colname", ['f8', 'f4', 'i1', 'i2', 'i4', 'i8', 'b'])
+def test_array_concat(heavydb, colname):
+
+    for n in [1, 3]:
+        _, result = heavydb.sql_execute(f'SELECT {colname} FROM {heavydb.table_name}array')
+        expected_result = [(item[0] * n,) for item in result]
+
+        colnames = ", ".join([colname] * n)
+        _, result = heavydb.sql_execute('SELECT * FROM table(array_concat(CURSOR('
+                                        f'SELECT {colnames} FROM {heavydb.table_name}array)))')
+        result = list(result)
+        assert result == expected_result

rbc/tests/heavydb/test_mlpack.py

@@ -7,26 +7,24 @@
 @pytest.fixture(scope='module')
 def heavydb():
 
-    for o in heavydb_fixture(globals(), minimal_version=(5, 6)):
+    for o in heavydb_fixture(globals(), minimal_version=(6, 1)):
         yield o
 
 
 @pytest.mark.parametrize("func", ['dbscan', 'kmeans'])
 def test_mlpack(heavydb, func):
     heavydb.require_version(
-        (5, 6),
+        (6, 1),
         'Requires heavydb-internal PR 5430 and heavydb built with -DENABLE_MLPACK')
 
-    extra_args = dict(dbscan='cast(1 as float), 1',
-                      kmeans='1')[func]
+    extra_args = dict(dbscan='cast(1 as float), 1, \'mlpack\'',
+                      kmeans='1, 1, \'default\', \'mlpack\'')[func]
     query = (f'select * from table({func}(cursor(select cast(rowid as int), f8, f8, f8 '
-             f'from {heavydb.table_name}), {extra_args}, 1))')
-
+             f'from {heavydb.table_name}), {extra_args}))')
     try:
         _, result = heavydb.sql_execute(query)
     except HeavyDBServerError as msg:
-        m = re.match(fr'.*Undefined function call {func!r}',
-                     msg.args[0])
+        m = re.match(r'.*Cannot find mlpack ML library to support', msg.args[0])
         if m is not None:
             pytest.skip(f'test requires heavydb server with MLPACK support: {msg}')
         raise

rbc/tests/heavydb/test_udtf.py

@@ -1,5 +1,5 @@
 import pytest
-from rbc.errors import NumbaTypeError
+from rbc.errors import NumbaTypeError, HeavyDBServerError
 from rbc.tests import heavydb_fixture, sql_execute
 from rbc.externals.heavydb import table_function_error
 import numpy as np
@@ -145,7 +145,7 @@ def my_manager_row_size(mgr, col, out):
             out[i] = col[i]
         return size
 
-    with pytest.raises(heavydb.thrift_client.thrift.TMapDException) as exc:
+    with pytest.raises(HeavyDBServerError) as exc:
         heavydb.sql_execute(
             f'select out0 from table(my_manager_error('
             f'cursor(select f8 from {heavydb.table_name})));')
@@ -250,7 +250,7 @@ def my_divide(column, k, row_multiplier, out):
     """)
     assert list(result) == [(0.0,), (0.5,), (1.0,), (1.5,), (2.0,)]
 
-    with pytest.raises(heavydb.thrift_client.thrift.TMapDException) as exc:
+    with pytest.raises(HeavyDBServerError) as exc:
         heavydb.sql_execute(f"""
             select *
             from table(