diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/GlutenStrideExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/GlutenStrideExecTransformer.scala new file mode 100644 index 00000000000..7e0feed7e25 --- /dev/null +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/GlutenStrideExecTransformer.scala @@ -0,0 +1,112 @@ +/* + * 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.execution + +import org.apache.gluten.metrics.MetricsUpdater +import org.apache.gluten.substrait.SubstraitContext +import org.apache.gluten.substrait.rel.{RelBuilder, RelNode} + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.execution.SparkPlan + +import scala.collection.JavaConverters._ + +/** + * Gluten-native "stride" operator: outputs every N-th row from each input batch. + * + * Row indices within each batch are 0-based; index 0 is always included. The counter resets per + * batch, so the operator is completely stateless and parallelism-friendly. + * + * This is a self-contained example of the Gluten custom-operator mechanism. It does NOT correspond + * to any existing Spark operator -- it exists purely to demonstrate how to add a backend-specific + * native operator without requiring a Velox upstream contribution. + * + * ==End-to-end flow== + * {{{ + * GlutenStrideExecTransformer(stride=3, child=scanPlan) + * --doTransform--> GlutenStrideRelNode(stride=3) + * --toProtobuf--> FetchRel { offset=3, advanced_extension.optimization="isGlutenStride=1" } + * --JNI--> SubstraitToVeloxPlanConverter::toVeloxPlan(FetchRel&) + * --builds--> GlutenStrideNode(stride=3, child) + * --exec--> GlutenStrideOperator: keeps rows 0, 3, 6, 9, ... + * }}} + * + * ==Usage== + * Instantiate directly in a unit test or wire it into a custom offload rule: + * {{{ + * val strider = GlutenStrideExecTransformer(stride = 3L, child = childPlan) + * }}} + */ +case class GlutenStrideExecTransformer(stride: Long, child: SparkPlan) + extends UnaryTransformSupport { + + require(stride >= 1L, s"stride must be >= 1, got $stride") + + // ------------------------------------------------------------------------- + // SparkPlan identity + // ------------------------------------------------------------------------- + + override def output: Seq[Attribute] = child.output + + override def metricsUpdater(): MetricsUpdater = MetricsUpdater.None + + override protected def withNewChildInternal(newChild: SparkPlan): GlutenStrideExecTransformer = + copy(child = newChild) + + // ------------------------------------------------------------------------- + // Validation + // ------------------------------------------------------------------------- + + override protected def doValidateInternal(): ValidationResult = { + val context = new SubstraitContext + val operatorId = context.nextOperatorId(this.nodeName) + val relNode = makeRelNode(context, operatorId, inputRelNode = null, validation = true) + doNativeValidation(context, relNode) + } + + // ------------------------------------------------------------------------- + // Transformation + // ------------------------------------------------------------------------- + + override protected def doTransform(context: SubstraitContext): TransformContext = { + val childCtx = child.asInstanceOf[TransformSupport].transform(context) + val operatorId = context.nextOperatorId(this.nodeName) + val relNode = makeRelNode(context, operatorId, inputRelNode = childCtx.root, validation = false) + TransformContext(output, relNode) + } + + // ------------------------------------------------------------------------- + // Private helpers + // ------------------------------------------------------------------------- + + private def makeRelNode( + context: SubstraitContext, + operatorId: Long, + inputRelNode: RelNode, + validation: Boolean): RelNode = { + if (validation) { + RelBuilder.makeGlutenStrideRel( + inputRelNode, + stride, + RelBuilder.createExtensionNode(output.asJava), + context, + operatorId) + } else { + RelBuilder.makeGlutenStrideRel(inputRelNode, stride, context, operatorId) + } + } +} diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/GlutenStrideExecSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/GlutenStrideExecSuite.scala new file mode 100644 index 00000000000..ac13b723bdd --- /dev/null +++ b/backends-velox/src/test/scala/org/apache/gluten/execution/GlutenStrideExecSuite.scala @@ -0,0 +1,176 @@ +/* + * 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.execution + +import org.apache.spark.SparkConf +import org.apache.spark.sql.execution.SparkPlan + +/** + * Tests for [[GlutenStrideExecTransformer]] - the example Gluten-native custom operator. + * + * The operator keeps every N-th row within each input batch (indices 0, stride, 2*stride, ...). It + * has no Spark logical-plan equivalent, so it cannot be exercised via SQL. Instead the tests + * construct the physical plan node directly and drive execution via [[SparkPlan.executeCollect()]], + * which triggers the full Velox pipeline. + */ +class GlutenStrideExecSuite extends VeloxWholeStageTransformerSuite { + + override protected val resourcePath: String = "N/A" + override protected val fileFormat: String = "N/A" + + override protected def sparkConf: SparkConf = + super.sparkConf + .set("spark.memory.offHeap.size", "512m") + // Single partition keeps tests deterministic - no per-partition counter reset. + .set("spark.sql.shuffle.partitions", "1") + .set("spark.default.parallelism", "1") + + // --------------------------------------------------------------------------- + // Helper: wrap `child` in a GlutenStrideExecTransformer, execute it, and + // return the collected Long values from the first (id) column. + // --------------------------------------------------------------------------- + private def strideIds(child: SparkPlan, stride: Long): Seq[Long] = { + val strider = GlutenStrideExecTransformer(stride = stride, child = child) + strider.executeCollect().map(_.getLong(0)).toSeq + } + + // --------------------------------------------------------------------------- + // Correctness tests + // --------------------------------------------------------------------------- + + test("stride=1 returns all rows unchanged") { + withTable("stride_t") { + spark.range(0, 6).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").queryExecution.executedPlan + val result = strideIds(child, stride = 1L) + assert( + result.sorted == Seq(0L, 1L, 2L, 3L, 4L, 5L), + s"stride=1 should return all 6 rows, got: $result") + } + } + + test("stride=2 returns every other row") { + withTable("stride_t") { + // Write exactly 5 rows so the stride-2 result is deterministic within one batch + spark.range(0, 5).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").orderBy("id").queryExecution.executedPlan + val result = strideIds(child, stride = 2L) + // Sorted input: 0,1,2,3,4 -> indices 0,2,4 -> values 0,2,4 + assert( + result == Seq(0L, 2L, 4L), + s"stride=2 over 5 rows should give [0,2,4], got: $result") + } + } + + test("stride=3 returns every third row") { + withTable("stride_t") { + // 9 rows sorted -> indices 0,3,6 -> values 0,3,6 + spark.range(0, 9).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").orderBy("id").queryExecution.executedPlan + val result = strideIds(child, stride = 3L) + assert( + result == Seq(0L, 3L, 6L), + s"stride=3 over 9 rows should give [0,3,6], got: $result") + } + } + + test("stride > row count returns only the first row") { + withTable("stride_t") { + spark.range(0, 5).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").orderBy("id").queryExecution.executedPlan + val result = strideIds(child, stride = 100L) + assert( + result == Seq(0L), + s"stride > row count should keep only row at index 0, got: $result") + } + } + + test("stride applied to single-row input returns that single row") { + withTable("stride_t") { + spark.range(42L, 43L).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").queryExecution.executedPlan + val result = strideIds(child, stride = 5L) + assert( + result == Seq(42L), + s"single-row input should be returned unchanged, got: $result") + } + } + + test("stride applied to empty input returns empty result") { + withTable("stride_t") { + spark.range(0L, 0L).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").queryExecution.executedPlan + val result = strideIds(child, stride = 2L) + assert( + result.isEmpty, + s"empty input should produce empty output, got: $result") + } + } + + // --------------------------------------------------------------------------- + // Plan-structure test: GlutenStrideExecTransformer appears in the plan tree + // --------------------------------------------------------------------------- + + test("GlutenStrideExecTransformer appears in the plan tree string") { + withTable("stride_t") { + spark.range(0, 4).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").queryExecution.executedPlan + val strider = GlutenStrideExecTransformer(stride = 2L, child = child) + val planStr = strider.treeString + assert( + planStr.contains("GlutenStride"), + s"Expected 'GlutenStride' in plan tree but got:\n$planStr") + } + } + + // --------------------------------------------------------------------------- + // Chained plan: GlutenStrideExecTransformer wrapping an already-filtered scan + // --------------------------------------------------------------------------- + + test("stride after filter produces correct subset") { + withTable("stride_t") { + // rows 0..9 -> filter keeps even rows: 0,2,4,6,8 -> stride=2 -> indices 0,2,4 -> values 0,4,8 + spark.range(0, 10).write.format("parquet").saveAsTable("stride_t") + val filteredPlan = spark + .table("stride_t") + .filter("id % 2 = 0") + .orderBy("id") + .queryExecution + .executedPlan + val result = strideIds(filteredPlan, stride = 2L) + // Even values sorted: 0,2,4,6,8 -> keep indices 0,2,4 -> 0,4,8 + assert( + result == Seq(0L, 4L, 8L), + s"stride=2 after even-filter mismatch: $result") + } + } + + // --------------------------------------------------------------------------- + // Argument validation + // --------------------------------------------------------------------------- + + test("stride=0 is rejected at construction time") { + withTable("stride_t") { + spark.range(0, 3).write.format("parquet").saveAsTable("stride_t") + val child = spark.table("stride_t").queryExecution.executedPlan + val ex = intercept[IllegalArgumentException] { + GlutenStrideExecTransformer(stride = 0L, child = child) + } + assert(ex.getMessage.contains("stride"), s"Unexpected error message: ${ex.getMessage}") + } + } +} diff --git a/cpp/velox/CMakeLists.txt b/cpp/velox/CMakeLists.txt index e347381887e..0e0d9284aaa 100644 --- a/cpp/velox/CMakeLists.txt +++ b/cpp/velox/CMakeLists.txt @@ -178,6 +178,7 @@ set(VELOX_SRCS operators/functions/delta/DeltaBitmapAggregator.cc operators/functions/RowConstructorWithNull.cc operators/functions/SparkExprToSubfieldFilterParser.cc + operators/plannodes/GlutenStrideNode.cc operators/plannodes/RowVectorStream.cc operators/hashjoin/HashTableBuilder.cc operators/hashjoin/HashTableSerializer.cc diff --git a/cpp/velox/compute/VeloxBackend.cc b/cpp/velox/compute/VeloxBackend.cc index 3b735601a1f..72680750a10 100644 --- a/cpp/velox/compute/VeloxBackend.cc +++ b/cpp/velox/compute/VeloxBackend.cc @@ -47,6 +47,7 @@ #include "jni/JniFileSystem.h" #include "memory/GlutenBufferedInputBuilder.h" #include "operators/functions/SparkExprToSubfieldFilterParser.h" +#include "operators/plannodes/GlutenStrideNode.h" #include "operators/plannodes/RowVectorStream.h" #include "shuffle/ArrowShuffleDictionaryWriter.h" #include "udf/UdfLoader.h" @@ -228,6 +229,9 @@ void VeloxBackend::init( } #endif + // Register Gluten-native custom operator translators. + facebook::velox::exec::Operator::registerOperator(std::make_unique()); + const int32_t numTaskSlotsPerExecutor = [&]() { if (!backendConf_->valueExists(kNumTaskSlotsPerExecutor)) { LOG(WARNING) << kNumTaskSlotsPerExecutor << " is not set. Falling back to 1."; diff --git a/cpp/velox/operators/plannodes/GlutenStrideNode.cc b/cpp/velox/operators/plannodes/GlutenStrideNode.cc new file mode 100644 index 00000000000..85344c90f66 --- /dev/null +++ b/cpp/velox/operators/plannodes/GlutenStrideNode.cc @@ -0,0 +1,65 @@ +/* + * 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. + */ + +#include "GlutenStrideNode.h" + +#include "velox/exec/OperatorUtils.h" + +namespace gluten { + +GlutenStrideOperator::GlutenStrideOperator( + int32_t operatorId, + facebook::velox::exec::DriverCtx* driverCtx, + std::shared_ptr node) + : facebook::velox::exec::Operator(driverCtx, node->outputType(), operatorId, node->id(), "GlutenStride"), + stride_(node->stride()) {} + +facebook::velox::RowVectorPtr GlutenStrideOperator::getOutput() { + if (!input_) { + return nullptr; + } + auto inputBatch = std::move(input_); + const auto numRows = static_cast(inputBatch->size()); + + // Count how many rows we will keep: rows at indices 0, stride_, 2*stride_, ... + int64_t numSelected = 0; + for (int64_t i = 0; i < numRows; i += stride_) { + ++numSelected; + } + + if (numSelected == numRows) { + // stride == 1: pass every row through unchanged. + return inputBatch; + } + if (numSelected == 0) { + return inputBatch; // empty batch — nothing to do + } + + // Build an index buffer selecting rows 0, stride_, 2*stride_, ... + facebook::velox::BufferPtr indices = + facebook::velox::allocateIndices(static_cast(numSelected), pool()); + auto* rawIndices = indices->asMutable(); + facebook::velox::vector_size_t idx = 0; + for (int64_t i = 0; i < numRows; i += stride_) { + rawIndices[idx++] = static_cast(i); + } + + return facebook::velox::exec::wrap( + static_cast(numSelected), std::move(indices), inputBatch); +} + +} // namespace gluten diff --git a/cpp/velox/operators/plannodes/GlutenStrideNode.h b/cpp/velox/operators/plannodes/GlutenStrideNode.h new file mode 100644 index 00000000000..cfd120f823f --- /dev/null +++ b/cpp/velox/operators/plannodes/GlutenStrideNode.h @@ -0,0 +1,125 @@ +/* + * 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. + */ + +#pragma once + +#include "velox/core/PlanNode.h" +#include "velox/exec/Operator.h" + +namespace gluten { + +/// --------------------------------------------------------------------------- +/// GlutenStrideNode — a Gluten-native "every-Nth-row" sampler. +/// +/// Outputs exactly the rows at positions 0, stride, 2*stride, … within each +/// input batch (0-based within the batch, not globally across batches). +/// This is a deterministic, state-free operator that demonstrates the full +/// Gluten custom-operator mechanism without requiring any Velox upstream change. +/// +/// Encoding on the wire (reuses FetchRel to avoid a new proto type): +/// FetchRel.offset = stride +/// FetchRel.count = 0 (unused) +/// AdvancedExtension.optimization[0] = "isGlutenStride=1" +/// --------------------------------------------------------------------------- +class GlutenStrideNode final : public facebook::velox::core::PlanNode { + public: + GlutenStrideNode( + const facebook::velox::core::PlanNodeId& id, + int64_t stride, + facebook::velox::core::PlanNodePtr child) + : PlanNode(id), stride_(stride), sources_({std::move(child)}) { + VELOX_USER_CHECK_GT(stride_, 0, "GlutenStrideNode: stride must be > 0"); + } + + const facebook::velox::RowTypePtr& outputType() const override { + return sources_[0]->outputType(); + } + + const std::vector& sources() const override { + return sources_; + } + + int64_t stride() const { + return stride_; + } + + std::string_view name() const override { + return "GlutenStride"; + } + + folly::dynamic serialize() const override { + VELOX_UNSUPPORTED("GlutenStrideNode serialization is not supported"); + } + + private: + void addDetails(std::stringstream& stream) const override { + stream << "stride=" << stride_; + } + + const int64_t stride_; + std::vector sources_; +}; + +/// Outputs every stride-th row from each input batch. +/// Row indices within the batch are 0-based; the first row (index 0) is always +/// included. The counter resets at the start of every batch, keeping the +/// operator stateless and parallelism-friendly. +class GlutenStrideOperator : public facebook::velox::exec::Operator { + public: + GlutenStrideOperator( + int32_t operatorId, + facebook::velox::exec::DriverCtx* driverCtx, + std::shared_ptr node); + + bool needsInput() const override { + return !noMoreInput_ && !input_; + } + + void addInput(facebook::velox::RowVectorPtr input) override { + input_ = std::move(input); + } + + facebook::velox::RowVectorPtr getOutput() override; + + facebook::velox::exec::BlockingReason isBlocked(facebook::velox::ContinueFuture*) override { + return facebook::velox::exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return noMoreInput_ && !input_; + } + + private: + const int64_t stride_; +}; + +/// Registers GlutenStrideNode → GlutenStrideOperator. +/// Call once at backend startup via Operator::registerOperator(). +class GlutenStrideTranslator : public facebook::velox::exec::Operator::PlanNodeTranslator { + public: + std::unique_ptr toOperator( + facebook::velox::exec::DriverCtx* ctx, + int32_t id, + const facebook::velox::core::PlanNodePtr& node) override { + if (auto n = std::dynamic_pointer_cast(node)) { + return std::make_unique(id, ctx, n); + } + return nullptr; + } +}; + +} // namespace gluten diff --git a/cpp/velox/substrait/SubstraitToVeloxPlan.cc b/cpp/velox/substrait/SubstraitToVeloxPlan.cc index 9e85daa2064..d20cb4a4f7f 100644 --- a/cpp/velox/substrait/SubstraitToVeloxPlan.cc +++ b/cpp/velox/substrait/SubstraitToVeloxPlan.cc @@ -24,6 +24,7 @@ #include "compute/iceberg/IcebergPlanConverter.h" #include "jni/JniHashTable.h" #include "operators/hashjoin/HashTableBuilder.h" +#include "operators/plannodes/GlutenStrideNode.h" #include "operators/plannodes/RowVectorStream.h" #include "velox/connectors/hive/HiveDataSink.h" #include "velox/exec/TableWriter.h" @@ -1384,6 +1385,15 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait:: core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::FetchRel& fetchRel) { auto childNode = convertSingleInput<::substrait::FetchRel>(fetchRel); + + // GlutenStride is encoded as a FetchRel carrying the stride in FetchRel.offset, + // with a "isGlutenStride=1" marker in AdvancedExtension.optimization[0]. + if (fetchRel.has_advanced_extension() && + SubstraitParser::configSetInOptimization(fetchRel.advanced_extension(), "isGlutenStride=")) { + const int64_t stride = static_cast(fetchRel.offset()); + return std::make_shared(nextPlanNodeId(), stride, childNode); + } + return std::make_shared( nextPlanNodeId(), static_cast(fetchRel.offset()), diff --git a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GlutenStrideRelNode.java b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GlutenStrideRelNode.java new file mode 100644 index 00000000000..515c8609c3f --- /dev/null +++ b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/GlutenStrideRelNode.java @@ -0,0 +1,114 @@ +/* + * 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.substrait.rel; + +import org.apache.gluten.substrait.extensions.AdvancedExtensionNode; + +import com.google.protobuf.StringValue; +import io.substrait.proto.AdvancedExtension; +import io.substrait.proto.FetchRel; +import io.substrait.proto.Rel; +import io.substrait.proto.RelCommon; + +import java.io.Serializable; +import java.util.Collections; +import java.util.List; + +/** + * Substrait representation of Gluten's native GlutenStride operator. + * + *

The operator outputs every {@code stride}-th row from each input batch (rows at indices 0, + * stride, 2*stride, …). It is deterministic, stateless, and has no equivalent in standard Spark or + * Velox, making it a clean demonstration of the Gluten custom-operator mechanism. + * + *

Wire encoding

+ * + *

We reuse {@link FetchRel} to avoid introducing a new proto type: + * + *

    + *
  • {@code FetchRel.offset} = {@code stride} (the step size, e.g. 2 means keep every other row) + *
  • {@code FetchRel.count} = 0 (unused) + *
  • {@code AdvancedExtension.optimization[0]} = {@code "isGlutenStride=1"} — the marker that + * tells the C++ dispatcher to build a {@code GlutenStrideNode} instead of a {@code LimitNode} + *
+ * + *

C++ side

+ * + *

{@code SubstraitToVeloxPlanConverter::toVeloxPlan(FetchRel&)} checks for {@code + * isGlutenStride=1} and constructs {@code GlutenStrideNode(stride, child)}, which is executed by + * {@code GlutenStrideOperator}. + */ +public class GlutenStrideRelNode implements RelNode, Serializable { + + /** Optimization marker recognized by the C++ Substrait converter. */ + private static final String MARKER = "isGlutenStride=1"; + + private final RelNode input; + private final long stride; + // Present only in validation mode; carries input column types for the native validator. + private final AdvancedExtensionNode extensionNode; + + GlutenStrideRelNode(RelNode input, long stride) { + this.input = input; + this.stride = stride; + this.extensionNode = null; + } + + GlutenStrideRelNode(RelNode input, long stride, AdvancedExtensionNode extensionNode) { + this.input = input; + this.stride = stride; + this.extensionNode = extensionNode; + } + + @Override + public Rel toProtobuf() { + RelCommon.Builder relCommonBuilder = + RelCommon.newBuilder().setDirect(RelCommon.Direct.newBuilder()); + + // Build the marker optimization Any. + com.google.protobuf.Any markerAny = + com.google.protobuf.Any.pack(StringValue.newBuilder().setValue(MARKER).build()); + + AdvancedExtension.Builder extBuilder = AdvancedExtension.newBuilder(); + extBuilder.addOptimization(markerAny); + // Merge in the validation-mode enhancement if present. + if (extensionNode != null) { + AdvancedExtension baseExt = extensionNode.toProtobuf(); + if (baseExt.hasEnhancement()) { + extBuilder.setEnhancement(baseExt.getEnhancement()); + } + } + + FetchRel.Builder fetchBuilder = + FetchRel.newBuilder() + .setCommon(relCommonBuilder.build()) + .setOffset(stride) // stride is stored in offset + .setCount(0L) // unused + .setAdvancedExtension(extBuilder.build()); + + if (input != null) { + fetchBuilder.setInput(input.toProtobuf()); + } + + return Rel.newBuilder().setFetch(fetchBuilder.build()).build(); + } + + @Override + public List childNodes() { + return Collections.singletonList(input); + } +} diff --git a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelBuilder.java b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelBuilder.java index 40723946241..67ee7be504c 100644 --- a/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelBuilder.java +++ b/gluten-substrait/src/main/java/org/apache/gluten/substrait/rel/RelBuilder.java @@ -388,4 +388,30 @@ public static RelNode makeSetRel( context.registerRelToOperator(operatorId); return new SetRelNode(inputs, setOp, extensionNode); } + + // --------------------------------------------------------------------------- + // GlutenStride — custom Gluten-native operator: output every stride-th row. + // Encoded as FetchRel with AdvancedExtension marker "isGlutenStride=1". + // --------------------------------------------------------------------------- + + /** Execution mode: no type-annotation; the native plan is run as-is. */ + public static RelNode makeGlutenStrideRel( + RelNode input, long stride, SubstraitContext context, Long operatorId) { + context.registerRelToOperator(operatorId); + return new GlutenStrideRelNode(input, stride); + } + + /** + * Validation mode: attach an AdvancedExtension enhancement carrying the input column types so + * that the native validator can type-check the plan node. + */ + public static RelNode makeGlutenStrideRel( + RelNode input, + long stride, + AdvancedExtensionNode extensionNode, + SubstraitContext context, + Long operatorId) { + context.registerRelToOperator(operatorId); + return new GlutenStrideRelNode(input, stride, extensionNode); + } }