[WIP] Rust Acceleration Kernel POC

pinot-core/pom.xml

@@ -106,6 +106,11 @@
       <groupId>org.apache.pinot</groupId>
       <artifactId>pinot-timeseries-spi</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.pinot</groupId>
+      <artifactId>pinot-native</artifactId>
+      <version>${project.version}</version>
+    </dependency>
 
     <dependency>
       <groupId>org.apache.pinot</groupId>

pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java

@@ -41,6 +41,8 @@
 import org.apache.pinot.core.query.aggregation.function.AggregationFunctionUtils.AggregationInfo;
 import org.apache.pinot.core.query.aggregation.groupby.DefaultGroupByExecutor;
 import org.apache.pinot.core.query.aggregation.groupby.GroupByExecutor;
+import org.apache.pinot.core.query.aggregation.groupby.NativeGroupByExecutor;
+import org.apache.pinot.core.query.aggregation.groupby.NativeGroupByRouter;
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.core.startree.executor.StarTreeGroupByExecutor;
 import org.apache.pinot.core.util.GroupByUtils;
@@ -108,6 +110,8 @@ protected GroupByResultsBlock getNextBlock() {
     GroupByExecutor groupByExecutor;
     if (_useStarTree) {
       groupByExecutor = new StarTreeGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator);
+    } else if (NativeGroupByRouter.shouldAccelerate(_queryContext, _groupByExpressions, _projectOperator)) {
+      groupByExecutor = new NativeGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator);
     } else {
       groupByExecutor = new DefaultGroupByExecutor(_queryContext, _groupByExpressions, _projectOperator);
     }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java

