feat(isthmus)!: preserve aggregate output types and semantics through Calcite

core/src/main/java/io/substrait/extension/InvalidFunctionBindingException.java

@@ -0,0 +1,40 @@
+package io.substrait.extension;
+
+import io.substrait.type.Type;
+
+/**
+ * Thrown when an extension function invocation cannot be resolved, or when the output type declared
+ * by a plan is inconsistent with the type derived from the function's extension declaration.
+ *
+ * <p>Resolution is fail-closed: rather than trusting a plan-supplied output type, the resolver
+ * derives the type from the declaration and raises this exception when they disagree.
+ */
+public class InvalidFunctionBindingException extends RuntimeException {
+
+  private static final long serialVersionUID = 1L;
+
+  /**
+   * Creates an exception with the given message.
+   *
+   * @param message the detail message
+   */
+  public InvalidFunctionBindingException(String message) {
+    super(message);
+  }
+
+  /**
+   * Creates an exception describing a mismatch between the declared and the derived output type.
+   *
+   * @param anchor the function anchor being resolved
+   * @param declared the output type declared by the plan
+   * @param derived the output type derived from the declaration
+   * @return the exception
+   */
+  public static InvalidFunctionBindingException outputTypeMismatch(
+      SimpleExtension.FunctionAnchor anchor, Type declared, Type derived) {
+    return new InvalidFunctionBindingException(
+        String.format(
+            "Declared output type %s for %s does not match the type %s derived from its declaration",
+            declared, anchor, derived));
+  }
+}

core/src/main/java/io/substrait/extension/ResolvedAggregateBinding.java

