fix(isthmus): allow for conversion of plans containing Calcite aggregate functions

isthmus/src/main/java/io/substrait/isthmus/AggregateFunctions.java

@@ -1,5 +1,6 @@
 package io.substrait.isthmus;
 
+import java.util.Optional;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlKind;
@@ -22,6 +23,35 @@ public class AggregateFunctions {
   public static SqlAggFunction SUM = new SubstraitSumAggFunction();
   public static SqlAggFunction SUM0 = new SubstraitSumEmptyIsZeroAggFunction();
 
+  /**
+   * Some Calcite rules, like {@link
+   * org.apache.calcite.rel.rules.AggregateExpandDistinctAggregatesRule}, introduce the default
+   * Calcite aggregate functions into plans.
+   *
+   * <p>When converting these Calcite plans to Substrait, we need to convert the default Calcite
+   * aggregate calls to the Substrait specific variants.
+   *
+   * <p>This function attempts to convert the given {@code aggFunction} to its Substrait equivalent
+   *
+   * @param aggFunction the {@link SqlAggFunction} to convert to a Substrait specific variant
+   * @return an optional containing the Substrait equivalent of the given {@code aggFunction} if
+   *     conversion was needed, empty otherwise.
+   */
+  public static Optional<SqlAggFunction> toSubstraitAggVariant(SqlAggFunction aggFunction) {
+    if (aggFunction instanceof SqlMinMaxAggFunction fun) {
+      return Optional.of(
+          fun.getKind() == SqlKind.MIN ? AggregateFunctions.MIN : AggregateFunctions.MAX);
+    } else if (aggFunction instanceof SqlAvgAggFunction) {
+      return Optional.of(AggregateFunctions.AVG);
+    } else if (aggFunction instanceof SqlSumAggFunction) {
+      return Optional.of(AggregateFunctions.SUM);
+    } else if (aggFunction instanceof SqlSumEmptyIsZeroAggFunction) {
+      return Optional.of(AggregateFunctions.SUM0);
+    } else {
+      return Optional.empty();
+    }
+  }
+
   /** Extension of {@link SqlMinMaxAggFunction} that ALWAYS infers a nullable return type */
   private static class SubstraitSqlMinMaxAggFunction extends SqlMinMaxAggFunction {
     public SubstraitSqlMinMaxAggFunction(SqlKind kind) {

isthmus/src/main/java/io/substrait/isthmus/expression/AggregateFunctionConverter.java

@@ -6,6 +6,7 @@
 import io.substrait.expression.ExpressionCreator;
 import io.substrait.expression.FunctionArg;
 import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.AggregateFunctions;
 import io.substrait.isthmus.SubstraitRelVisitor;
 import io.substrait.isthmus.TypeConverter;
 import io.substrait.type.Type;
@@ -80,13 +81,7 @@ public Optional<AggregateFunctionInvocation> convert(
       AggregateCall call,
       Function<RexNode, Expression> topLevelConverter) {
 
-    // replace COUNT() + distinct == true and approximate == true with APPROX_COUNT_DISTINCT
-    // before converting into substrait function
-    SqlAggFunction aggFunction = call.getAggregation();
-    if (aggFunction == SqlStdOperatorTable.COUNT && call.isDistinct() && call.isApproximate()) {
-      aggFunction = SqlStdOperatorTable.APPROX_COUNT_DISTINCT;
-    }
-    FunctionFinder m = signatures.get(aggFunction);
+    var m = getFunctionFinder(call);
     if (m == null) {
       return Optional.empty();
     }
@@ -98,6 +93,21 @@ public Optional<AggregateFunctionInvocation> convert(
     return m.attemptMatch(wrapped, topLevelConverter);
   }
 
+  protected FunctionFinder getFunctionFinder(AggregateCall call) {
+    // replace COUNT() + distinct == true and approximate == true with APPROX_COUNT_DISTINCT
+    // before converting into substrait function
+    SqlAggFunction aggFunction = call.getAggregation();
+    if (aggFunction == SqlStdOperatorTable.COUNT && call.isDistinct() && call.isApproximate()) {
+      aggFunction = SqlStdOperatorTable.APPROX_COUNT_DISTINCT;
+    }
+
+    SqlAggFunction lookupFunction =
+        // Replace default Calcite aggregate calls with Substrait specific variants.
+        // See toSubstraitAggVariant for more details.
+        AggregateFunctions.toSubstraitAggVariant(aggFunction).orElse(aggFunction);
+    return signatures.get(lookupFunction);
+  }
+
   static class WrappedAggregateCall implements GenericCall {
     private final AggregateCall call;
     private final RelNode input;

isthmus/src/main/java/io/substrait/isthmus/expression/FunctionConverter.java

@@ -354,6 +354,14 @@ public Optional<T> attemptMatch(C call, Function<RexNode, Expression> topLevelCo
       }
       return Optional.empty();
     }
+
+    protected String getName() {
+      return name;
+    }
+
+    public SqlOperator getOperator() {
+      return operator;
+    }
   }
 
   public interface GenericCall {

isthmus/src/test/java/io/substrait/isthmus/OptimizerIntegrationTest.java

@@ -0,0 +1,51 @@
+package io.substrait.isthmus;
+
+import static io.substrait.isthmus.SqlConverterBase.EXTENSION_COLLECTION;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.io.IOException;
+import java.util.List;
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgram;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelRoot;
+import org.apache.calcite.rel.rules.CoreRules;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.junit.jupiter.api.Test;
+
+public class OptimizerIntegrationTest extends PlanTestBase {
+
+  @Test
+  void conversionHandlesBuiltInSum0CallAddedByRule() throws SqlParseException, IOException {
+    var query =
+        "select O_CUSTKEY, count(distinct O_ORDERKEY), count(*) from orders group by O_CUSTKEY";
+    // verify that the query works generally
+    assertFullRoundTrip(query);
+
+    SqlToSubstrait sqlConverter = new SqlToSubstrait();
+    List<RelRoot> relRoots = sqlConverter.sqlToRelNode(query, tpchSchemaCreateStatements());
+    assertEquals(1, relRoots.size());
+    RelRoot planRoot = relRoots.get(0);
+    RelNode originalPlan = planRoot.rel;
+
+    // Create a program to apply the AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN rule.
+    // This will introduce a SqlSumEmptyIsZeroAggFunction to the plan.
+    // This function does not have a mapping to Substrait.
+    // SubstraitSumEmptyIsZeroAggFunction is the variant which has a mapping.
+    // See io.substrait.isthmus.AggregateFunctions for details
+    HepProgram program =
+        new HepProgramBuilder()
+            .addRuleInstance(CoreRules.AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN)
+            .build();
+    HepPlanner planner = new HepPlanner(program);
+    planner.setRoot(originalPlan);
+    var newPlan = planner.findBestExp();
+
+    assertDoesNotThrow(
+        () ->
+            // Conversion of the new plan should succeed
+            SubstraitRelVisitor.convert(RelRoot.of(newPlan, planRoot.kind), EXTENSION_COLLECTION));
+  }
+}

isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java

@@ -32,7 +32,7 @@
 import org.junit.jupiter.api.Assertions;
 
 public class PlanTestBase {
-  final SimpleExtension.ExtensionCollection extensions;
+  protected final SimpleExtension.ExtensionCollection extensions;
 
   {
     try {

isthmus/src/test/java/io/substrait/isthmus/expression/AggregateFunctionConverterTest.java

@@ -0,0 +1,35 @@
+package io.substrait.isthmus.expression;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+import io.substrait.isthmus.AggregateFunctions;
+import io.substrait.isthmus.PlanTestBase;
+import io.substrait.isthmus.TypeConverter;
+import java.util.List;
+import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.junit.jupiter.api.Test;
+
+public class AggregateFunctionConverterTest extends PlanTestBase {
+
+  @Test
+  void testFunctionFinderMatch() {
+    AggregateFunctionConverter converter =
+        new AggregateFunctionConverter(
+            extensions.aggregateFunctions(), List.of(), typeFactory, TypeConverter.DEFAULT);
+
+    var functionFinder =
+        converter.getFunctionFinder(
+            AggregateCall.create(
+                new SqlSumEmptyIsZeroAggFunction(),
+                true,
+                List.of(1),
+                0,
+                typeFactory.createSqlType(SqlTypeName.VARCHAR),
+                null));
+    assertNotNull(functionFinder);
+    assertEquals("sum0", functionFinder.getName());
+    assertEquals(AggregateFunctions.SUM0, functionFinder.getOperator());
+  }
+}