@@ -67,6 +67,12 @@ public static AggregationFunction getAggregationFunction(FunctionContext functio
       List<ExpressionContext> arguments = function.getArguments();
       int numArguments = arguments.size();
       ExpressionContext firstArgument = arguments.get(0);
+      // Single integration point for the native (Rust+JNI) acceleration path. See
+      // NativeAggregationRouter for eligibility rules. When disabled or ineligible,
+      // construction falls through to the standard Java factory branches below.
+      if (NativeAggregationRouter.shouldAccelerate(upperCaseFunctionName, arguments, nullHandlingEnabled)) {
+        return NativeAggregationRouter.createNative(upperCaseFunctionName, arguments, nullHandlingEnabled);
+      }
       if (upperCaseFunctionName.startsWith("PERCENTILE")) {
         String remainingFunctionName = upperCaseFunctionName.substring(10);
         if (remainingFunctionName.equals("SMARTTDIGEST")) {

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NativeAggregationRouter.java

@@ -0,0 +1,131 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.function;
+
+import java.util.List;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.nativeengine.agg.PinotNativeAgg;
+
+
+/**
+ * Routes eligible aggregation function constructions to the native (Rust+JNI) acceleration
+ * path. This is the single integration point between Pinot's aggregation engine and
+ * {@code pinot-native}.
+ *
+ * <p>Eligibility checks (short-circuiting):
+ * <ol>
+ *   <li>Feature flag {@code pinot.native.aggregation.enabled} is set to {@code true}</li>
+ *   <li>The native library loaded successfully ({@link PinotNativeAgg#isAvailable()})</li>
+ *   <li>Null handling is disabled (no native null path yet)</li>
+ *   <li>Function type is in the Phase 1.B scope: {@code SUM} / {@code SUM0} / {@code MIN} /
+ *       {@code MAX} / {@code COUNT}</li>
+ *   <li>The aggregated expression is a simple column identifier (no transforms) — for COUNT
+ *       this includes {@code *}</li>
+ * </ol>
+ *
+ * <p>When any check fails the caller must construct the original Java AggregationFunction
+ * unchanged. The router never throws.
+ *
+ * <p>Routing here covers all aggregation contexts in Pinot — SSE V1, MSE leaf (which
+ * delegates to V1), MSE intermediate, star-tree, realtime consuming segments, and
+ * Materialized View refresh — because every one of them obtains AggregationFunction
+ * instances from {@link AggregationFunctionFactory}. See {@code docs/native/phase-1-design.md}
+ * §2 for the engine landscape.
+ */
+public final class NativeAggregationRouter {
+
+  /** System property gating the native engine. Default {@code false}. */
+  public static final String ENABLED_PROPERTY = "pinot.native.aggregation.enabled";
+
+  private NativeAggregationRouter() {
+  }
+
+  /**
+   * @return {@code true} if a native AggregationFunction should be constructed for the
+   *         given function name + arguments. Caller must use {@link #createNative} to do so.
+   */
+  public static boolean shouldAccelerate(
+      String upperCaseFunctionName, List<ExpressionContext> arguments,
+      boolean nullHandlingEnabled) {
+    if (!enabled()) {
+      return false;
+    }
+    if (!PinotNativeAgg.isAvailable()) {
+      return false;
+    }
+    if (nullHandlingEnabled) {
+      return false;
+    }
+    if (!isInScopeFunction(upperCaseFunctionName)) {
+      return false;
+    }
+    return isSimpleColumnArg(arguments);
+  }
+
+  /**
+   * Constructs a native AggregationFunction for a name + arguments combination previously
+   * accepted by {@link #shouldAccelerate}. Caller is responsible for the eligibility check.
+   *
+   * @throws IllegalStateException if the function name is not in scope (programming error)
+   */
+  @SuppressWarnings("rawtypes")
+  public static AggregationFunction createNative(
+      String upperCaseFunctionName, List<ExpressionContext> arguments,
+      boolean nullHandlingEnabled) {
+    switch (upperCaseFunctionName) {
+      case "SUM":
+      case "SUM0":
+        return new NativeSumAggregationFunction(arguments, nullHandlingEnabled);
+      case "MIN":
+        return new NativeMinAggregationFunction(arguments, nullHandlingEnabled);
+      case "MAX":
+        return new NativeMaxAggregationFunction(arguments, nullHandlingEnabled);
+      case "COUNT":
+        return new NativeCountAggregationFunction(arguments, nullHandlingEnabled);
+      default:
+        throw new IllegalStateException(
+            "Native AggregationFunction requested for unsupported function: "
+                + upperCaseFunctionName);
+    }
+  }
+
+  static boolean enabled() {
+    return Boolean.getBoolean(ENABLED_PROPERTY);
+  }
+
+  private static boolean isInScopeFunction(String upperCaseFunctionName) {
+    switch (upperCaseFunctionName) {
+      case "SUM":
+      case "SUM0":
+      case "MIN":
+      case "MAX":
+      case "COUNT":
+        return true;
+      default:
+        return false;
+    }
+  }
+
+  private static boolean isSimpleColumnArg(List<ExpressionContext> arguments) {
+    if (arguments.size() != 1) {
+      return false;
+    }
+    return arguments.get(0).getType() == ExpressionContext.Type.IDENTIFIER;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NativeCountAggregationFunction.java

@@ -0,0 +1,70 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.function;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
+
+
+/**
+ * COUNT aggregation accelerated by the native engine — but notably without a kernel.
+ *
+ * <p>Under {@link NativeAggregationRouter}'s gating (null handling disabled, simple
+ * column identifier or absent expression), both {@code COUNT(*)} and {@code COUNT(col)}
+ * collapse to "block length added to the result holder". Crossing JNI for this would
+ * add the ~85 ns FFI fixed cost from §11.A of the design doc with zero kernel benefit.
+ *
+ * <p>So the "native" path here is just: {@code holder.setValue(holder.getDoubleResult() + length)},
+ * inline in Java. The class still exists so the routing infra in
+ * {@link NativeAggregationRouter} can be uniform across SUM/MIN/MAX/COUNT — every routed
+ * function gets its own {@code Native*AggregationFunction} class, even when the
+ * implementation is trivial.
+ *
+ * <p>Star-tree pre-aggregated case (where {@code blockValSetMap} is not empty and
+ * {@code nullHandlingEnabled=false} but the expression is the star-tree-internal
+ * {@code COUNT_STAR_EXPRESSION}, requiring a {@code long[]} sum) and null-handling case
+ * both fall through to the Java parent. {@link NativeAggregationRouter#shouldAccelerate}
+ * already disqualifies them, but we double-guard here for safety.
+ */
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class NativeCountAggregationFunction extends CountAggregationFunction {
+
+  public NativeCountAggregationFunction(List<ExpressionContext> arguments,
+      boolean nullHandlingEnabled) {
+    super(arguments, nullHandlingEnabled);
+  }
+
+  @Override
+  public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
+      Map<ExpressionContext, BlockValSet> blockValSetMap) {
+    // Fast path: COUNT(*) — blockValSetMap is empty per CountAggregationFunction's contract,
+    // so the value to add IS the block length.
+    if (blockValSetMap.isEmpty()) {
+      aggregationResultHolder.setValue(aggregationResultHolder.getDoubleResult() + length);
+      return;
+    }
+    // Anything else (null-handling enabled, star-tree pre-aggregated) — defer to Java.
+    // NativeAggregationRouter.shouldAccelerate should disqualify these at construction time,
+    // but we double-guard.
+    super.aggregate(length, aggregationResultHolder, blockValSetMap);
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/NativeMaxAggregationFunction.java

@@ -0,0 +1,77 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.query.aggregation.function;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
+import org.apache.pinot.nativeengine.agg.PinotNativeAgg;
+
+
+/**
+ * MAX aggregation accelerated by the native (Rust+JNI) engine. Mirror of
+ * {@link NativeMinAggregationFunction} for {@code MAX} semantics: dispatches
+ * single-value INT/LONG/FLOAT/DOUBLE columns to the matching {@link PinotNativeAgg}
+ * kernel; other type / encoding combinations fall through to the Java parent.
+ *
+ * <p>NaN semantics match the parent's {@code Math.max(...)} — NaN propagates.
+ * Empty input returns {@link Double#NEGATIVE_INFINITY}.
+ */
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class NativeMaxAggregationFunction extends MaxAggregationFunction {
+
+  public NativeMaxAggregationFunction(List<ExpressionContext> arguments,
+      boolean nullHandlingEnabled) {
+    super(arguments, nullHandlingEnabled);
+  }
+
+  @Override
+  public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
+      Map<ExpressionContext, BlockValSet> blockValSetMap) {
+    BlockValSet blockValSet = blockValSetMap.get(_expression);
+    if (blockValSet.isSingleValue()) {
+      double blockMax = Double.NaN;
+      switch (blockValSet.getValueType().getStoredType()) {
+        case INT:
+          blockMax = PinotNativeAgg.maxInt(blockValSet.getIntValuesSV(), length);
+          break;
+        case LONG:
+          blockMax = PinotNativeAgg.maxLong(blockValSet.getLongValuesSV(), length);
+          break;
+        case FLOAT:
+          blockMax = PinotNativeAgg.maxFloat(blockValSet.getFloatValuesSV(), length);
+          break;
+        case DOUBLE:
+          blockMax = PinotNativeAgg.maxDouble(blockValSet.getDoubleValuesSV(), length);
+          break;
+        default:
+          break;
+      }
+      // See NativeMinAggregationFunction for the NaN-vs-error rationale.
+      if (!Double.isNaN(blockMax)) {
+        double prev = aggregationResultHolder.getDoubleResult();
+        aggregationResultHolder.setValue(Math.max(blockMax, prev));
+        return;
+      }
+    }
+    super.aggregate(length, aggregationResultHolder, blockValSetMap);
+  }
+}