@@ -0,0 +1,154 @@
+package io.substrait.extension;
+
+import io.substrait.expression.AggregateFunctionInvocation;
+import io.substrait.expression.EnumArg;
+import io.substrait.expression.Expression;
+import io.substrait.expression.FunctionArg;
+import io.substrait.type.Type;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import org.immutables.value.Value;
+
+/**
+ * A fully-resolved binding of an aggregate function invocation.
+ *
+ * <p>Wraps a {@link ResolvedFunctionBinding} with the aggregate-specific phase and invocation
+ * semantics. The output type is derived from the function declaration (never from a plan-supplied
+ * value) and depends on the phase, so it does not participate in the binding's identity.
+ */
+@Value.Immutable
+public abstract class ResolvedAggregateBinding {
+
+  /**
+   * Returns the resolved function binding.
+   *
+   * @return the function binding
+   */
+  public abstract ResolvedFunctionBinding function();
+
+  /**
+   * Returns the aggregation phase.
+   *
+   * @return the aggregation phase
+   */
+  public abstract Expression.AggregationPhase phase();
+
+  /**
+   * Returns the aggregation invocation semantics (all vs. distinct).
+   *
+   * @return the aggregation invocation
+   */
+  public abstract Expression.AggregationInvocation invocation();
+
+  /**
+   * Returns the intermediate type, when it has been resolved explicitly. When empty, {@link
+   * #outputType()} derives it from the declaration on demand. Not part of the binding identity.
+   *
+   * @return the intermediate type, if known
+   */
+  @Value.Auxiliary
+  public abstract Optional<Type> intermediateType();
+
+  /**
+   * Returns the type this aggregate produces <em>in its phase</em>: a phase that stops at the
+   * intermediate state produces the declaration's intermediate type, while a phase that runs to the
+   * result produces its return type. Derived from the declaration, never from a plan-supplied
+   * value.
+   *
+   * @return the output type of this phase
+   * @throws InvalidFunctionBindingException if an intermediate type is required but the declaration
+   *     is not an aggregate variant, or its type expression cannot be derived
+   */
+  public Type outputType() {
+    return producesIntermediateState()
+        ? intermediateType().orElseGet(this::deriveIntermediateType)
+        : function().outputType();
+  }
+
+  /**
+   * Returns whether this phase stops at the declaration's intermediate state instead of producing
+   * its result.
+   *
+   * @return {@code true} for the initial-to-intermediate and intermediate-to-intermediate phases
+   */
+  public boolean producesIntermediateState() {
+    return phase() == Expression.AggregationPhase.INITIAL_TO_INTERMEDIATE
+        || phase() == Expression.AggregationPhase.INTERMEDIATE_TO_INTERMEDIATE;
+  }
+
+  /**
+   * Returns whether this phase consumes the declaration's intermediate state instead of its
+   * declared arguments. The arguments of such an invocation are intermediate values, so they are
+   * not expected to match the declaration's argument list.
+   *
+   * @return {@code true} for the intermediate-to-intermediate and intermediate-to-result phases
+   */
+  public boolean consumesIntermediateState() {
+    return phase() == Expression.AggregationPhase.INTERMEDIATE_TO_INTERMEDIATE
+        || phase() == Expression.AggregationPhase.INTERMEDIATE_TO_RESULT;
+  }
+
+  private Type deriveIntermediateType() {
+    SimpleExtension.Function declaration = function().declaration();
+    if (!(declaration instanceof SimpleExtension.AggregateFunctionVariant)) {
+      throw new InvalidFunctionBindingException(
+          String.format(
+              "%s is not an aggregate declaration and has no intermediate type",
+              function().anchor()));
+    }
+    return FunctionBindingResolver.deriveIntermediateType(
+        (SimpleExtension.AggregateFunctionVariant) declaration, function().arguments());
+  }
+
+  /**
+   * Creates a builder for {@link ResolvedAggregateBinding}.
+   *
+   * @return a new builder
+   */
+  public static ImmutableResolvedAggregateBinding.Builder builder() {
+    return ImmutableResolvedAggregateBinding.builder();
+  }
+
+  /**
+   * Resolves an aggregate function invocation into a binding, capturing its semantic identity
+   * (anchor, arguments and options). This does <em>not</em> validate the invocation against the
+   * declaration; signature, options and output-type validation is a separate, opt-in step (see
+   * {@link FunctionBindingResolver#validate}). The output type is derived on demand via {@link
+   * #outputType()}.
+   *
+   * @param invocation the aggregate invocation to resolve
+   * @return the resolved aggregate binding
+   */
+  public static ResolvedAggregateBinding resolve(AggregateFunctionInvocation invocation) {
+    List<ResolvedArgument> arguments = resolvedArguments(invocation.arguments());
+    ResolvedFunctionBinding function =
+        FunctionBindingResolver.resolve(invocation.declaration(), arguments, invocation.options());
+    return builder()
+        .function(function)
+        .phase(invocation.aggregationPhase())
+        .invocation(invocation.invocation())
+        .intermediateType(Optional.empty())
+        .build();
+  }
+
+  private static List<ResolvedArgument> resolvedArguments(List<FunctionArg> arguments) {
+    List<ResolvedArgument> resolved = new ArrayList<>();
+    for (FunctionArg arg : arguments) {
+      if (arg instanceof Expression) {
+        resolved.add(ResolvedArgument.value(((Expression) arg).getType()));
+      } else if (arg instanceof Type) {
+        resolved.add(ResolvedArgument.type((Type) arg));
+      } else if (arg instanceof EnumArg) {
+        // An enum argument may carry no option; keep that distinct from an empty one so the
+        // identity is faithful and validation can report it against the declaration.
+        resolved.add(
+            ((EnumArg) arg)
+                .value()
+                .map(ResolvedArgument::enumOption)
+                .orElseGet(ResolvedArgument::unspecifiedEnumOption));
+      }
+    }
+    return resolved;
+  }
+}

core/src/main/java/io/substrait/extension/ResolvedArgument.java

