Generate the scalar-valued array accessors

src/test/java/org/mobilitydb/spark/GeneratedSurfaceTest.java

@@ -148,6 +148,23 @@ void json_path_query_surface() {
             "SELECT jsonb_path_exists('{\"a\": 1}', '$.b', '{}', false, false)")).intValue());
     }
 
+    @Test
+    void scalar_value_array_accessors() {
+        // The value-array accessors (tint_values / tfloat_values / *set_values / ...) return
+        // a native scalar array with the element count via an int* out-param, marshalled to
+        // a Spark array<int|long|double>. Inspect via Spark SQL size()/array_contains() (no
+        // Java-side Seq casting). tint_values of [1@.., 2@.., 1@..] = the distinct set {1, 2}.
+        assertEquals(2, ((Number) scalar(
+            "SELECT size(tint_values('" + TINT_HEX + "'))")).intValue());
+        assertEquals(Boolean.TRUE, scalar(
+            "SELECT array_contains(tint_values('" + TINT_HEX + "'), 1)"));
+        assertEquals(Boolean.TRUE, scalar(
+            "SELECT array_contains(tint_values('" + TINT_HEX + "'), 2)"));
+        // double element type: tfloat_values of [1.5@.., 2.5@..] = {1.5, 2.5}
+        assertEquals(2, ((Number) scalar(
+            "SELECT size(tfloat_values('" + TFLOAT_HEX + "'))")).intValue());
+    }
+
     @Test
     void portable_bare_name_dispatch_surface() {
         // The portable bare-name operator dialect, emitted by the generator's DISPATCH pass

tools/codegen_spark_udfs.py

@@ -591,6 +591,87 @@ def emit_tgeoarr(name, f, shape):
     return "\n".join(L)
 
 
+# Return-element base -> (Spark element type, UdfMarshal reader). ONLY fixed-width scalar
+# elements whose C array has a known stride; struct elements (Span*/STBox*/TBox*) are NOT
+# here (their per-element size is opaque to the binding), nor TimestampTz*/DateADT* (which
+# need MEOS-owned rendering, not raw longs).
+SCALAR_ELEM = {
+    "int":      ("IntegerType", "readIntArr"),
+    "int64_t":  ("LongType",    "readLongArr"),
+    "uint64_t": ("LongType",    "readLongArr"),
+    "double":   ("DoubleType",  "readDoubleArr"),
+}
+
+
+def scalar_values_shape(f):
+    """Recognize an array-returning accessor `(inputs..., int *count) -> <scalar>*`: the
+    trailing non-const `int *` is the element-count out-param and the return is a
+    contiguous C array of a fixed-width scalar (int/int64/uint64/double). JMEOS returns the
+    array Pointer and the caller allocates the count. Returns a shape dict or None."""
+    params, out = classify(f)
+    if out is not None or not params:
+        return None
+    last = params[-1]
+    if "const" in last["canonical"] or norm(last["canonical"]) != "int *":
+        return None
+    rt = norm(f["returnType"]["canonical"])
+    rb = base(f["returnType"]["canonical"])
+    if not rt.endswith("*") or rt.count("*") != 1 or rb not in SCALAR_ELEM:
+        return None
+    ins = params[:-1]
+    # every non-count input must marshal (a stray `int *`/pointer-out among the inputs is
+    # arg_kind None -> excluded, so array-in kernels like intset_make never match here).
+    for p in ins:
+        if arg_kind(p["canonical"]) is None:
+            return None
+    return {"elem": SCALAR_ELEM[rb], "ins": ins}
+
+
+def emit_scalar_values(name, f, shape):
+    """Emit a Spark UDF for a scalar array-returning accessor: marshal the inputs (mirroring
+    emit_single), allocate the `int *count` out-param, call, then read `count` fixed-width
+    elements from the returned array (freeing it) into a Spark array<int|long|double>."""
+    ins = shape["ins"]
+    ret_dt, reader = shape["elem"]
+    argnames = [_javaid(p["name"] or ("a%d" % i)) for i, p in enumerate(ins)]
+    argboxes, kinds = [], []
+    for p in ins:
+        k = arg_kind(p["canonical"]); kinds.append(k)
+        argboxes.append({"ptr": "String", "ts": "Object",
+                         "scalar": k[2] if k[0] == "scalar" else "String"}[k[0]])
+    retbox = "java.util.List<%s>" % {"IntegerType": "Integer", "LongType": "Long",
+                                     "DoubleType": "Double"}[ret_dt]
+    iface = ("UDF%d<%s, %s>" % (len(ins), ", ".join(argboxes), retbox)) if ins else "UDF0<%s>" % retbox
+    L = ['        spark.udf().register("%s", (%s) (%s) -> {' % (name, iface, ", ".join(argnames))]
+    if argnames:
+        L.append("        if (" + " || ".join("%s == null" % a for a in argnames) + ") return null;")
+    L.append("        MeosThread.ensureReady();")
+    callargs, frees = [], []
+    for a, p, k in zip(argnames, ins, kinds):
+        if k[0] == "ptr":
+            L.append("        jnr.ffi.Pointer p_%s = %s;" % (a, k[1] % a))
+            L.append("        if (p_%s == null) return null;" % a)
+            callargs.append("p_%s" % a); frees.append("p_%s" % a)
+        elif k[0] == "ts":
+            L.append("        java.time.OffsetDateTime dt_%s = UdfMarshal.tsOdt(%s);" % (a, a))
+            callargs.append("dt_%s" % a)
+        else:
+            callargs.append(a)
+    L.append("        jnr.ffi.Runtime _rt = jnr.ffi.Runtime.getSystemRuntime();")
+    L.append("        jnr.ffi.Pointer _cnt = jnr.ffi.Memory.allocateDirect(_rt, 4);")
+    callargs.append("_cnt")
+    call = "GeneratedFunctions.%s(%s)" % (f["name"], ", ".join(callargs))
+    L.append("        try {")
+    L.append("            jnr.ffi.Pointer _res = %s;" % call)
+    L.append("            return UdfMarshal.%s(_res, _cnt.getInt(0L));" % reader)
+    L.append("        } finally {")
+    for fr in frees:
+        L.append("            MeosMemory.free(%s);" % fr)
+    L.append("        }")
+    L.append("        }, DataTypes.createArrayType(DataTypes.%s));" % ret_dt)
+    return "\n".join(L)
+
+
 GEN_NOTE = "// GENERATED by tools/codegen_spark_udfs.py from the MEOS-API catalog. DO NOT EDIT.\n"
 IMPORTS = """\
 package org.mobilitydb.spark.generated;
@@ -744,6 +825,31 @@ def emit_tgeoarr(name, f, shape):
         return out;
     }
 
+    // Scalar array-return accessors (*_values, pose_orientation, quadbin_k_ring, ...) return
+    // a contiguous native array of `cnt` fixed-width elements (cnt from the int* out-param);
+    // read each at its known stride and free the array (malloc-backed, like readPairs).
+    static java.util.List<Integer> readIntArr(Pointer p, int cnt) {
+        java.util.ArrayList<Integer> o = new java.util.ArrayList<>();
+        if (p == null) return o;
+        for (int i = 0; i < cnt; i++) o.add(p.getInt((long) i * 4L));
+        MeosMemory.free(p);
+        return o;
+    }
+    static java.util.List<Long> readLongArr(Pointer p, int cnt) {
+        java.util.ArrayList<Long> o = new java.util.ArrayList<>();
+        if (p == null) return o;
+        for (int i = 0; i < cnt; i++) o.add(p.getLongLong((long) i * 8L));
+        MeosMemory.free(p);
+        return o;
+    }
+    static java.util.List<Double> readDoubleArr(Pointer p, int cnt) {
+        java.util.ArrayList<Double> o = new java.util.ArrayList<>();
+        if (p == null) return o;
+        for (int i = 0; i < cnt; i++) o.add(p.getDouble((long) i * 8L));
+        MeosMemory.free(p);
+        return o;
+    }
+
     // True iff a hex-WKB temporal pointer is a tnumber (tint/tfloat): only those have
     // a TBox value extent. Used SOLELY to select which existing backing to delegate to
     // for the axis-ambiguous space-X operators — never to compute a result.
@@ -1016,6 +1122,27 @@ def main():
         narr += 1
     print("  NxN array (tgeoarr) UDFs          : %d" % narr, file=sys.stderr)
 
+    # ── scalar-value array pass: (inputs..., int *count) -> <scalar>* accessors ──
+    # The value-array accessors (tint_values / tfloat_values / *set_values / th3index_values
+    # / tquadbin_values / pose_orientation / quadbin_k_ring) carry an `int *count` out-param,
+    # so the 1:1 and @sqlfn passes exclude them (arg:int *). Emit each under its C name: the
+    # return is a contiguous fixed-width scalar array -> a Spark array<int|long|double>.
+    nval = 0
+    for f in fns:
+        nm = f["name"]
+        if nm in names:
+            continue
+        if jar_syms is not None and nm not in jar_syms:
+            continue
+        shape = scalar_values_shape(f)
+        if shape is None:
+            continue
+        grouped.setdefault(class_for(f.get("group")), []).append(emit_scalar_values(nm, f, shape))
+        names.add(nm)
+        cov += 1
+        nval += 1
+    print("  scalar-value array UDFs           : %d" % nval, file=sys.stderr)
+
     # ── @sqlfn CANONICAL-NAME pass: emit the MobilityDB SQL surface ──
     # Every catalog function carries the canonical MobilityDB SQL spelling in its
     # @sqlfn tag (numInstants, eIntersects, atTime, asHexWKB ...). That is the name a