[GLUTEN-12597][CORE] Migrate FetchRel offset/count to expressions (Substrait 0.98)

cpp-ch/local-engine/Parser/RelParsers/FetchRelParser.cpp

@@ -31,7 +31,9 @@ class FetchRelParser : public RelParser
     DB::QueryPlanPtr parse(DB::QueryPlanPtr query_plan, const substrait::Rel & rel, std::list<const substrait::Rel *> &)
     {
         const auto & limit = rel.fetch();
-        auto limit_step = std::make_unique<DB::LimitStep>(query_plan->getCurrentHeader(), limit.count(), limit.offset());
+        size_t count = limit.has_count_expr() ? limit.count_expr().literal().i64() : 0;
+        size_t offset = limit.has_offset_expr() ? limit.offset_expr().literal().i64() : 0;
+        auto limit_step = std::make_unique<DB::LimitStep>(query_plan->getCurrentHeader(), count, offset);
         limit_step->setStepDescription("LIMIT");
         steps.push_back(limit_step.get());
         query_plan->addStep(std::move(limit_step));

cpp-ch/local-engine/Parser/RelParsers/SortRelParser.cpp

@@ -61,7 +61,7 @@ size_t SortRelParser::parseLimit(std::list<const substrait::Rel *> & rel_stack_)
     if (last_rel.has_fetch())
     {
         const auto & fetch_rel = last_rel.fetch();
-        return fetch_rel.count();
+        return fetch_rel.has_count_expr() ? fetch_rel.count_expr().literal().i64() : 0;
     }
     return 0;
 }

cpp/velox/substrait/SubstraitToVeloxPlan.cc

@@ -1384,12 +1384,13 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
 
 core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::FetchRel& fetchRel) {
   auto childNode = convertSingleInput<::substrait::FetchRel>(fetchRel);
-  return std::make_shared<core::LimitNode>(
-      nextPlanNodeId(),
-      static_cast<int32_t>(fetchRel.offset()),
-      static_cast<int32_t>(fetchRel.count()),
-      false /*isPartial*/,
-      childNode);
+  int32_t offset = fetchRel.has_offset_expr()
+      ? static_cast<int32_t>(SubstraitParser::getLiteralValue<int64_t>(fetchRel.offset_expr().literal()))
+      : 0;
+  int32_t count = fetchRel.has_count_expr()
+      ? static_cast<int32_t>(SubstraitParser::getLiteralValue<int64_t>(fetchRel.count_expr().literal()))
+      : 0;
+  return std::make_shared<core::LimitNode>(nextPlanNodeId(), offset, count, false /*isPartial*/, childNode);
 }
 
 core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::TopNRel& topNRel) {

cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc

@@ -482,7 +482,11 @@ bool SubstraitToVeloxPlanValidator::validate(const ::substrait::FetchRel& fetchR
     }
   }
 