@@ -0,0 +1,142 @@
+package io.substrait.extension;
+
+import io.substrait.type.Type;
+import java.util.Objects;
+import java.util.Optional;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * An ordered, kind-aware argument of a resolved function binding.
+ *
+ * <p>Unlike a bare {@link Type} list, this preserves the argument <em>kind</em> (value, type or
+ * enum) and the selected enum option, so that two invocations that differ only by an enum argument
+ * — e.g. {@code std_dev(POPULATION, fp32)} vs {@code std_dev(SAMPLE, fp32)} — are not conflated.
+ *
+ * <p>An enum argument may also carry <em>no</em> option, for a plan that leaves it unspecified;
+ * that is distinct from any specified option. An option is kept exactly as the plan spelled it, so
+ * identity is case-sensitive and merely conservative: two invocations differing only in the case of
+ * an enum symbol stay distinct, which costs a missed deduplication and never rewrites the plan's
+ * data. Matching an option against its declaration is case-insensitive, as the spec requires.
+ */
+public final class ResolvedArgument {
+
+  /** The kind of a function argument. */
+  public enum Kind {
+    /** A value argument, carrying a data {@link Type}. */
+    VALUE,
+    /** A type argument, carrying a {@link Type}. */
+    TYPE,
+    /** An enum argument, carrying a selected option. */
+    ENUM
+  }
+
+  private final Kind kind;
+  private final @Nullable Type type;
+  private final @Nullable String enumValue;
+
+  private ResolvedArgument(Kind kind, @Nullable Type type, @Nullable String enumValue) {
+    this.kind = kind;
+    this.type = type;
+    this.enumValue = enumValue;
+  }
+
+  /**
+   * Creates a value argument.
+   *
+   * @param type the value type
+   * @return the resolved argument
+   */
+  public static ResolvedArgument value(Type type) {
+    return new ResolvedArgument(Kind.VALUE, Objects.requireNonNull(type), null);
+  }
+
+  /**
+   * Creates a type argument.
+   *
+   * @param type the argument type
+   * @return the resolved argument
+   */
+  public static ResolvedArgument type(Type type) {
+    return new ResolvedArgument(Kind.TYPE, Objects.requireNonNull(type), null);
+  }
+
+  /**
+   * Creates an enum argument.
+   *
+   * @param option the selected enum option
+   * @return the resolved argument
+   */
+  public static ResolvedArgument enumOption(String option) {
+    return new ResolvedArgument(Kind.ENUM, null, Objects.requireNonNull(option));
+  }
+
+  /**
+   * Creates an enum argument whose option the plan left unspecified. This is distinct from an
+   * argument carrying an empty option, and is only valid where the declaration does not require an
+   * option.
+   *
+   * @return the resolved argument
+   */
+  public static ResolvedArgument unspecifiedEnumOption() {
+    return new ResolvedArgument(Kind.ENUM, null, null);
+  }
+
+  /**
+   * Returns the argument kind.
+   *
+   * @return the kind
+   */
+  public Kind kind() {
+    return kind;
+  }
+
+  /**
+   * Returns the data type of a value or type argument.
+   *
+   * @return the type, if this is a value or type argument
+   */
+  public Optional<Type> type() {
+    return Optional.ofNullable(type);
+  }
+
+  /**
+   * Returns the selected option of an enum argument.
+   *
+   * @return the enum option, or empty if this is not an enum argument or its option was left
+   *     unspecified
+   */
+  public Optional<String> enumValue() {
+    return Optional.ofNullable(enumValue);
+  }
+
+  @Override
+  public boolean equals(@Nullable Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof ResolvedArgument)) {
+      return false;
+    }
+    ResolvedArgument other = (ResolvedArgument) o;
+    return kind == other.kind
+        && Objects.equals(type, other.type)
+        && Objects.equals(enumValue, other.enumValue);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(kind, type, enumValue);
+  }
+
+  @Override
+  public String toString() {
+    switch (kind) {
+      case VALUE:
+        return "value(" + type + ")";
+      case TYPE:
+        return "type(" + type + ")";
+      default:
+        return "enum(" + (enumValue == null ? "unspecified" : enumValue) + ")";
+    }
+  }
+}