diff --git a/src/test/java/org/mobilitydb/spark/GeneratedSurfaceTest.java b/src/test/java/org/mobilitydb/spark/GeneratedSurfaceTest.java index 28dfb7a6..a503cf4e 100644 --- a/src/test/java/org/mobilitydb/spark/GeneratedSurfaceTest.java +++ b/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. 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 diff --git a/tools/codegen_spark_udfs.py b/tools/codegen_spark_udfs.py index 1ff622ea..4dbe5347 100644 --- a/tools/codegen_spark_udfs.py +++ b/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) -> *`: 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.""" + 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 readIntArr(Pointer p, int cnt) { + java.util.ArrayList 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 readLongArr(Pointer p, int cnt) { + java.util.ArrayList 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 readDoubleArr(Pointer p, int cnt) { + java.util.ArrayList 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) -> * 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. + 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