-  if (fetchRel.offset() < 0 || fetchRel.count() < 0) {
+  int64_t offset =
+      fetchRel.has_offset_expr() ? SubstraitParser::getLiteralValue<int64_t>(fetchRel.offset_expr().literal()) : 0;
+  int64_t count =
+      fetchRel.has_count_expr() ? SubstraitParser::getLiteralValue<int64_t>(fetchRel.count_expr().literal()) : 0;
+  if (offset < 0 || count < 0) {
     LOG_VALIDATION_MSG("Offset and count should be valid in FetchRel.");
     return false;
   }

cpp/velox/substrait/VeloxToSubstraitPlan.cc

@@ -356,8 +356,8 @@ void VeloxToSubstraitPlanConvertor::toSubstrait(
 
   VELOX_CHECK(!topNNode->isPartial(), "Substrait doesn't support partial topN yet");
 
-  fetchRel->set_offset(0);
-  fetchRel->set_count(topNNode->count());
+  fetchRel->mutable_offset_expr()->mutable_literal()->set_i64(0);
+  fetchRel->mutable_count_expr()->mutable_literal()->set_i64(topNNode->count());
   fetchRel->mutable_common()->mutable_direct();
 }
 
@@ -388,8 +388,8 @@ void VeloxToSubstraitPlanConvertor::toSubstrait(
   const auto& source = getSingleSource(limitNode);
   toSubstrait(arena, source, fetchRel->mutable_input());
 
-  fetchRel->set_offset(limitNode->offset());
-  fetchRel->set_count(limitNode->count());
+  fetchRel->mutable_offset_expr()->mutable_literal()->set_i64(limitNode->offset());
+  fetchRel->mutable_count_expr()->mutable_literal()->set_i64(limitNode->count());
 
   VELOX_CHECK(!limitNode->isPartial(), "Substrait doesn't support partial limit yet");
 

gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java

@@ -16,6 +16,7 @@
  */
 package org.apache.gluten.substrait.rel;
 
+import org.apache.gluten.substrait.expression.ExpressionBuilder;
 import org.apache.gluten.substrait.extensions.AdvancedExtensionNode;
 
 import io.substrait.proto.FetchRel;
@@ -57,8 +58,8 @@ public Rel toProtobuf() {
     if (input != null) {
       fetchRelBuilder.setInput(input.toProtobuf());
     }
-    fetchRelBuilder.setOffset(offset);
-    fetchRelBuilder.setCount(count);
+    fetchRelBuilder.setOffsetExpr(ExpressionBuilder.makeLongLiteral(offset).toProtobuf());
+    fetchRelBuilder.setCountExpr(ExpressionBuilder.makeLongLiteral(count).toProtobuf());
 
     if (extensionNode != null) {
       fetchRelBuilder.setAdvancedExtension(extensionNode.toProtobuf());

gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto

@@ -319,12 +319,23 @@ message CrossRel {
 
 // The relational operator representing LIMIT/OFFSET or TOP type semantics.
 message FetchRel {
+  reserved 3, 4;
+  reserved "offset", "count";
+
   RelCommon common = 1;
   Rel input = 2;
-  // the offset expressed in number of records
-  int64 offset = 3;
-  // the amount of records to return
-  int64 count = 4;
+  // Expression evaluated into a non-negative integer specifying the number
+  // of records to skip. An expression evaluating to null is treated as 0.
+  // Evaluating to a negative integer should result in an error.
+  // Recommended type for offset is int64. Unset is treated as 0.
+  Expression offset_expr = 5;
+  // Expression evaluated into a non-negative integer specifying the number
+  // of records to return. An expression evaluating to null signals that ALL
+  // records should be returned.
+  // Evaluating to a negative integer should result in an error.
+  // Recommended type for count is int64. Unset signals that ALL records
+  // should be returned.
+  Expression count_expr = 6;
   substrait.extensions.AdvancedExtension advanced_extension = 10;
 }
 

gluten-substrait/src/test/scala/org/apache/gluten/utils/FetchRelProtoSuite.scala

@@ -0,0 +1,43 @@
+/*
+ * 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.gluten.utils
+
+import org.apache.gluten.substrait.SubstraitContext
+import org.apache.gluten.substrait.rel.RelBuilder
+
+import org.scalatest.funsuite.AnyFunSuite
+
+/**
+ * Locks the FetchRel producer contract after the Substrait 0.98 migration. 0.98 removed the scalar
+ * `int64 offset = 3` / `int64 count = 4` fields in favor of `Expression offset_expr = 5` /
+ * `Expression count_expr = 6`. Gluten's only FetchRel producer feeds literal `Long`s (from Spark
+ * Limit/Offset), so the producer now wraps each into an i64-literal `Expression`. This suite pins
+ * that the values land in the new expression carriers as i64 literals.
+ */
+class FetchRelProtoSuite extends AnyFunSuite {
+
+  test("makeFetchRel emits offset/count as i64 literal expressions") {
+    val context = new SubstraitContext
+    val rel = RelBuilder.makeFetchRel(null, 5L, 10L, context, 0L)
+    val fetchRel = rel.toProtobuf.getFetch
+
+    assert(fetchRel.hasOffsetExpr)
+    assert(fetchRel.hasCountExpr)
+    assert(fetchRel.getOffsetExpr.getLiteral.getI64 === 5L)
+    assert(fetchRel.getCountExpr.getLiteral.getI64 === 10L)
+  }
+}