From fed43e7de638eb6861f79bf84aab20a896e58202 Mon Sep 17 00:00:00 2001 From: Esteban Zimanyi Date: Sat, 10 Oct 2026 10:53:12 +0200 Subject: [PATCH] Run the aggregates of the catalog from the roles it states on both typed SQL surfaces An aggregate of the catalog's aggregates section runs from the public MEOS function of each role it states: the transition folds a value into the state, the combine joins two partial states, the final answers the result, and a partial state crosses between workers as the bytes the serialize function writes and the deserialize function reads, or, for a state of a SQL type (the box of extent), as the form that type's codec writes. _catalog_aggregates of tools/codegen_jvm.py chooses, for both engines, the aggregates taking one value with a combine and a state either internal with a serialize and a deserialize function or of a SQL type the surface holds, and MeosAggregate holds their roles: a combine answers in the place of its first state, returning it or freeing it for a larger one, except that a skip-list combine answers its second state when the first is empty, and the state it does not answer is released by releasing what its final answers, the final consuming its state. The typed Flink surface registers each as an AggregateFunction whose accumulator, MeosAggState, holds the bytes of the state and no native pointer; the typed Spark surface registers each through MeosAggregates, a partition keeping its state between two values and writing it when it leaves the executor. Both register extent, tCount, tSum, tAvg, tCentroid, tdensity, tnpoints, mergeAgg, tMinAgg, tMaxAgg, tAndAgg, tOrAgg, setUnionAgg, spanUnionAgg and spansetUnionAgg, a name a scalar of the surface carries going to the scalar as main's generator names the aggregates, setUnion, spanUnion and spansetUnion being the unions of a value and a set. The Spark UDF arm takes its temporal aggregates from the same section, a partial state crossing as the bytes of its serialize function, so the pairing of temporal_tagg_finalfn with temporal_to_taggstate leaves both arms and temporal_to_taggstate joins the gaps ledger. Witness: over MobilityDB bdc4e41234 main's generator registers no aggregate on the typed Flink surface, and on the typed Spark surface and the Spark UDF arm 7 whose final function is temporal_tagg_finalfn, a partial state crossing as the temporal value that function answers, which carries neither the sums and counts of tAvg nor the coordinates of tCentroid; MobilitySpark fails on "Cannot resolve routine `tAvg`" and MobilityFlink's test does not compile, the surface holding no class TAvg. Measured over MobilityDB bdc4e41234 with the catalog of MEOS-API d259ddbfea: both typed surfaces register 15 aggregates of 141 overloads besides the 1163 SQL functions and 8931 overloads main's generator registers; the 67 overloads of more than one argument (wCount, appendInstant), the 40 without a combine (appendInstant, appendSequence) and the 12 over a value of no surface type (setUnion over an integer) are counted as left out. The Spark UDF arm registers 11 temporal aggregates against 7, adding tAvg, tCentroid, tdensity and tnpoints, and the gaps ledger loses the 11 role functions it reaches and gains temporal_to_taggstate. MobilitySpark passes its 38 tests and MobilityFlink its GeneratedSqlSurfaceTest with assertions that tCount, tAvg, tCentroid, setUnionAgg, spanUnionAgg and extent answer what MobilityDB answers over values in partitions of their own, in a streaming job and in a two-phase batch job, and that tAvg joins two partial states one of which crosses the accumulator's serializer. Why: an aggregate SQL states answers on both engines as PostgreSQL answers it, from the roles MEOS carries, with no hand rule pairing two functions. --- tools/codegen_jvm.py | 621 ++++++++++++++++++++++++++++++------ tools/codegen_spark_udfs.py | 168 +++++----- tools/spark-udf-gaps.txt | 12 +- 3 files changed, 614 insertions(+), 187 deletions(-) diff --git a/tools/codegen_jvm.py b/tools/codegen_jvm.py index f0ed44d81..d3df9a149 100644 --- a/tools/codegen_jvm.py +++ b/tools/codegen_jvm.py @@ -1845,6 +1845,16 @@ def _value_class_src(m, sql): .build(); } + /** The inference of #inference for an aggregate, whose accumulator is of the type acc. */ + public static TypeInference aggregateInference(DataType[][] args, DataType[] rets, DataType acc) { + TypeInference t = inference(args, rets); + return TypeInference.newBuilder() + .inputTypeStrategy(t.getInputTypeStrategy()) + .outputTypeStrategy(t.getOutputTypeStrategy()) + .accumulatorTypeStrategy(TypeStrategies.explicit(acc)) + .build(); + } + ''' _RUNTIME_CORE = ''' /** Free each value MEOS allocated. Null-safe. */ public static void free(Pointer[] ps) { @@ -2225,6 +2235,18 @@ def run_flink_sql(args): body.append('}') (root / 'functions' / f'{cls}.java').write_text('\n'.join(body) + '\n') + aggs, agg_left = _catalog_aggregates(m, cat, jmeos, {n.lower() for n in names}) + (root / 'MeosAggregate.java').write_text(_AGGREGATE_ROLES.format(pkg=SQL_PKG)) + (root / 'MeosAggState.java').write_text(_FLINK_AGG_STATE) + for sqlname in sorted(aggs): + cls = _javaid(sqlname) + while cls.lower() in taken: + cls += '_' + taken.add(cls.lower()) + classes[sqlname] = cls + (root / 'functions' / f'{cls}.java').write_text( + _flink_aggregate_src(sqlname, cls, aggs[sqlname])) + reg = [f'package {SQL_PKG};', '', 'import org.apache.flink.table.api.TableEnvironment;', '', '/* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand.', @@ -2246,10 +2268,14 @@ def run_flink_sql(args): (root / 'MobilityFlinkSql.java').write_text('\n'.join(reg) + '\n') n_ov = sum(len(v) for v in names.values()) - print(f'flink-sql: {len(classes)} SQL functions ({n_ov} overloads, {nset} of them ' + print(f'flink-sql: {len(names)} SQL functions ({n_ov} overloads, {nset} of them ' f'set-returning signatures), {len(m.value_class)} value types into {root}') for why, n in sorted(skipped.items(), key=lambda x: -x[1]): print(f' skipped {n:5d} {why}') + print(f'flink-sql: {len(aggs)} aggregates ' + f'({sum(len(v) for v in aggs.values())} overloads)') + for why, n in sorted(agg_left.items()): + print(f' aggregate left out {n:3d} {why}') # ───────────────────────── spark-sql: the typed Spark SQL surface ───────────────────────── @@ -2539,14 +2565,408 @@ def _spark_planning(arity): ''' -# The temporal aggregates of the typed surface, emitted by #_spark_aggregates_src. Each overload -# delegates to #TemporalAggregate, the aggregate the UDF surface generates (codegen_spark_udfs.py), -# over the bytes a typed value carries, so a partial aggregate travels between executors as that -# class makes it travel; this class adds only the typed value at each end and the resolution by -# type, which mirrors #_spark_planning's register. +# ── aggregates: the catalog's `aggregates` over the typed values of both surfaces ── +# An aggregate of the catalog (A9) states, per CREATE AGGREGATE, its argument and result types, +# the type of its state, and the public MEOS function of each role it defines: transition, +# combine, final, serialize and deserialize. Both engines run it from those roles alone: the +# transition folds a value into the state, the combine joins two partial states, the final +# answers the result, and a partial state travels as the bytes the serialize function writes and +# the deserialize function reads, or, for a state of a SQL type (the box of `extent`), as the form +# that type's codec writes. #MeosAggregate holds the roles and both engines share it; Spark keeps +# the state between two values and writes it only when a partial aggregate leaves the executor, +# while Flink holds only the bytes, a native state never being a field of a Flink accumulator. + +# The roles of one aggregate overload over a MEOS state, shared by both surfaces. A combine +# answers in the place of its first state, returning it or freeing it for a larger one +# (set_union_combinefn), except that a skip-list combine answers its second state when the first +# is empty (temporal_tagg_combinefn): the state it does not answer is released. A final function +# consumes its state, so a state is released by releasing what its final answers. +_AGGREGATE_ROLES = '''package {pkg}; + +import java.io.Serializable; +import jnr.ffi.Pointer; + +/* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand. + * The roles of one aggregate of the catalog over its MEOS state: the transition folding a value + * into the state, the combine joining two partial states, the final answering the result, and the + * writing of a state as bytes and its reading back. */ +public final class MeosAggregate implements Serializable {{ + + /** The transition: the state with the value folded in, the value of the overload's class. */ + public interface Transition extends Serializable {{ + Pointer apply(Pointer state, Object value); + }} + + /** The combine: the two partial states joined. */ + public interface Combine extends Serializable {{ + Pointer apply(Pointer a, Pointer b); + }} + + /** The final: the result over a state, which it consumes. */ + public interface Final extends Serializable {{ + Pointer apply(Pointer state); + }} + + /** The bytes of a state, which is kept. */ + public interface Write extends Serializable {{ + byte[] apply(Pointer state); + }} + + /** The state the bytes hold. */ + public interface Read extends Serializable {{ + Pointer apply(byte[] form); + }} + + /** The value of the result type holding what the final answers. */ + public interface Make extends Serializable {{ + MeosValue apply(Pointer result); + }} + + private final Transition transition; + private final Combine combine; + private final Final fin; + private final Write write; + private final Read read; + private final Make make; + + /** fin is null when the aggregate states no final function: its result is its state. */ + public MeosAggregate(Transition transition, Combine combine, Final fin, Write write, Read read, + Make make) {{ + this.transition = transition; + this.combine = combine; + this.fin = fin; + this.write = write; + this.read = read; + this.make = make; + }} + + /** The state with the value folded in; the state given is the transition's to keep or free. */ + public Pointer step(Pointer state, Object value) {{ + return transition.apply(state, value); + }} + + /** The two partial states joined; both are consumed. */ + public Pointer combine(Pointer a, Pointer b) {{ + if (a == null) {{ + return b; + }} + if (b == null) {{ + return a; + }} + Pointer r = combine.apply(a, b); + if (r != null && r.address() == b.address()) {{ + release(a); + }} else {{ + release(b); + }} + return r; + }} + + /** Release a state. */ + public void release(Pointer state) {{ + if (state != null) {{ + MeosSqlRuntime.free(fin == null ? state : fin.apply(state)); + }} + }} + + /** The result over a state, which is consumed. */ + public MeosValue finish(Pointer state) {{ + if (state == null) {{ + return null; + }} + Pointer r = fin == null ? state : fin.apply(state); + try {{ + return r == null ? null : make.apply(r); + }} finally {{ + MeosSqlRuntime.free(r); + }} + }} + + /** The bytes of a state, which is kept. */ + public byte[] write(Pointer state) {{ + return state == null ? null : write.apply(state); + }} + + /** The state the bytes of #write hold. */ + public Pointer read(byte[] form) {{ + return form == null ? null : read.apply(form); + }} +}} +''' + + +def _catalog_aggregates(m, cat, jmeos, scalars): + """The aggregates of the catalog's `aggregates` section both typed surfaces carry: + {SQL name: [overload]} and {reason: count} for what is left out. An overload is a dict of its + argument and result classes, the transition's inputs as #_inputs passes them, and the Java of + its roles. An aggregate is carried when it takes one value of a type the surface holds, + states a transition and a combine, and its state is either internal with a serialize and a + deserialize function or of a SQL type the surface holds, each role a MEOS function of the + jar. Each engine holds one function per name regardless of case, so a name a scalar of the + surface carries stays the scalar's and the aggregate takes the `Agg` suffix of its canonical + name: the catalog's own `Agg` name where it states one beside the bare one (mergeAgg, + tMinAgg), else the bare name with the suffix (tAndAgg).""" + aggs, left, seen = defaultdict(list), defaultdict(int), set() + stated = {a['sqlName'].lower() for a in cat.get('aggregates') or []} + for a in cat.get('aggregates') or []: + meos = lambda k: (a.get(k) or {}).get('meos') # noqa: E731 + args = [_deployed_type(t) for t in a.get('args') or []] + name = a['sqlName'] + if name.lower() in scalars: + if (name + 'Agg').lower() in stated: + continue + name += 'Agg' + if len(args) != 1: + left[f'{len(args)} arguments'] += 1 + continue + if not meos('combine'): + left['no combine'] += 1 + continue + stype, ret = _deployed_type(a['stype']), _deployed_type(a['ret']) + if stype == 'internal': + if not (meos('serialize') and meos('deserialize')): + left['internal state without serialize'] += 1 + continue + elif stype not in m.value_class: + left[f'state {stype}'] += 1 + continue + if args[0] not in m.value_class or ret not in m.value_class: + left[f'argument {args[0]} or result {ret} of no value type'] += 1 + continue + roles = [meos(k) for k in ('transition', 'combine', 'final', 'serialize', 'deserialize')] + if any(r and r not in jmeos for r in roles): + left['not in jar'] += 1 + continue + f = m.by_name.get(meos('transition')) + jt = jmeos[meos('transition')]['arg_types'] + got, why = _inputs(m, dict(f, params=f['params'][1:]), args, jt[1:]) + if got is None or any(kind != 'value' for _, _, kind in got[2]): + left[f'transition input {why or "of no value"}'] += 1 + continue + if (name.lower(), args[0]) in seen: + continue + seen.add((name.lower(), args[0])) + cls = lambda sql: f'{m.pkg}.types.{m.value_class[sql]}' # noqa: E731 + if stype == 'internal': + write = f'GeneratedFunctions::{meos("serialize")}' + read = f'GeneratedFunctions::{meos("deserialize")}' + else: + write = f'p -> {cls(stype)}.encode(p).form' + read = f'b -> new {cls(stype)}(b).decode()' + fin = f'GeneratedFunctions::{meos("final")}' if meos('final') else 'null' + aggs[name].append({'arg': cls(args[0]), 'ret': cls(ret), + 'transition': meos('transition'), 'inputs': got, + 'roles': (f'GeneratedFunctions::{meos("combine")}', fin, write, read, + f'{cls(ret)}::encode')}) + return aggs, left + + +def _transition_src(k, ov): + """The static transition of overload k of an aggregate: its value decoded as #_inputs decodes + an argument, folded into the state, and released after the call.""" + params, call, temps = ov['inputs'] + (fc, name), = params + L = [f' static Pointer t{k}(Pointer state, Object value) {{', + f' {fc} {name} = ({fc}) value;'] + L += [f' Pointer {t} = {e};' for t, e, _ in temps] + L.append(f' Pointer[] _in = {{{", ".join(t for t, _, _ in temps)}}};') + L += [' try {', + f' return GeneratedFunctions.{ov["transition"]}(state, {", ".join(call)});', + ' } finally {', ' MeosSqlRuntime.free(_in);', ' }', ' }', ''] + return L + + +def _roles_src(cls, k, ov, indent=' '): + """The MeosAggregate of overload k of the aggregate class cls.""" + comb, fin, write, read, make = ov['roles'] + return (f'{indent}new MeosAggregate({cls}::t{k}, {comb}, {fin},\n' + f'{indent} {write}, {read}, {make})') + + +# The Flink accumulator of every aggregate: the bytes of its state and the overload that wrote +# them, carried as a RAW type, so Flink copies and checkpoints it as bytes. Each value is folded +# into the state the bytes hold, which is then written back and released. +_FLINK_AGG_STATE = f'''package {SQL_PKG}; + +import java.io.IOException; +import jnr.ffi.Pointer; +import org.apache.flink.api.common.typeutils.SimpleTypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot; +import org.apache.flink.api.common.typeutils.base.TypeSerializerSingleton; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.types.DataType; + +/* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand. + * The accumulator of an aggregate: the bytes of its MEOS state, as MeosAggregate#write writes + * them, and the overload that wrote them; Flink holds no native state between two calls. */ +public final class MeosAggState {{ + + /** The Flink SQL type of an accumulator. */ + public static final DataType TYPE = DataTypes.RAW(MeosAggState.class, new Serializer()); + + byte[] state; + int overload = -1; + + /** Fold a value into the state, by the roles of overload k. */ + public void add(MeosAggregate[] roles, int k, Object value) {{ + if (value == null) {{ + return; + }} + MeosAggregate r = roles[k]; + Pointer s = r.step(r.read(state), value); + try {{ + state = r.write(s); + }} finally {{ + r.release(s); + }} + overload = k; + }} + + /** Join the partial states of others into this one. */ + public void merge(MeosAggregate[] roles, Iterable others) {{ + for (MeosAggState o : others) {{ + if (o.state == null) {{ + continue; + }} + if (state == null) {{ + state = o.state.clone(); + overload = o.overload; + continue; + }} + MeosAggregate r = roles[overload]; + Pointer s = r.combine(r.read(state), r.read(o.state)); + try {{ + state = r.write(s); + }} finally {{ + r.release(s); + }} + }} + }} + + /** The result over the state, null over no value. */ + public MeosValue value(MeosAggregate[] roles) {{ + return state == null ? null : roles[overload].finish(roles[overload].read(state)); + }} + + public static final class Serializer extends TypeSerializerSingleton {{ + @Override public boolean isImmutableType() {{ return false; }} + @Override public MeosAggState createInstance() {{ return new MeosAggState(); }} + @Override public int getLength() {{ return -1; }} + + @Override + public MeosAggState copy(MeosAggState from) {{ + MeosAggState c = new MeosAggState(); + c.state = from.state == null ? null : from.state.clone(); + c.overload = from.overload; + return c; + }} + + @Override + public MeosAggState copy(MeosAggState from, MeosAggState reuse) {{ + return copy(from); + }} + + @Override + public void serialize(MeosAggState v, DataOutputView out) throws IOException {{ + out.writeInt(v.overload); + out.writeInt(v.state == null ? -1 : v.state.length); + if (v.state != null) {{ + out.write(v.state); + }} + }} + + @Override + public MeosAggState deserialize(DataInputView in) throws IOException {{ + MeosAggState v = new MeosAggState(); + v.overload = in.readInt(); + int n = in.readInt(); + if (n >= 0) {{ + v.state = new byte[n]; + in.readFully(v.state); + }} + return v; + }} + + @Override + public MeosAggState deserialize(MeosAggState reuse, DataInputView in) throws IOException {{ + return deserialize(in); + }} + + @Override + public void copy(DataInputView in, DataOutputView out) throws IOException {{ + serialize(deserialize(in), out); + }} + + @Override + public TypeSerializerSnapshot snapshotConfiguration() {{ + return new Snapshot(); + }} + }} + + public static final class Snapshot extends SimpleTypeSerializerSnapshot {{ + public Snapshot() {{ + super(Serializer::new); + }} + }} +}} +''' + + +def _flink_aggregate_src(sqlname, cls, ovs): + """The Flink AggregateFunction of the SQL aggregate sqlname: one accumulate per overload, + which Flink chooses from the class of the argument, the roles of each, and the type inference + over them.""" + body = [f'package {SQL_PKG}.functions;', '', + 'import functions.GeneratedFunctions;', + 'import jnr.ffi.Pointer;', + 'import org.apache.flink.table.catalog.DataTypeFactory;', + 'import org.apache.flink.table.functions.AggregateFunction;', + 'import org.apache.flink.table.types.DataType;', + 'import org.apache.flink.table.types.inference.TypeInference;', + f'import {SQL_PKG}.MeosAggState;', + f'import {SQL_PKG}.MeosAggregate;', + f'import {SQL_PKG}.MeosSqlRuntime;', + f'import {SQL_PKG}.MeosValue;', '', + '/* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand.', + f' * The MobilityDB SQL aggregate {sqlname}: {len(ovs)} overload(s). */', + f'public final class {cls} extends AggregateFunction {{', '', + ' private static final MeosAggregate[] ROLES = {', + ',\n'.join(_roles_src(cls, k, ov) for k, ov in enumerate(ovs)), + ' };', ''] + for k, ov in enumerate(ovs): + body += _transition_src(k, ov) + body += [' @Override', ' public MeosAggState createAccumulator() {', + ' return new MeosAggState();', ' }', ''] + for k, ov in enumerate(ovs): + body += [f' public void accumulate(MeosAggState acc, {ov["arg"]} a0) {{', + f' acc.add(ROLES, {k}, a0);', ' }', ''] + body += [' public void merge(MeosAggState acc, Iterable others) {', + ' acc.merge(ROLES, others);', ' }', '', + ' @Override', ' public MeosValue getValue(MeosAggState acc) {', + ' return acc.value(ROLES);', ' }', '', + ' @Override', + ' public TypeInference getTypeInference(DataTypeFactory typeFactory) {', + ' return MeosSqlRuntime.aggregateInference(new DataType[][] {'] + body += [f' {{{_datatype(ov["arg"])}}},' for ov in ovs] + body += [' }, new DataType[] {'] + body += [f' {_datatype(ov["ret"])},' for ov in ovs] + body += [' }, MeosAggState.TYPE);', ' }', '}'] + return '\n'.join(body) + '\n' + + +# The aggregates of the typed Spark surface: each SQL name registered once with the builder that +# chooses its overload from the type of its argument, as #_spark_planning's register chooses a +# function's, a call over any other argument going to the function held before under the name, +# the aggregate of the UDF surface over binary values. _SPARK_AGGREGATES = '''package {pkg}; -import functions.GeneratedFunctions; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.io.Serializable; +import jnr.ffi.Pointer; import org.apache.spark.sql.Encoder; import org.apache.spark.sql.Encoders; import org.apache.spark.sql.SparkSession; @@ -2556,54 +2976,62 @@ def _spark_planning(arity): import org.apache.spark.sql.execution.aggregate.ScalaAggregator; import org.apache.spark.sql.expressions.Aggregator; import org.apache.spark.sql.types.DataType; -import org.mobilitydb.spark.generated.TemporalAggregate; import scala.Function1; import scala.Option; import scala.collection.immutable.Seq; import scala.runtime.AbstractFunction1; /* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand. - * The temporal aggregates of the catalog over the typed values of this surface, each SQL name - * registered once with the builder that chooses its overload from the type of its argument, as - * MeosSqlRuntime.register chooses a function's. An overload runs the TemporalAggregate of the UDF - * surface over the bytes the value carries, and a call over any other argument goes to the - * function held before under the name, the aggregate of the UDF surface over binary values. */ -public final class TemporalAggregates {{ - - private TemporalAggregates() {{ }} - - /** The value of a result type, made from the bytes the aggregate answers. */ - public interface Make extends java.io.Serializable {{ - MeosValue apply(byte[] form); - }} - - /** One overload: the SQL type it takes, the classes of its argument and result, and the - * aggregate computing it. */ - public static final class Overload implements java.io.Serializable {{ + * The aggregates of the catalog over the typed values of this surface, each run from the roles + * it states (MeosAggregate) and each SQL name registered once with the builder that chooses its + * overload from the type of its argument. A partition keeps its MEOS state between two values, + * and the state travels between executors as the bytes MeosAggregate#write writes. A call over + * any other argument goes to the function held before under the name, the aggregate of the UDF + * surface over binary values. */ +public final class MeosAggregates {{ + + private MeosAggregates() {{ }} + + /** One overload: the SQL type it takes, the classes of its argument and result, and its + * roles. */ + public static final class Overload implements Serializable {{ final DataType arg; final Class in; final Class out; - final Make make; - final TemporalAggregate agg; + final MeosAggregate roles; public Overload(DataType arg, Class in, Class out, - Make make, TemporalAggregate agg) {{ + MeosAggregate roles) {{ this.arg = arg; this.in = in; this.out = out; - this.make = make; - this.agg = agg; + this.roles = roles; }} }} - public static void registerAll(SparkSession spark) {{ -{registrations} }} + /** The state of a partition, written as the bytes of its roles when it leaves the executor. */ + public static final class Buffer implements Serializable {{ + final MeosAggregate roles; + transient Pointer state; + + Buffer(MeosAggregate roles) {{ + this.roles = roles; + }} - private static TemporalAggregate steps(TemporalAggregate.Step trans, TemporalAggregate.Step comb) {{ - return new TemporalAggregate(new int[] {{-1}}, new TemporalAggregate.Step[] {{trans}}, - new TemporalAggregate.Step[] {{comb}}); + private void writeObject(ObjectOutputStream out) throws IOException {{ + out.defaultWriteObject(); + out.writeObject(roles.write(state)); + }} + + private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {{ + in.defaultReadObject(); + state = roles.read((byte[]) in.readObject()); + }} }} + public static void registerAll(SparkSession spark) {{ +{registrations} }} + private static void register(SparkSession spark, String name, Overload... overloads) {{ FunctionIdentifier id = FunctionIdentifier.apply(name); Option, Expression>> before = @@ -2637,85 +3065,69 @@ def _spark_planning(arity): }}, "scala_udf"); }} - /** An aggregate over typed values: the TemporalAggregate of the overload over their bytes. */ - static final class Typed extends Aggregator {{ - private final Overload o; + /** An aggregate over typed values, run by the roles of its overload. */ + static final class Typed extends Aggregator {{ + private final MeosAggregate r; private final Encoder outEnc; @SuppressWarnings("unchecked") Typed(Overload o, Encoder outEnc) {{ - this.o = o; + this.r = o.roles; this.outEnc = (Encoder) outEnc; }} - @Override public TemporalAggregate.Buffer zero() {{ return o.agg.zero(); }} + @Override public Buffer zero() {{ return new Buffer(r); }} - @Override public TemporalAggregate.Buffer reduce(TemporalAggregate.Buffer b, MeosValue v) {{ - return o.agg.reduce(b, v == null ? null : v.form); + @Override public Buffer reduce(Buffer b, MeosValue v) {{ + if (v != null) {{ + b.state = r.step(b.state, v); + }} + return b; }} - @Override public TemporalAggregate.Buffer merge(TemporalAggregate.Buffer b1, TemporalAggregate.Buffer b2) {{ - return o.agg.merge(b1, b2); + @Override public Buffer merge(Buffer b1, Buffer b2) {{ + b1.state = r.combine(b1.state, b2.state); + b2.state = null; + return b1; }} - @Override public MeosValue finish(TemporalAggregate.Buffer b) {{ - byte[] form = o.agg.finish(b); - return form == null ? null : o.make.apply(form); + @Override public MeosValue finish(Buffer b) {{ + Pointer s = b.state; + b.state = null; + return r.finish(s); }} - @Override public Encoder bufferEncoder() {{ return o.agg.bufferEncoder(); }} + @Override public Encoder bufferEncoder() {{ return Encoders.javaSerialization(Buffer.class); }} @Override public Encoder outputEncoder() {{ return outEnc; }} }} }} ''' -def _spark_aggregates(m, cat, jmeos, scalars): - """The temporal aggregates of the catalog's `aggregates` section the typed surface carries: - {SQL name: [(argument SQL type, result SQL type, transition, combine)]}, and what is left out. - Mirrors the temporal aggregate pass of codegen_spark_udfs.py, whose #TemporalAggregate runs - them: an aggregate is carried when it takes one temporal value, its transition and combine - functions are MEOS's and its final one is temporal_tagg_finalfn, so that a partial aggregate - travels between executors as the temporal value that function returns. Spark holds one - function per name regardless of case, so a name a scalar of this surface carries stays the - scalar's and the aggregate takes the `Agg` suffix of its canonical name, as the UDF surface - names it: the catalog's own `Agg` name where it states one beside the bare one (mergeAgg, - tMinAgg), else the bare name with the suffix (tAndAgg).""" - aggs, left = defaultdict(list), defaultdict(int) - need = ('temporal_tagg_finalfn', 'temporal_to_taggstate') - stated = {a['sqlName'].lower() for a in cat.get('aggregates') or []} - for a in cat.get('aggregates') or []: - meos = lambda k: (a.get(k) or {}).get('meos') # noqa: E731 - args = [_deployed_type(t) for t in a.get('args') or []] - if len(args) != 1 or meos('final') != 'temporal_tagg_finalfn' \ - or not meos('transition') or not meos('combine'): - continue - name = a['sqlName'] - if name.lower() in scalars: - if (name + 'Agg').lower() in stated: - continue - name += 'Agg' - if args[0] not in m.value_class or _deployed_type(a['ret']) not in m.value_class: - left[f"{name}: no value type"] += 1 - continue - if any(s not in jmeos for s in (meos('transition'), meos('combine')) + need): - left[f"{name}: not in jar"] += 1 - continue - aggs[name].append((args[0], a['ret'], meos('transition'), meos('combine'))) - return aggs, left - - -def _spark_aggregates_src(m, aggs): - """TemporalAggregates.java over the aggregates #_spark_aggregates carries; the twin of - #_spark_value_class_src for the aggregates.""" - t = lambda sql: f'{SPARK_SQL_PKG}.types.{m.value_class[sql]}' # noqa: E731 - regs = [] - for name in sorted(aggs): - ovs = [f' new Overload({t(arg)}.TYPE, {t(arg)}.class, {t(ret)}.class, {t(ret)}::new,\n' - f' steps(GeneratedFunctions::{trans}, GeneratedFunctions::{comb}))' - for arg, ret, trans, comb in aggs[name]] - regs.append(f' register(spark, "{name}",\n' + ',\n'.join(ovs) + ');\n') - return _SPARK_AGGREGATES.format(pkg=SPARK_SQL_PKG, registrations=''.join(regs)) +def _spark_aggregate_src(sqlname, cls, ovs): + """The class of the Spark aggregate sqlname: the transition of each overload and the + overloads MeosAggregates registers, the twin of #_flink_aggregate_src.""" + body = [f'package {SPARK_SQL_PKG}.functions;', '', + 'import functions.GeneratedFunctions;', + 'import jnr.ffi.Pointer;', + f'import {SPARK_SQL_PKG}.MeosAggregate;', + f'import {SPARK_SQL_PKG}.MeosAggregates;', + f'import {SPARK_SQL_PKG}.MeosSqlRuntime;', '', + '/* AUTO-GENERATED by tools/codegen_jvm.py — do not edit by hand.', + f' * The MobilityDB SQL aggregate {sqlname}: {len(ovs)} overload(s). */', + f'public final class {cls} {{', '', + f' private {cls}() {{ }}', ''] + for k, ov in enumerate(ovs): + body += _transition_src(k, ov) + body += [' /** The overloads of this aggregate, in the order the builder tries them. */', + ' public static MeosAggregates.Overload[] overloads() {', + ' return new MeosAggregates.Overload[] {', + ',\n'.join(f' new MeosAggregates.Overload({ov["arg"]}.TYPE, ' + f'{ov["arg"]}.class, {ov["ret"]}.class,\n' + + _roles_src(cls, k, ov, ' ') + ')' + for k, ov in enumerate(ovs)), + ' };', ' }', '}'] + return '\n'.join(body) + '\n' def _spark_value_class_src(m, sql): @@ -2929,9 +3341,19 @@ def run_spark_sql(args): body += [' };', ' }', '}'] (root / 'functions' / f'{cls}.java').write_text('\n'.join(body) + '\n') - aggs, agg_left = _spark_aggregates(m, cat, jmeos, {n.lower() for n in names}) - if aggs: - (root / 'TemporalAggregates.java').write_text(_spark_aggregates_src(m, aggs)) + aggs, agg_left = _catalog_aggregates(m, cat, jmeos, {n.lower() for n in names}) + (root / 'MeosAggregate.java').write_text(_AGGREGATE_ROLES.format(pkg=SPARK_SQL_PKG)) + regs = [] + for sqlname in sorted(aggs): + cls = _javaid(sqlname) + while cls.lower() in taken: + cls += '_' + taken.add(cls.lower()) + (root / 'functions' / f'{cls}.java').write_text( + _spark_aggregate_src(sqlname, cls, aggs[sqlname])) + regs.append(f' register(spark, "{sqlname}", {SPARK_SQL_PKG}.functions.{cls}.overloads());\n') + (root / 'MeosAggregates.java').write_text( + _SPARK_AGGREGATES.format(pkg=SPARK_SQL_PKG, registrations=''.join(regs))) reg = [f'package {SPARK_SQL_PKG};', '', 'import org.apache.spark.sql.SparkSession;', '', @@ -2944,8 +3366,7 @@ def run_spark_sql(args): chunks = [items[i:i + 200] for i in range(0, len(items), 200)] reg.append(' public static void registerAll(SparkSession spark) {') reg += [f' register{i}(spark);' for i in range(len(chunks))] - if aggs: - reg.append(' TemporalAggregates.registerAll(spark);') + reg.append(' MeosAggregates.registerAll(spark);') reg += [' }', ''] for i, chunk in enumerate(chunks): reg.append(f' private static void register{i}(SparkSession spark) {{') @@ -2960,7 +3381,7 @@ def run_spark_sql(args): f'set-returning signatures), {len(m.value_class)} value types into {root}') for why, n in sorted(skipped.items(), key=lambda x: -x[1]): print(f' skipped {n:5d} {why}') - print(f'spark-sql: {len(aggs)} temporal aggregates ' + print(f'spark-sql: {len(aggs)} aggregates ' f'({sum(len(v) for v in aggs.values())} overloads)') for why, n in sorted(agg_left.items()): print(f' aggregate left out {n:3d} {why}') diff --git a/tools/codegen_spark_udfs.py b/tools/codegen_spark_udfs.py index bf96e1e02..405d4c9d9 100644 --- a/tools/codegen_spark_udfs.py +++ b/tools/codegen_spark_udfs.py @@ -1532,12 +1532,12 @@ def emit_setret(name, cands, colnames): """ # A temporal aggregate of MEOS as a Spark Aggregator over WKB temporals, a generated -# helper written beside UdfMarshal like the MARSHAL template above. The transition and -# combine functions of the family are handed in per temporal type, and the buffer holds -# the native state of the partition. When Spark moves a buffer between executors it writes -# the state as the temporal value temporal_tagg_finalfn returns and reads it back through -# temporal_to_taggstate, and a combine function answers one of the two states it takes -# while temporal_tagg_finalfn releases the other. +# helper written beside UdfMarshal like the MARSHAL template above. The roles the catalog +# states for the aggregate are handed in per temporal type, and the buffer holds the native +# state of the partition. When Spark moves a buffer between executors it writes the state as +# the bytes the serialize function writes and reads it back through the deserialize function, +# and the state a combine function does not answer is released through the final function, +# which consumes it. AGGREGATE = GEN_NOTE + """\ package org.mobilitydb.spark.generated; @@ -1553,62 +1553,78 @@ def emit_setret(name, cands, colnames): import org.mobilitydb.spark.MeosMemory; public final class TemporalAggregate extends Aggregator { - /** A transition or a combine function of the family: (state, value) or (state, state). */ + /** A transition or a combine function: (state, value) or (state, state). */ public interface Step extends Serializable { Pointer apply(Pointer a, Pointer b); } - /** The state of a partition, and the index of the temporal type it aggregates. */ + /** A final function, which consumes the state. */ + public interface Final extends Serializable { + Pointer apply(Pointer state); + } + + /** A serialize function: the bytes of a state, which is kept. */ + public interface Write extends Serializable { + byte[] apply(Pointer state); + } + + /** A deserialize function: the state the bytes hold. */ + public interface Read extends Serializable { + Pointer apply(byte[] form); + } + + /** The state of a partition, the index of the temporal type it aggregates, and the + * functions writing the state as bytes and reading it back when it leaves the executor. */ public static final class Buffer implements Serializable { transient Pointer state; int step = -1; + Write write; + Read read; private void writeObject(ObjectOutputStream out) throws IOException { out.defaultWriteObject(); - byte[] wkb = null; - if (state != null) { - Pointer t = GeneratedFunctions.temporal_tagg_finalfn(state); - state = null; - if (t != null) { - wkb = GeneratedFunctions.temporal_as_wkb(t, (byte) 4); - state = GeneratedFunctions.temporal_to_taggstate(t); - MeosMemory.free(t); - } - } - out.writeObject(wkb); + out.writeObject(state == null ? null : write.apply(state)); } private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException { in.defaultReadObject(); - byte[] wkb = (byte[]) in.readObject(); - state = null; - if (wkb != null) { - Pointer t = GeneratedFunctions.temporal_from_wkb(wkb); - state = GeneratedFunctions.temporal_to_taggstate(t); - MeosMemory.free(t); - } + byte[] form = (byte[]) in.readObject(); + state = form == null ? null : read.apply(form); } } private final int[] codes; private final Step[] trans; private final Step[] combs; + private final Final[] fins; + private final Write[] writes; + private final Read[] reads; - /** codes[i] is the WKB type byte trans[i] and combs[i] take, or -1 for every temporal. */ - public TemporalAggregate(int[] codes, Step[] trans, Step[] combs) { + /** codes[i] is the WKB type byte of the values the roles of index i take. */ + public TemporalAggregate(int[] codes, Step[] trans, Step[] combs, Final[] fins, Write[] writes, + Read[] reads) { this.codes = codes; this.trans = trans; this.combs = combs; + this.fins = fins; + this.writes = writes; + this.reads = reads; } private int pick(Object value) { int type = UdfMarshal.wkbType(value); for (int i = 0; i < codes.length; i++) - if (codes[i] < 0 || codes[i] == type) + if (codes[i] == type) return i; throw new IllegalArgumentException("the aggregate takes no temporal value of type " + type); } + /** Release a state by releasing what its final function answers. */ + private void release(int i, Pointer state) { + if (state != null) + MeosMemory.free(fins[i].apply(state)); + } + @Override public Buffer zero() { return new Buffer(); } @Override public Buffer reduce(Buffer b, byte[] wkb) { @@ -1624,12 +1640,17 @@ def emit_setret(name, cands, colnames): try { b.state = trans[i].apply(b.state, t); b.step = i; + b.write = writes[i]; + b.read = reads[i]; } finally { MeosMemory.free(t); } return b; } + /** The two partial states joined. A combine answers in the place of its first state, except + * that a skip-list combine answers its second state when the first is empty: the state it + * does not answer is released. */ @Override public Buffer merge(Buffer b1, Buffer b2) { if (b2.state == null) return b1; @@ -1638,8 +1659,7 @@ def emit_setret(name, cands, colnames): if (b1.step != b2.step) throw new IllegalArgumentException("the aggregate takes values of one temporal type"); Pointer r = combs[b1.step].apply(b1.state, b2.state); - Pointer other = r.address() == b1.state.address() ? b2.state : b1.state; - MeosMemory.free(GeneratedFunctions.temporal_tagg_finalfn(other)); + release(b1.step, r != null && r.address() == b2.state.address() ? b1.state : b2.state); b1.state = r; b2.state = null; return b1; @@ -1648,7 +1668,7 @@ def emit_setret(name, cands, colnames): @Override public byte[] finish(Buffer b) { if (b.state == null) return null; - Pointer t = GeneratedFunctions.temporal_tagg_finalfn(b.state); + Pointer t = fins[b.step].apply(b.state); b.state = null; if (t == null) return null; @@ -2049,65 +2069,61 @@ def main(): print(" @sqlfn canonical names : %d (%d arg-kind-dispatched, %d subtype-siblings + %d unsafe-overload names to C-name)" % (nsql, nsqldisp, ndropped, nskip), file=sys.stderr) - # ── temporal aggregate pass: the catalog's sqlAgg families as Spark aggregators ── - # An aggregate is the group of functions the catalog tags with one sqlAgg name. It is - # emitted when each transition function takes (SkipList *, Temporal *), has the combine - # function of the same stem, and the final function of the family is the generic - # temporal_tagg_finalfn: then a partial aggregate travels between executors as the - # temporal value that final function returns, and temporal_to_taggstate rebuilds the - # state from it (TemporalAggregate.java). A family whose own final function reads more - # than the temporal values (tAvg, tCentroid) keeps a state that value cannot carry, so it - # is left out and reported. The typed transition functions of one family are chosen by - # the WKB type byte of the value, like the typed overloads of a dispatcher; the generic - # ones (temporal_tcount_transfn) take every temporal. Spark resolves function names - # regardless of case and holds one function per name, so an aggregate whose name a - # scalar of the catalog or of this surface carries takes the `Agg` suffix of its - # canonical name (mergeAgg, tMinAgg), and the scalar keeps the bare name. + # ── temporal aggregate pass: the catalog's aggregates over temporal values as Spark aggregators ── + # An aggregate of the catalog's `aggregates` section states the public MEOS function of each + # role it defines. It is emitted when it takes one temporal value and answers one, states a + # combine, and keeps an internal state its serialize and deserialize functions write and + # read: a partition keeps the state between two values, and a partial aggregate travels + # between executors as the bytes the serialize function writes (TemporalAggregate.java). The + # overloads of one name are chosen by the WKB type byte of the value, like the typed + # overloads of a dispatcher. Spark resolves function names regardless of case and holds one + # function per name, so an aggregate whose name a scalar of the catalog or of this surface + # carries takes the `Agg` suffix of its canonical name: the catalog's own `Agg` name where it + # states one beside the bare one (mergeAgg, tMinAgg), else the bare name with the suffix + # (tAndAgg), and the scalar keeps the bare name. registered = {m.lower() for part in grouped.values() for code in part for m in re.findall(r'register\("([^"]+)"', code)} registered |= {f["sqlfn"].lower() for f in fns if isinstance(f.get("sqlfn"), str) and not f.get("sqlAgg")} have = lambda n: n in by_name and (jar_syms is None or n in jar_syms) - aggs = {} - for f in fns: - for a in f.get("sqlAgg") or []: - aggs.setdefault(a, []).append(f) - state_canon = lambda p: norm(p["canonical"]) == "SkipList *" - nagg, agg_left = 0, [] - for a in sorted(aggs): - grp = {f["name"]: f for f in aggs[a]} - finals = [n for n in grp if n.endswith("_finalfn")] - steps = [] - for n in sorted(grp): - f = grp[n] - ps = f.get("params") or [] - if not (n.endswith("_transfn") and len(ps) == 2 and state_canon(ps[0]) - and base(ps[1]["canonical"]) == "Temporal" - and norm(f["returnType"]["canonical"]) == "SkipList *"): + stated = {a["sqlName"].lower() for a in cat.get("aggregates") or []} + temporal = set(cat.get("temporalTypes") or {}) + aggs, agg_left = {}, set() + for a in cat.get("aggregates") or []: + role = lambda k: (a.get(k) or {}).get("meos") + name = a["sqlName"] + if name.lower() in registered: + if (name + "Agg").lower() in stated: continue - comb = n[:-len("_transfn")] + "_combinefn" - t = _expected_temptype(f) - steps.append((TEMPTYPE_CODE[t] if t else -1, n, comb)) - if (not steps or any(n != "temporal_tagg_finalfn" for n in finals) - or not all(have(n) and have(c) for _, n, c in steps) - or not (have("temporal_tagg_finalfn") and have("temporal_to_taggstate"))): - agg_left.append(a) + name += "Agg" + takes = a.get("args") or [] + roles = [role(k) for k in ("transition", "combine", "final", "serialize", "deserialize")] + if (len(takes) != 1 or takes[0] not in temporal or a.get("ret") not in temporal + or a.get("stype") != "internal" or not all(roles) or not all(map(have, roles))): + agg_left.add(name) continue - sname = a + "Agg" if a.lower() in registered else a - codes = ", ".join(str(c) for c, _, _ in steps) - trans = ", ".join("GeneratedFunctions::%s" % n for _, n, _ in steps) - combs = ", ".join("GeneratedFunctions::%s" % c for _, _, c in steps) + steps = aggs.setdefault(name, {}) + steps.setdefault(TEMPTYPE_CODE[takes[0]], roles) + nagg = 0 + for sname in sorted(aggs): + steps = sorted(aggs[sname].items()) + role = lambda i: ", ".join("GeneratedFunctions::%s" % r[i] for _, r in steps) grouped.setdefault("GeneratedUdfs_aggregate", []).append( ' spark.udf().register("%s", org.apache.spark.sql.functions.udaf(\n' ' new TemporalAggregate(new int[] {%s},\n' ' new TemporalAggregate.Step[] {%s},\n' - ' new TemporalAggregate.Step[] {%s}),\n' - ' org.apache.spark.sql.Encoders.BINARY()));' % (sname, codes, trans, combs)) + ' new TemporalAggregate.Step[] {%s},\n' + ' new TemporalAggregate.Final[] {%s},\n' + ' new TemporalAggregate.Write[] {%s},\n' + ' new TemporalAggregate.Read[] {%s}),\n' + ' org.apache.spark.sql.Encoders.BINARY()));' + % (sname, ", ".join(str(c) for c, _ in steps), role(0), role(1), role(2), role(3), + role(4))) registered.add(sname.lower()) nagg += 1 cov += 1 print(" temporal aggregates : %d (left out: %s)" - % (nagg, ", ".join(agg_left) or "none"), file=sys.stderr) + % (nagg, ", ".join(sorted(agg_left - set(aggs))) or "none"), file=sys.stderr) # Organize by doxygen module group (@ingroup), one class per group — the SAME # structure as the MEOS reference manual / XML docs, so a function is found in the diff --git a/tools/spark-udf-gaps.txt b/tools/spark-udf-gaps.txt index d4954e16c..2f6afd5f6 100644 --- a/tools/spark-udf-gaps.txt +++ b/tools/spark-udf-gaps.txt @@ -467,8 +467,6 @@ stbox_tmin arg:TimestampTz * stbox_to_box3d internal stbox_to_gbox internal stboxarr_round array:STBox * -taggstate_deserialize array-or-out-param:bytes; no-encoder:SkipList -taggstate_serialize no-decoder:SkipList; array-or-out-param:size_out; unsupported-return:uint8_t * tbigint_time_boxes array-or-out-param:count tbigint_value_bins array-or-out-param:count tbigint_value_boxes array-or-out-param:count @@ -495,6 +493,7 @@ temporal_split_n_spans array-or-out-param:count temporal_time_bins array-or-out-param:count temporal_timestamps array-or-out-param:count temporal_timestamptz_n arg:TimestampTz * +temporal_to_taggstate no-encoder:SkipList temporal_wcount_transfn no-decoder:SkipList; no-encoder:SkipList teq_posechain_tposechain arg:PoseChain * teq_text_ttext internal @@ -641,15 +640,11 @@ tne_text_ttext internal tne_tposechain_posechain arg:PoseChain * tne_ttext_text internal tnpoint_positions internal -tnpoint_tcentroid_transfn no-decoder:SkipList; no-encoder:SkipList tnpoint_value_at_timestamptz internal tnpoint_value_n internal tnpoint_values internal tnumber_split_each_n_tboxes array-or-out-param:count tnumber_split_n_tboxes array-or-out-param:count -tnumber_tavg_combinefn no-decoder:SkipList; no-encoder:SkipList -tnumber_tavg_finalfn no-decoder:SkipList -tnumber_tavg_transfn no-decoder:SkipList; no-encoder:SkipList tnumber_tboxes array-or-out-param:count tnumber_wavg_transfn no-decoder:SkipList; no-encoder:SkipList tpc_extent_transfn ret:TPCBox * @@ -690,8 +685,6 @@ tpcbox_ymax arg:TPCBox * tpcbox_ymin arg:TPCBox * tpcbox_zmax arg:TPCBox * tpcbox_zmin arg:TPCBox * -tpcpatch_tdensity_transfn no-decoder:SkipList; no-encoder:SkipList -tpcpatch_tnpoints_transfn no-decoder:SkipList; no-encoder:SkipList tpcpatch_value_at_timestamptz internal tpcpatch_value_n internal tpcpatch_values internal @@ -700,9 +693,6 @@ tpcpoint_value_n internal tpcpoint_values internal tpoint_as_mvtgeom array-or-out-param:gsarr; array-or-out-param:timesarr; array-or-out-param:count tpoint_make_simple internal -tpoint_tcentroid_combinefn no-decoder:SkipList; no-encoder:SkipList -tpoint_tcentroid_finalfn no-decoder:SkipList -tpoint_tcentroid_transfn no-decoder:SkipList; no-encoder:SkipList tpoint_tfloat_to_geomeas internal tpointseq_make_coords array-or-out-param:xcoords; array-or-out-param:ycoords; array-or-out-param:zcoords tpose_value_at_timestamptz internal