diff --git a/cpp-ch/local-engine/Parser/RelParsers/FetchRelParser.cpp b/cpp-ch/local-engine/Parser/RelParsers/FetchRelParser.cpp index 1b4c5e037e3..bd5b5f1dc3c 100644 --- a/cpp-ch/local-engine/Parser/RelParsers/FetchRelParser.cpp +++ b/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 auto & limit = rel.fetch(); - auto limit_step = std::make_unique(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(query_plan->getCurrentHeader(), count, offset); limit_step->setStepDescription("LIMIT"); steps.push_back(limit_step.get()); query_plan->addStep(std::move(limit_step)); diff --git a/cpp-ch/local-engine/Parser/RelParsers/SortRelParser.cpp b/cpp-ch/local-engine/Parser/RelParsers/SortRelParser.cpp index ce964da04b9..30e7c8db892 100644 --- a/cpp-ch/local-engine/Parser/RelParsers/SortRelParser.cpp +++ b/cpp-ch/local-engine/Parser/RelParsers/SortRelParser.cpp @@ -61,7 +61,7 @@ size_t SortRelParser::parseLimit(std::list & 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; } diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc index 9e85daa2064..8975d556c1b 100644 --- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc +++ b/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( - nextPlanNodeId(), - static_cast(fetchRel.offset()), - static_cast(fetchRel.count()), - false /*isPartial*/, - childNode); + int32_t offset = fetchRel.has_offset_expr() + ? static_cast(SubstraitParser::getLiteralValue(fetchRel.offset_expr().literal())) + : 0; + int32_t count = fetchRel.has_count_expr() + ? static_cast(SubstraitParser::getLiteralValue(fetchRel.count_expr().literal())) + : 0; + return std::make_shared(nextPlanNodeId(), offset, count, false /*isPartial*/, childNode); } core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::TopNRel& topNRel) { diff --git a/cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc b/cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc index 8a55d623070..bbce17159ad 100644 --- a/cpp/velox/substrait/SubstraitToVeloxPlanValidator.cc +++ b/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(fetchRel.offset_expr().literal()) : 0; + int64_t count = + fetchRel.has_count_expr() ? SubstraitParser::getLiteralValue(fetchRel.count_expr().literal()) : 0; + if (offset < 0 || count < 0) { LOG_VALIDATION_MSG("Offset and count should be valid in FetchRel."); return false; } diff --git a/cpp/velox/substrait/VeloxToSubstraitPlan.cc b/cpp/velox/substrait/VeloxToSubstraitPlan.cc index a1378724117..f21ba82dabf 100644 --- a/cpp/velox/substrait/VeloxToSubstraitPlan.cc +++ b/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"); diff --git a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java index 91f5a00ce56..f83ef143121 100644 --- a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/FetchRelNode.java +++ b/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()); diff --git a/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto b/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto index 6619a0395fc..4b8ad721473 100644 --- a/gluten-substrait/src/main/resources/substrait/proto/substrait/algebra.proto +++ b/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; } diff --git a/gluten-substrait/src/test/scala/org/apache/gluten/utils/FetchRelProtoSuite.scala b/gluten-substrait/src/test/scala/org/apache/gluten/utils/FetchRelProtoSuite.scala new file mode 100644 index 00000000000..4bcd1101fcd --- /dev/null +++ b/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) + } +}