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