Lay out a returned SQL array as the catalog's arrayShape states on the typed SQL surfaces

codegen/src/test/java/JvmArrayShapeTest.java

@@ -0,0 +1,135 @@
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.*;
+
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.TimeUnit;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * A signature returning a SQL array whose layout its wrapper gives beyond the C array, the
+ * catalog's arrayShape, on the Flink and Spark surfaces of tools/codegen_jvm.py
+ * (#_array_shape_overload there): each element is read from the C array, NULL where the output
+ * array the shape names as nulls is true, and with a rowLength the elements are cut into rows of
+ * the length the function it names answers, an array of arrays, as dumpValues(raster, integer,
+ * boolean) answers a height by width array of the values of a band, NULL for a nodata pixel.
+ * The generator runs over a catalog of two functions, as JvmStructFieldColumnTest runs it, so no
+ * jar and no libmeos.
+ */
+class JvmArrayShapeTest {
+
+    private static final Path TOOLS = Path.of("..", "tools").toAbsolutePath().normalize();
+
+    /** values(v const Temporal *, band int, nodata bool **, count int *) -> double *, of the SQL
+     * signature values(tint, integer) RETURNS float[] stating its null mask and its rows, and
+     * width(v const Temporal *) -> int answering the length of a row. */
+    private static final String CATALOG = "{\"functions\": ["
+        + "{\"name\": \"values\", \"file\": \"meos.h\", \"api\": \"public\", \"sqlfn\": \"values\","
+        + " \"returnType\": {\"c\": \"double *\", \"canonical\": \"double *\"},"
+        + " \"params\": [{\"name\": \"v\", \"cType\": \"const Temporal *\", \"canonical\": \"const Temporal *\"},"
+        + "  {\"name\": \"band\", \"cType\": \"int\", \"canonical\": \"int\"},"
+        + "  {\"name\": \"nodata\", \"cType\": \"bool **\", \"canonical\": \"bool **\"},"
+        + "  {\"name\": \"count\", \"cType\": \"int *\", \"canonical\": \"int *\"}],"
+        + " \"shape\": {\"arrayReturn\": {\"lengthFrom\": {\"kind\": \"param\", \"name\": \"count\"},"
+        + "  \"element\": {\"c\": \"double\", \"canonical\": \"double\"}},"
+        + "  \"outputArrays\": [{\"param\": \"nodata\"}], \"outParams\": [\"nodata\", \"count\"]},"
+        + " \"sqlSignatures\": [{\"args\": [\"tint\", \"integer\"], \"ret\": \"float[]\","
+        + "  \"arrayShape\": {\"nulls\": \"nodata\", \"rowLength\": {\"call\": \"width\", \"args\": [\"v\"]}}}]},"
+        + "{\"name\": \"width\", \"file\": \"meos.h\", \"api\": \"public\", \"sqlfn\": \"width\","
+        + " \"returnType\": {\"c\": \"int\", \"canonical\": \"int\"},"
+        + " \"params\": [{\"name\": \"v\", \"cType\": \"const Temporal *\", \"canonical\": \"const Temporal *\"}],"
+        + " \"sqlSignatures\": [{\"args\": [\"tint\"], \"ret\": \"integer\"}]}],"
+        + " \"typeEncodings\": {"
+        + "  \"Temporal\": {\"bytes\": {\"decoder\": \"temporal_from_wkb\", \"encoder\": \"temporal_as_wkb\"}}},"
+        + " \"structs\": [], \"objectModel\": {\"classes\": {}}, \"enums\": [], \"macros\": []}";
+
+    private static final String JAR = "{"
+        + "\"values\": {\"ret\": \"jnr.ffi.Pointer\", \"arg_types\": [\"jnr.ffi.Pointer\", \"int\", \"jnr.ffi.Pointer\", \"jnr.ffi.Pointer\"]},"
+        + "\"width\": {\"ret\": \"int\", \"arg_types\": [\"jnr.ffi.Pointer\"]},"
+        + "\"temporal_from_wkb\": {\"ret\": \"jnr.ffi.Pointer\", \"arg_types\": [\"byte[]\"]},"
+        + "\"temporal_as_wkb\": {\"ret\": \"byte[]\", \"arg_types\": [\"jnr.ffi.Pointer\", \"byte\"]}}";
+
+    /** Prints {engine: [{args, why, body, eval}]}, one entry per overload of values. */
+    private static final String DRIVER = String.join("\n",
+        "import json, sys",
+        "sys.path.insert(0, sys.argv[1])",
+        "import codegen_jvm as g",
+        "from collections import defaultdict",
+        "cat = json.load(open(sys.argv[2]))",
+        "jar = json.load(open(sys.argv[3]))",
+        "out = {}",
+        "for engine, m in (('flink', g.SqlModel(cat, jar)),",
+        "                  ('spark', g.SqlModel(cat, jar, g.SPARK_SQL_PKG, 'spark'))):",
+        "    rows = []",
+        "    for name, sargs, sdef, ov, why, isset, arity in g._sql_overloads(",
+        "            m, cat, jar, m.layout, defaultdict(int), lambda f: f['name'] == 'values'):",
+        "        if name == 'values':",
+        "            rows.append({'args': sargs, 'why': why, 'body': [] if ov is None else ov[3],",
+        "                         'cls': None if ov is None else ov[2],",
+        "                         'type': None if ov is None else ov[4]})",
+        "    out[engine] = rows",
+        "print(json.dumps(out))",
+        "");
+
+    private static JsonNode results;
+
+    @BeforeAll
+    static void generate() throws Exception {
+        Path dir = Files.createTempDirectory("array-shape");
+        try {
+            Files.writeString(dir.resolve("catalog.json"), CATALOG);
+            Files.writeString(dir.resolve("jar.json"), JAR);
+            Files.writeString(dir.resolve("driver.py"), DRIVER);
+            Process p = new ProcessBuilder("python3", "-B", dir.resolve("driver.py").toString(),
+                TOOLS.toString(), dir.resolve("catalog.json").toString(),
+                dir.resolve("jar.json").toString()).redirectErrorStream(true).start();
+            String out = new String(p.getInputStream().readAllBytes(), StandardCharsets.UTF_8);
+            assertTrue(p.waitFor(120, TimeUnit.SECONDS), "the generator ran");
+            assertEquals(0, p.exitValue(), out);
+            results = new ObjectMapper().readTree(out);
+        } finally {
+            try (var stream = Files.walk(dir)) {
+                stream.sorted(java.util.Comparator.reverseOrder()).forEach(q -> q.toFile().delete());
+            }
+        }
+    }
+
+    private static JsonNode built(String engine) {
+        JsonNode all = results.get(engine);
+        assertEquals(1, all.size(), engine + " " + all);
+        JsonNode one = all.get(0);
+        assertEquals("[\"tint\",\"integer\"]", one.get("args").toString(), engine);
+        assertTrue(one.get("why").isNull(), engine + " " + one);
+        return one;
+    }
+
+    @Test
+    void theArrayIsAnArrayOfRowsOfTheLengthTheShapeNames() {
+        assertEquals("DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.DOUBLE()))",
+            built("flink").get("type").asText());
+        assertEquals("DataTypes.createArrayType(DataTypes.createArrayType(DataTypes.DoubleType))",
+            built("spark").get("type").asText());
+        for (String engine : new String[] {"flink", "spark"}) {
+            JsonNode one = built(engine);
+            assertEquals("Double[][]", one.get("cls").asText(), engine);
+            String body = one.get("body").toString();
+            assertTrue(body.contains("_r = GeneratedFunctions.values(_p0, a1, _n, _cnt);"), engine + " " + body);
+            assertTrue(body.contains("int _w = GeneratedFunctions.width(_p0);"), engine + " " + body);
+            assertTrue(body.contains("Double[][] _rows = new Double[_c / _w][];"), engine + " " + body);
+        }
+    }
+
+    @Test
+    void anElementTheMaskStatesIsNullAndTheMaskIsFreedWithTheArray() {
+        for (String engine : new String[] {"flink", "spark"}) {
+            String body = built(engine).get("body").toString();
+            assertTrue(body.contains("_rows[_y][_x] = (_m.getByte((long) _i) != 0 ? null : "
+                + "_r.getDouble((long) _i * 8L));"), engine + " " + body);
+            assertTrue(body.contains("MeosSqlRuntime.free(_m);") && body.contains("MeosSqlRuntime.free(_r);"),
+                engine + " " + body);
+        }
+    }
+}

tools/codegen_jvm.py

@@ -1413,6 +1413,9 @@ def _signature_overload(m, f, sig, args, ret, jsig, bound, layout):
     if sig.get('retSet'):
         ov, why = _setret_overload(m, f, sig, jsig, bound, layout)
         return ov, why, ov is not None
+    if sig.get('arrayShape') and ret and ret.endswith('[]'):
+        ov, why = _array_shape_overload(m, f, sig, jsig, bound)
+        return ov, why, False
     order = sig.get('sqlArgParams') or (f.get('shape') or {}).get('sqlArgParams')
     ov, why = _overload(m, f, args, ret, jsig, bound, order)
     if ov is None and ret and ret.endswith('[]'):
@@ -1797,6 +1800,77 @@ def _setret_overload(m, f, sig, jsig, bound, layout):
     return (params, temps, f'{ecls}[]', body, m.array_type(etype)), None
 
 
+def _array_shape_overload(m, f, sig, jsig, bound):
+    """The eval of a signature returning a SQL array whose layout its wrapper gives beyond the
+    C array (the catalog's arrayShape), or (None, reason), as #_setret_overload builds the eval
+    of a set-returning signature: the C array read element by element, NULL where the output
+    array `nulls` is true, and, with `rowLength`, cut into rows of the length the public
+    function it names answers for the parameters it names, an array of arrays as PostgreSQL's
+    two-dimensional array: dumpValues(raster, integer, boolean) answers a row of values per
+    row of the raster, raster_width of it long, NULL for a nodata pixel."""
+    shape = sig['arrayShape']
+    vis, outs = m.visible(f)
+    if len(jsig['arg_types']) != len(f['params']):
+        return None, 'arrayShape:arity'
+    if jsig['ret'] != 'jnr.ffi.Pointer':
+        return None, 'arrayShape:ret'
+    counts = [p for p in outs if _norm(p['canonical']) == 'int *'
+              and 'const' not in p['canonical']]
+    nulls = shape.get('nulls')
+    if len(counts) != 1 or any(p is not counts[0] and p['name'] != nulls for p in outs):
+        return None, 'arrayShape:out'
+    elem = _norm(f['returnType']['canonical'])[:-1].strip()
+    hit = SETRET_SCALAR.get((elem, _deployed_type(sig['ret'][:-2])))
+    if not hit:
+        return None, f'arrayShape:element/{elem}'
+    rl = shape.get('rowLength')
+    if rl and ((m.jmeos.get(rl['call']) or {}).get('ret') != 'int'
+               or any(a not in [p['name'] for p in vis] for a in rl['args'])):
+        return None, 'arrayShape:rowLength'
+    jt = {p['name']: t for p, t in zip(f['params'], jsig['arg_types'])}
+    got, why = _inputs(m, f, [_deployed_type(a) for a in sig.get('args') or []],
+                       [jt[p['name']] for p in vis], bound,
+                       sig.get('sqlArgParams') or (f.get('shape') or {}).get('sqlArgParams'))
+    if got is None:
+        return None, why
+    params, call_vis, temps = got
+    vcall = dict(zip([p['name'] for p in vis], call_vis))
+    call = [vcall.get(p['name']) or ('_n' if p['name'] == nulls else '_cnt') for p in f['params']]
+    ecls = hit[0]
+    read = hit[1].format(p='_r', i='_i')
+    if nulls:
+        read = f'(_m.getByte((long) _i) != 0 ? null : {read})'
+    body = ['Pointer _cnt = MeosSqlRuntime.cell(4);']
+    body += ['Pointer _n = MeosSqlRuntime.cell(8);'] if nulls else []
+    body += ['Pointer _r = null;', 'Pointer _m = null;', 'try {',
+             f'    _r = GeneratedFunctions.{f["name"]}({", ".join(call)});',
+             '    if (_r == null) return null;',
+             '    int _c = _cnt.getInt(0L);']
+    body += ['    _m = _n.getPointer(0L);'] if nulls else []
+    if rl:
+        args = ', '.join(vcall[a] for a in rl['args'])
+        body += [f'    int _w = GeneratedFunctions.{rl["call"]}({args});',
+                 '    if (_w <= 0) return null;',
+                 f'    {ecls}[][] _rows = new {ecls}[_c / _w][];',
+                 '    for (int _y = 0; _y < _rows.length; _y++) {',
+                 f'        _rows[_y] = new {ecls}[_w];',
+                 '        for (int _x = 0; _x < _w; _x++) {',
+                 '            int _i = _y * _w + _x;',
+                 f'            _rows[_y][_x] = {read};',
+                 '        }',
+                 '    }']
+        cls, etype = f'{ecls}[][]', m.array_type(m.array_type(m.datatype(ecls)))
+    else:
+        body += [f'    {ecls}[] _rows = new {ecls}[_c];',
+                 '    for (int _i = 0; _i < _c; _i++) {',
+                 f'        _rows[_i] = {read};',
+                 '    }']
+        cls, etype = f'{ecls}[]', m.array_type(m.datatype(ecls))
+    body += ['    return _rows;', '} finally {', '    MeosSqlRuntime.free(_m);',
+             '    MeosSqlRuntime.free(_r);', '}']
+    return (params, temps, cls, body, etype), None
+
+
 def _aux_literal(a):
     """The Java literal of the value the catalog states for a trailing input: a number as
     written, a string quoted, and null where it states none (geo_as_hexewkb's endian, which