Copilot commented on code in PR #4884: URL: https://github.com/apache/datafusion-comet/pull/4884#discussion_r3600616188
########## native/shuffle/src/writers/rss/rss_partition_pusher.rs: ########## @@ -0,0 +1,85 @@ +// 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. + +use datafusion_comet_jni_bridge::{jni_call, JVMClasses}; +use jni::objects::{Global, JObject}; +use std::io::Write; +use std::sync::Arc; + +/// Forwards encoded shuffle data to a Java `ShufflePartitionPusher`. +/// +/// The global JNI reference keeps the Java pusher alive and is shared by partition-specific clones +/// created with [`Self::clone_with_pid`]. The [`Write`] implementation copies each buffer into a +/// Java byte array and invokes `ShufflePartitionPusher.pushPartitionData`. +#[derive(Debug)] +pub struct RssPartitionPusher { + pid: i32, + jobject: Arc<Global<JObject<'static>>>, +} + +impl RssPartitionPusher { + /// Creates a pusher without an assigned output partition. + /// + /// Use [`Self::clone_with_pid`] to create the partition-specific instances that write data. + pub fn try_new(jobject: Arc<Global<JObject<'static>>>) -> datafusion::common::Result<Self> { + Ok(RssPartitionPusher { pid: -1, jobject }) + } + + /// Creates a pusher for `pid` that shares the same Java pusher reference. + pub fn clone_with_pid(&self, pid: i32) -> Self { + RssPartitionPusher { + pid, + jobject: self.jobject.clone(), + } + } + + /// Pushes `buf` to the Java pusher for this instance's output partition. + /// + /// Returns the number of bytes reported as accepted by the Java callback. + pub fn push_partition_data(&mut self, buf: &[u8]) -> datafusion::common::Result<i32> { + let length = buf.len() as i32; + JVMClasses::with_env(|env| { + let jbytes = env.byte_array_from_slice(buf).unwrap(); + let length: i32 = unsafe { + jni_call!(env, + shuffle_partition_pusher(self.jobject.as_ref()).push_partition_data(self.pid, &jbytes, length) -> i32)? + }; + Ok(length) + }) Review Comment: `env.byte_array_from_slice(buf).unwrap()` can panic on JNI errors (for example OOM or a pending exception), which would abort the shuffle write. It is safer to convert that error into a `DataFusionError` and return it. ########## thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala: ########## @@ -0,0 +1,367 @@ +/* + * 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.spark.sql.comet.uniffle + +import java.nio.{ByteBuffer, ByteOrder} +import java.util +import java.util.{Collections, Optional} +import java.util.function.Supplier + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.{InterruptibleIterator, TaskContext} +import org.apache.spark.executor.ShuffleReadMetrics +import org.apache.spark.internal.Logging +import org.apache.spark.shuffle.{RssShuffleHandle, RssSparkConfig} +import org.apache.spark.shuffle.reader.RssShuffleReader +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleDependency +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.CompletionIterator +import org.apache.uniffle.client.api.{ShuffleManagerClient, ShuffleReadClient} +import org.apache.uniffle.client.factory.ShuffleClientFactory +import org.apache.uniffle.client.response.ShuffleBlock +import org.apache.uniffle.common.{ShuffleDataDistributionType, ShuffleServerInfo} +import org.apache.uniffle.common.compression.Codec +import org.apache.uniffle.common.config.RssConf +import org.apache.uniffle.shaded.org.roaringbitmap.longlong.Roaring64NavigableMap +import org.apache.uniffle.storage.handler.impl.ShuffleServerReadCostTracker + +import org.apache.comet.Native +import org.apache.comet.shuffle.{CometNativeShuffleReader, CometShuffleBlockIterator} +import org.apache.comet.vector.NativeUtil + +class CometUniffleShuffleReader[K, C]( + startPartition: Int, + endPartition: Int, + mapStartIndex: Int, + mapEndIndex: Int, + context: TaskContext, + rssShuffleHandle: RssShuffleHandle[K, _, C], + basePath: String, + hadoopConf: Configuration, + partitionNum: Int, + partitionToExpectBlocks: util.Map[Integer, Roaring64NavigableMap], + taskIdBitmap: Roaring64NavigableMap, + readMetrics: ShuffleReadMetrics, + managerClientSupplier: Supplier[ShuffleManagerClient], + rssConf: RssConf, + dataDistributionType: ShuffleDataDistributionType, + allPartitionToServers: util.Map[Integer, util.List[ShuffleServerInfo]]) + extends RssShuffleReader[K, C]( + startPartition, + endPartition, + mapStartIndex, + mapEndIndex, + context, + rssShuffleHandle, + basePath, + hadoopConf, + partitionNum, + partitionToExpectBlocks, + taskIdBitmap, + readMetrics, + managerClientSupplier, + rssConf, + dataDistributionType, + allPartitionToServers) + with CometNativeShuffleReader { + + private val shuffleServerReadCostTracker = new ShuffleServerReadCostTracker + private val dep = rssShuffleHandle.getDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + + override def read(): Iterator[Product2[K, C]] = { + val nativeLib = new Native() + val nativeUtil = new NativeUtil() + + val shuffleBlockIterator = new CometUniffleShuffleBlockIterator() + + context.addTaskCompletionListener[Unit] { _ => + shuffleBlockIterator.close() + nativeUtil.close() + } + + val recordIter: Iterator[(Int, ColumnarBatch)] = new Iterator[(Int, ColumnarBatch)] + with AutoCloseable { + private var currentBatch: ColumnarBatch = null + + // To avoid calling hasNext() multiple times for the same iterator, + // we cache the result of the last hasNext() call. + private var lastHasNext: Option[Boolean] = None + override def hasNext: Boolean = { + if (lastHasNext.isDefined) { + return lastHasNext.get + } + lastHasNext = Some(shuffleBlockIterator.hasNext != -1) + lastHasNext.get + } + + override def next(): (Int, ColumnarBatch) = { + lastHasNext = None + if (currentBatch != null) { + currentBatch.close() + } + + val dataBuf = shuffleBlockIterator.getBuffer + val bytesToRead = shuffleBlockIterator.getCurrentBlockLength + val fieldCount = shuffleBlockIterator.currentBatchFieldCount + + currentBatch = nativeUtil.getNextBatch( + fieldCount, + (arrayAddrs, schemaAddrs) => { + nativeLib.decodeShuffleBlock( + dataBuf, + bytesToRead, + arrayAddrs, + schemaAddrs, + tracingEnabled = false) + }) match { + case Some(batch) => batch + case None => + throw new IllegalStateException( + "Unexpected end of shuffle data while reading block of length " + bytesToRead) + } + + (0, currentBatch) + } + + override def close(): Unit = { + if (currentBatch != null) { + currentBatch.close() + currentBatch = null + } + shuffleBlockIterator.close() + } + } + + // Update the context task metrics for each record read. + val metricIter = CompletionIterator[(Any, Any), Iterator[(Any, Any)]]( + recordIter.map { record => + readMetrics.incRecordsRead(record._2.numRows()) + record + }, + context.taskMetrics().mergeShuffleReadMetrics()) + + // An interruptible iterator must be used here in order to support task cancellation + val interruptibleIter = new InterruptibleIterator[(Any, Any)](context, metricIter) + + val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) { + throw new UnsupportedOperationException("aggregate not allowed") + } else { + interruptibleIter.asInstanceOf[Iterator[Product2[K, C]]] + } + + // Sort the output if there is a sort ordering defined. + val resultIter = dep.keyOrdering match { + case Some(_: Ordering[K]) => + throw new UnsupportedOperationException("order not allowed") + case None => + aggregatedIter + } + + resultIter match { + case _: InterruptibleIterator[Product2[K, C]] => resultIter + case _ => + // Use another interruptible iterator here to support task cancellation as aggregator + // or(and) sorter may have consumed previous interruptible iterator. + new InterruptibleIterator[Product2[K, C]](context, resultIter) + } + } + + private def createShuffleReadClient(partition: Int): ShuffleReadClient = { + if (!partitionToExpectBlocks.containsKey(partition) + || partitionToExpectBlocks.get(partition).isEmpty) { + return null + } + + val shuffleServerInfoList = allPartitionToServers.get(partition) + if (shuffleServerInfoList != null && shuffleServerInfoList.size > 1 && rssConf.getBoolean( + RssSparkConfig.RSS_READ_REORDER_MULTI_SERVERS_ENABLED)) { + Collections.shuffle(shuffleServerInfoList) + } + + val isReplicaFilterEnabled = + rssConf.getInteger("rss.data.replica", 1) > 1 && shuffleServerInfoList.size > 1 + val expectedTaskIdsBitmapFilterEnable = + mapStartIndex != 0 || mapEndIndex != Integer.MAX_VALUE || isReplicaFilterEnabled + val retryMax = rssConf.getInteger("rss.client.retry.max", 50) + val retryIntervalMax = rssConf.getLong("rss.client.retry.interval.max", 10000L) + val compress = rssConf.getBoolean("spark.shuffle.compress".substring("spark.".length), true) + val codec = + if (compress) Codec.newInstance(rssConf) + else Optional.empty + val builder = ShuffleClientFactory.newReadBuilder + .readCostTracker(shuffleServerReadCostTracker) + .appId(rssShuffleHandle.getAppId) + .shuffleId(rssShuffleHandle.getDependency.shuffleId) + .partitionId(partition) + .basePath(basePath) + .partitionNumPerRange(1) + .partitionNum(partitionNum) + .blockIdBitmap(partitionToExpectBlocks.get(partition)) + .taskIdBitmap(taskIdBitmap) + .shuffleServerInfoList(shuffleServerInfoList) + .hadoopConf(hadoopConf) + .shuffleDataDistributionType(dataDistributionType) + .expectedTaskIdsBitmapFilterEnable(expectedTaskIdsBitmapFilterEnable) + .retryMax(retryMax) + .retryIntervalMax(retryIntervalMax) + .rssConf(rssConf) + .taskAttemptId(context.taskAttemptId()) + if (codec.isPresent && rssConf + .get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_ENABLED) + .asInstanceOf[Boolean]) { + builder + .overlappingDecompressionEnabled(true) + .codec(codec.get) + .overlappingDecompressionThreadNum( + rssConf.get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_THREADS)) + } + + ShuffleClientFactory.getInstance.createShuffleReadClient(builder) + } + + override def readAsShuffleBlockIterator(): CometShuffleBlockIterator = { + new CometUniffleShuffleBlockIterator() + } + + class CometUniffleShuffleBlockIterator extends CometShuffleBlockIterator with Logging { + private var currentPartition: Int = startPartition + private var current: ByteBuffer = _ + private var currentShuffleReadClient: ShuffleReadClient = createShuffleReadClient( + currentPartition) + + private val headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN) + private val INITIAL_BUFFER_SIZE = 128 * 1024 + private var dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) + private var currentBlockLength = 0 + + var currentBatchFieldCount: Int = _ + + override def hasNext: Int = { + if (current == null || !current.hasRemaining) { + val nextBlock = nextShuffleBlock() + if (nextBlock.isEmpty) { + return -1 + } + current = nextBlock.get + } + + val oldLimit = current.limit() + try { + current.limit(current.position() + headerBuf.capacity()) + headerBuf.clear() + headerBuf.put(current) + headerBuf.flip() + } finally { + current.limit(oldLimit) + } Review Comment: Reading the 16-byte shuffle block header assumes the current Uniffle `ByteBuffer` has at least 16 remaining bytes. If the header is split across Uniffle blocks, `current.limit(current.position() + 16)` can throw and the reader will fail even though the data is valid. Consider reading the header in a loop across blocks, similar to how the body is assembled. ########## thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleManager.scala: ########## @@ -0,0 +1,244 @@ +/* + * 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.spark.sql.comet.uniffle + +import java.util + +import scala.collection.JavaConverters._ + +import org.apache.spark.{SparkConf, TaskContext} +import org.apache.spark.executor.{ShuffleReadMetrics, ShuffleWriteMetrics} +import org.apache.spark.shuffle.{RssShuffleHandle, RssShuffleManager, RssSparkConfig, RssSparkShuffleUtils, ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} +import org.apache.spark.shuffle.handle.{ShuffleHandleInfo, SimpleShuffleHandleInfo} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleDependency} +import org.apache.uniffle.client.PartitionDataReplicaRequirementTracking +import org.apache.uniffle.common.ShuffleServerInfo +import org.apache.uniffle.common.exception.{RssException, RssFetchFailedException} +import org.apache.uniffle.common.util.RssUtils +import org.apache.uniffle.shaded.com.google.common.collect.Sets +import org.apache.uniffle.shaded.org.roaringbitmap.longlong.Roaring64NavigableMap + +import org.apache.comet.CometConf + +private[uniffle] object CometUniffleShuffleManager { + private[uniffle] class ReadMetrics(reporter: ShuffleReadMetricsReporter) + extends ShuffleReadMetrics { + override def incRemoteBytesRead(v: Long): Unit = reporter.incRemoteBytesRead(v) + + override def incFetchWaitTime(v: Long): Unit = reporter.incFetchWaitTime(v) + + override def incRecordsRead(v: Long): Unit = reporter.incRecordsRead(v) + } + + private[uniffle] def cometRssSparkConf(sparkConf: SparkConf): SparkConf = sparkConf.clone + .set(RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssSparkConfig.RSS_ROW_BASED.key, "false") + .set("spark.rss.client.io.compression.codec", "NONE") +} + +class CometUniffleShuffleManager(conf: SparkConf, isDriver: Boolean) + extends RssShuffleManager(conf, isDriver) { + + import CometUniffleShuffleManager._ + + assert( + conf.get(CometConf.COMET_SHUFFLE_MODE.key, "") == "native", + "CometUniffleShuffleManager only supports native shuffle mode") Review Comment: Using `assert` for a user-facing configuration check is risky because assertions can be disabled, which would allow this shuffle manager to run in unsupported modes. `require` keeps the validation enabled in production. ########## native/core/src/execution/jni_api.rs: ########## @@ -1351,3 +1352,46 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } + +/// Create a native RSS partition pusher backed by the supplied Java pusher. +/// +/// The returned handle owns the native pusher and must eventually be passed to +/// `releaseRssPartitionPusher`. +/// +/// # Safety +/// `pusher` must be a valid JNI reference for the duration of this call. +#[no_mangle] +pub unsafe extern "system" fn Java_org_apache_comet_Native_createRssPartitionPusher( + e: EnvUnowned, + _class: JClass, + pusher: JObject, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + let rss_partition_pusher = Box::new( + RssPartitionPusher::try_new(Arc::new(jni_new_global_ref!(env, pusher)?)).unwrap(), + ); + + Ok(Box::into_raw(rss_partition_pusher) as i64) + }) Review Comment: `RssPartitionPusher::try_new(...).unwrap()` can panic, which would unwind through JNI and crash the JVM. Since this function already uses `try_unwrap_or_throw`, it should return the error instead of panicking. ########## thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleWriter.scala: ########## @@ -0,0 +1,152 @@ +/* + * 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.spark.sql.comet.uniffle + +import java.util.function.Function + +import org.apache.spark.{SparkConf, TaskContext} +import org.apache.spark.executor.ShuffleWriteMetrics +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.shuffle.{RssShuffleHandle, RssShuffleManager} +import org.apache.spark.shuffle.writer.RssShuffleWriter +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.comet.execution.shuffle.{CometBaseNativeShuffleWriter, NativeShuffleSpec} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.uniffle.client.api.ShuffleWriteClient + +import org.apache.comet.Native +import org.apache.comet.shuffle.ShufflePartitionPusher + +case class CometUniffleShuffleWriter[K, V, C]( + spec: NativeShuffleSpec, + outputPartitioning: Partitioning, + outputAttributes: Seq[Attribute], + numParts: Int, + nativeMetrics: Map[String, SQLMetric], + appId: String, + shuffleId: Int, + taskId: String, + rssTaskAttemptId: Long, + rssShuffleWriteMetrics: ShuffleWriteMetrics, + shuffleManager: RssShuffleManager, + sparkConf: SparkConf, + shuffleWriteClient: ShuffleWriteClient, + rssHandle: RssShuffleHandle[K, V, C], + taskFailureCallback: Function[String, java.lang.Boolean], + context: TaskContext, + rangePartitionBounds: Option[Seq[InternalRow]] = None) + extends RssShuffleWriter[K, V, C]( + appId, + shuffleId, + taskId, + rssTaskAttemptId, + rssShuffleWriteMetrics, + shuffleManager, + sparkConf, + shuffleWriteClient, + rssHandle, + taskFailureCallback, + context) + with CometBaseNativeShuffleWriter + with ShufflePartitionPusher { + + private val nativeLib = new Native() + private var nativeShufflePusherHandle: Long = nativeLib.createRssPartitionPusher(this) + + override def writeImpl(inputs: Iterator[Product2[K, V]]): Unit = { + val detailedMetrics = Seq( + "elapsed_compute", + "encode_time", + "repart_time", + "input_batches", + "spill_count", + "spilled_bytes") + val metricsOutputRows = new SQLMetric("outputRows") + val metricsWriteTime = new SQLMetric("writeTime") + val shuffleWriterSQLMetrics = Map( + "output_rows" -> metricsOutputRows, + "data_size" -> nativeMetrics("dataSize"), + "write_time" -> metricsWriteTime) ++ + nativeMetrics.filterKeys(detailedMetrics.contains) + + nativeWrite(inputs, shuffleWriterSQLMetrics, nativeShufflePusherHandle) + + rssShuffleWriteMetrics.incRecordsWritten(metricsOutputRows.value) + rssShuffleWriteMetrics.incWriteTime(metricsWriteTime.value) + + val pushMergedDataTime = System.nanoTime() + // clear all + sendRestBlockAndWait() + sendCommit() + val writeDurationNanos = System.nanoTime() - pushMergedDataTime + rssShuffleWriteMetrics.incWriteTime(writeDurationNanos) + + } + + override def sendCommit(): Unit = { + if (!this.isMemoryShuffleEnabled) { + super.sendCommit() + } + } + + override def pushPartitionData(partitionId: Int, bytes: Array[Byte], length: Int): Int = { + val shuffleBlockInfos = + super.getBufferManager.addPartitionData( + partitionId, + bytes, + length, + System.currentTimeMillis) + super.processShuffleBlockInfos(shuffleBlockInfos) + // fast fail or resend data// fast fail or resend data + super.checkDataIfAnyFailure() Review Comment: The comment `// fast fail or resend data// fast fail or resend data` is duplicated, which looks accidental and makes the code harder to read. ########## thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala: ########## @@ -0,0 +1,367 @@ +/* + * 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.spark.sql.comet.uniffle + +import java.nio.{ByteBuffer, ByteOrder} +import java.util +import java.util.{Collections, Optional} +import java.util.function.Supplier + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.{InterruptibleIterator, TaskContext} +import org.apache.spark.executor.ShuffleReadMetrics +import org.apache.spark.internal.Logging +import org.apache.spark.shuffle.{RssShuffleHandle, RssSparkConfig} +import org.apache.spark.shuffle.reader.RssShuffleReader +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleDependency +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.CompletionIterator +import org.apache.uniffle.client.api.{ShuffleManagerClient, ShuffleReadClient} +import org.apache.uniffle.client.factory.ShuffleClientFactory +import org.apache.uniffle.client.response.ShuffleBlock +import org.apache.uniffle.common.{ShuffleDataDistributionType, ShuffleServerInfo} +import org.apache.uniffle.common.compression.Codec +import org.apache.uniffle.common.config.RssConf +import org.apache.uniffle.shaded.org.roaringbitmap.longlong.Roaring64NavigableMap +import org.apache.uniffle.storage.handler.impl.ShuffleServerReadCostTracker + +import org.apache.comet.Native +import org.apache.comet.shuffle.{CometNativeShuffleReader, CometShuffleBlockIterator} +import org.apache.comet.vector.NativeUtil + +class CometUniffleShuffleReader[K, C]( + startPartition: Int, + endPartition: Int, + mapStartIndex: Int, + mapEndIndex: Int, + context: TaskContext, + rssShuffleHandle: RssShuffleHandle[K, _, C], + basePath: String, + hadoopConf: Configuration, + partitionNum: Int, + partitionToExpectBlocks: util.Map[Integer, Roaring64NavigableMap], + taskIdBitmap: Roaring64NavigableMap, + readMetrics: ShuffleReadMetrics, + managerClientSupplier: Supplier[ShuffleManagerClient], + rssConf: RssConf, + dataDistributionType: ShuffleDataDistributionType, + allPartitionToServers: util.Map[Integer, util.List[ShuffleServerInfo]]) + extends RssShuffleReader[K, C]( + startPartition, + endPartition, + mapStartIndex, + mapEndIndex, + context, + rssShuffleHandle, + basePath, + hadoopConf, + partitionNum, + partitionToExpectBlocks, + taskIdBitmap, + readMetrics, + managerClientSupplier, + rssConf, + dataDistributionType, + allPartitionToServers) + with CometNativeShuffleReader { + + private val shuffleServerReadCostTracker = new ShuffleServerReadCostTracker + private val dep = rssShuffleHandle.getDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + + override def read(): Iterator[Product2[K, C]] = { + val nativeLib = new Native() + val nativeUtil = new NativeUtil() + + val shuffleBlockIterator = new CometUniffleShuffleBlockIterator() + + context.addTaskCompletionListener[Unit] { _ => + shuffleBlockIterator.close() + nativeUtil.close() + } + + val recordIter: Iterator[(Int, ColumnarBatch)] = new Iterator[(Int, ColumnarBatch)] + with AutoCloseable { + private var currentBatch: ColumnarBatch = null + + // To avoid calling hasNext() multiple times for the same iterator, + // we cache the result of the last hasNext() call. + private var lastHasNext: Option[Boolean] = None + override def hasNext: Boolean = { + if (lastHasNext.isDefined) { + return lastHasNext.get + } + lastHasNext = Some(shuffleBlockIterator.hasNext != -1) + lastHasNext.get + } + + override def next(): (Int, ColumnarBatch) = { + lastHasNext = None + if (currentBatch != null) { + currentBatch.close() + } + + val dataBuf = shuffleBlockIterator.getBuffer + val bytesToRead = shuffleBlockIterator.getCurrentBlockLength + val fieldCount = shuffleBlockIterator.currentBatchFieldCount + + currentBatch = nativeUtil.getNextBatch( + fieldCount, + (arrayAddrs, schemaAddrs) => { + nativeLib.decodeShuffleBlock( + dataBuf, + bytesToRead, + arrayAddrs, + schemaAddrs, + tracingEnabled = false) + }) match { + case Some(batch) => batch + case None => + throw new IllegalStateException( + "Unexpected end of shuffle data while reading block of length " + bytesToRead) + } + + (0, currentBatch) + } + + override def close(): Unit = { + if (currentBatch != null) { + currentBatch.close() + currentBatch = null + } + shuffleBlockIterator.close() + } + } + + // Update the context task metrics for each record read. + val metricIter = CompletionIterator[(Any, Any), Iterator[(Any, Any)]]( + recordIter.map { record => + readMetrics.incRecordsRead(record._2.numRows()) + record + }, + context.taskMetrics().mergeShuffleReadMetrics()) + + // An interruptible iterator must be used here in order to support task cancellation + val interruptibleIter = new InterruptibleIterator[(Any, Any)](context, metricIter) + + val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) { + throw new UnsupportedOperationException("aggregate not allowed") + } else { + interruptibleIter.asInstanceOf[Iterator[Product2[K, C]]] + } + + // Sort the output if there is a sort ordering defined. + val resultIter = dep.keyOrdering match { + case Some(_: Ordering[K]) => + throw new UnsupportedOperationException("order not allowed") + case None => + aggregatedIter + } + + resultIter match { + case _: InterruptibleIterator[Product2[K, C]] => resultIter + case _ => + // Use another interruptible iterator here to support task cancellation as aggregator + // or(and) sorter may have consumed previous interruptible iterator. + new InterruptibleIterator[Product2[K, C]](context, resultIter) + } + } + + private def createShuffleReadClient(partition: Int): ShuffleReadClient = { + if (!partitionToExpectBlocks.containsKey(partition) + || partitionToExpectBlocks.get(partition).isEmpty) { + return null + } + + val shuffleServerInfoList = allPartitionToServers.get(partition) + if (shuffleServerInfoList != null && shuffleServerInfoList.size > 1 && rssConf.getBoolean( + RssSparkConfig.RSS_READ_REORDER_MULTI_SERVERS_ENABLED)) { + Collections.shuffle(shuffleServerInfoList) + } + + val isReplicaFilterEnabled = + rssConf.getInteger("rss.data.replica", 1) > 1 && shuffleServerInfoList.size > 1 + val expectedTaskIdsBitmapFilterEnable = + mapStartIndex != 0 || mapEndIndex != Integer.MAX_VALUE || isReplicaFilterEnabled + val retryMax = rssConf.getInteger("rss.client.retry.max", 50) + val retryIntervalMax = rssConf.getLong("rss.client.retry.interval.max", 10000L) + val compress = rssConf.getBoolean("spark.shuffle.compress".substring("spark.".length), true) + val codec = + if (compress) Codec.newInstance(rssConf) + else Optional.empty + val builder = ShuffleClientFactory.newReadBuilder + .readCostTracker(shuffleServerReadCostTracker) + .appId(rssShuffleHandle.getAppId) + .shuffleId(rssShuffleHandle.getDependency.shuffleId) + .partitionId(partition) + .basePath(basePath) + .partitionNumPerRange(1) + .partitionNum(partitionNum) + .blockIdBitmap(partitionToExpectBlocks.get(partition)) + .taskIdBitmap(taskIdBitmap) + .shuffleServerInfoList(shuffleServerInfoList) + .hadoopConf(hadoopConf) + .shuffleDataDistributionType(dataDistributionType) + .expectedTaskIdsBitmapFilterEnable(expectedTaskIdsBitmapFilterEnable) + .retryMax(retryMax) + .retryIntervalMax(retryIntervalMax) + .rssConf(rssConf) + .taskAttemptId(context.taskAttemptId()) + if (codec.isPresent && rssConf + .get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_ENABLED) + .asInstanceOf[Boolean]) { + builder + .overlappingDecompressionEnabled(true) + .codec(codec.get) + .overlappingDecompressionThreadNum( + rssConf.get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_THREADS)) + } + + ShuffleClientFactory.getInstance.createShuffleReadClient(builder) + } + + override def readAsShuffleBlockIterator(): CometShuffleBlockIterator = { + new CometUniffleShuffleBlockIterator() + } + + class CometUniffleShuffleBlockIterator extends CometShuffleBlockIterator with Logging { + private var currentPartition: Int = startPartition + private var current: ByteBuffer = _ + private var currentShuffleReadClient: ShuffleReadClient = createShuffleReadClient( + currentPartition) + + private val headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN) + private val INITIAL_BUFFER_SIZE = 128 * 1024 + private var dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) + private var currentBlockLength = 0 + + var currentBatchFieldCount: Int = _ + + override def hasNext: Int = { + if (current == null || !current.hasRemaining) { + val nextBlock = nextShuffleBlock() + if (nextBlock.isEmpty) { + return -1 + } + current = nextBlock.get + } + + val oldLimit = current.limit() + try { + current.limit(current.position() + headerBuf.capacity()) + headerBuf.clear() + headerBuf.put(current) + headerBuf.flip() + } finally { + current.limit(oldLimit) + } + + val compressedLength = headerBuf.getLong + currentBatchFieldCount = headerBuf.getLong.toInt + + // Subtract 8 because compressedLength includes the 8-byte field count we already read + val bytesToRead = compressedLength - 8 + if (bytesToRead > Integer.MAX_VALUE) { + throw new IllegalStateException( + "Native shuffle block size of " + bytesToRead + " exceeds maximum of " + + Integer.MAX_VALUE + ". Try reducing spark.comet.columnar.shuffle.batch.size.") + } + currentBlockLength = bytesToRead.toInt + Review Comment: `bytesToRead` can become negative if the stream is corrupt (for example if `compressedLength < 8`). That leads to confusing `ByteBuffer.limit` / allocation failures. It is safer to validate the header and throw a clear exception before using the length. ########## thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala: ########## @@ -0,0 +1,367 @@ +/* + * 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.spark.sql.comet.uniffle + +import java.nio.{ByteBuffer, ByteOrder} +import java.util +import java.util.{Collections, Optional} +import java.util.function.Supplier + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.{InterruptibleIterator, TaskContext} +import org.apache.spark.executor.ShuffleReadMetrics +import org.apache.spark.internal.Logging +import org.apache.spark.shuffle.{RssShuffleHandle, RssSparkConfig} +import org.apache.spark.shuffle.reader.RssShuffleReader +import org.apache.spark.sql.comet.execution.shuffle.CometShuffleDependency +import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.util.CompletionIterator +import org.apache.uniffle.client.api.{ShuffleManagerClient, ShuffleReadClient} +import org.apache.uniffle.client.factory.ShuffleClientFactory +import org.apache.uniffle.client.response.ShuffleBlock +import org.apache.uniffle.common.{ShuffleDataDistributionType, ShuffleServerInfo} +import org.apache.uniffle.common.compression.Codec +import org.apache.uniffle.common.config.RssConf +import org.apache.uniffle.shaded.org.roaringbitmap.longlong.Roaring64NavigableMap +import org.apache.uniffle.storage.handler.impl.ShuffleServerReadCostTracker + +import org.apache.comet.Native +import org.apache.comet.shuffle.{CometNativeShuffleReader, CometShuffleBlockIterator} +import org.apache.comet.vector.NativeUtil + +class CometUniffleShuffleReader[K, C]( + startPartition: Int, + endPartition: Int, + mapStartIndex: Int, + mapEndIndex: Int, + context: TaskContext, + rssShuffleHandle: RssShuffleHandle[K, _, C], + basePath: String, + hadoopConf: Configuration, + partitionNum: Int, + partitionToExpectBlocks: util.Map[Integer, Roaring64NavigableMap], + taskIdBitmap: Roaring64NavigableMap, + readMetrics: ShuffleReadMetrics, + managerClientSupplier: Supplier[ShuffleManagerClient], + rssConf: RssConf, + dataDistributionType: ShuffleDataDistributionType, + allPartitionToServers: util.Map[Integer, util.List[ShuffleServerInfo]]) + extends RssShuffleReader[K, C]( + startPartition, + endPartition, + mapStartIndex, + mapEndIndex, + context, + rssShuffleHandle, + basePath, + hadoopConf, + partitionNum, + partitionToExpectBlocks, + taskIdBitmap, + readMetrics, + managerClientSupplier, + rssConf, + dataDistributionType, + allPartitionToServers) + with CometNativeShuffleReader { + + private val shuffleServerReadCostTracker = new ShuffleServerReadCostTracker + private val dep = rssShuffleHandle.getDependency.asInstanceOf[CometShuffleDependency[_, _, _]] + + override def read(): Iterator[Product2[K, C]] = { + val nativeLib = new Native() + val nativeUtil = new NativeUtil() + + val shuffleBlockIterator = new CometUniffleShuffleBlockIterator() + + context.addTaskCompletionListener[Unit] { _ => + shuffleBlockIterator.close() + nativeUtil.close() + } + + val recordIter: Iterator[(Int, ColumnarBatch)] = new Iterator[(Int, ColumnarBatch)] + with AutoCloseable { + private var currentBatch: ColumnarBatch = null + + // To avoid calling hasNext() multiple times for the same iterator, + // we cache the result of the last hasNext() call. + private var lastHasNext: Option[Boolean] = None + override def hasNext: Boolean = { + if (lastHasNext.isDefined) { + return lastHasNext.get + } + lastHasNext = Some(shuffleBlockIterator.hasNext != -1) + lastHasNext.get + } + + override def next(): (Int, ColumnarBatch) = { + lastHasNext = None + if (currentBatch != null) { + currentBatch.close() + } + + val dataBuf = shuffleBlockIterator.getBuffer + val bytesToRead = shuffleBlockIterator.getCurrentBlockLength + val fieldCount = shuffleBlockIterator.currentBatchFieldCount + + currentBatch = nativeUtil.getNextBatch( + fieldCount, + (arrayAddrs, schemaAddrs) => { + nativeLib.decodeShuffleBlock( + dataBuf, + bytesToRead, + arrayAddrs, + schemaAddrs, + tracingEnabled = false) + }) match { + case Some(batch) => batch + case None => + throw new IllegalStateException( + "Unexpected end of shuffle data while reading block of length " + bytesToRead) + } + + (0, currentBatch) + } + + override def close(): Unit = { + if (currentBatch != null) { + currentBatch.close() + currentBatch = null + } + shuffleBlockIterator.close() + } + } + + // Update the context task metrics for each record read. + val metricIter = CompletionIterator[(Any, Any), Iterator[(Any, Any)]]( + recordIter.map { record => + readMetrics.incRecordsRead(record._2.numRows()) + record + }, + context.taskMetrics().mergeShuffleReadMetrics()) + + // An interruptible iterator must be used here in order to support task cancellation + val interruptibleIter = new InterruptibleIterator[(Any, Any)](context, metricIter) + + val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) { + throw new UnsupportedOperationException("aggregate not allowed") + } else { + interruptibleIter.asInstanceOf[Iterator[Product2[K, C]]] + } + + // Sort the output if there is a sort ordering defined. + val resultIter = dep.keyOrdering match { + case Some(_: Ordering[K]) => + throw new UnsupportedOperationException("order not allowed") + case None => + aggregatedIter + } + + resultIter match { + case _: InterruptibleIterator[Product2[K, C]] => resultIter + case _ => + // Use another interruptible iterator here to support task cancellation as aggregator + // or(and) sorter may have consumed previous interruptible iterator. + new InterruptibleIterator[Product2[K, C]](context, resultIter) + } + } + + private def createShuffleReadClient(partition: Int): ShuffleReadClient = { + if (!partitionToExpectBlocks.containsKey(partition) + || partitionToExpectBlocks.get(partition).isEmpty) { + return null + } + + val shuffleServerInfoList = allPartitionToServers.get(partition) + if (shuffleServerInfoList != null && shuffleServerInfoList.size > 1 && rssConf.getBoolean( + RssSparkConfig.RSS_READ_REORDER_MULTI_SERVERS_ENABLED)) { + Collections.shuffle(shuffleServerInfoList) + } + + val isReplicaFilterEnabled = + rssConf.getInteger("rss.data.replica", 1) > 1 && shuffleServerInfoList.size > 1 + val expectedTaskIdsBitmapFilterEnable = + mapStartIndex != 0 || mapEndIndex != Integer.MAX_VALUE || isReplicaFilterEnabled + val retryMax = rssConf.getInteger("rss.client.retry.max", 50) + val retryIntervalMax = rssConf.getLong("rss.client.retry.interval.max", 10000L) + val compress = rssConf.getBoolean("spark.shuffle.compress".substring("spark.".length), true) + val codec = + if (compress) Codec.newInstance(rssConf) + else Optional.empty + val builder = ShuffleClientFactory.newReadBuilder + .readCostTracker(shuffleServerReadCostTracker) + .appId(rssShuffleHandle.getAppId) + .shuffleId(rssShuffleHandle.getDependency.shuffleId) + .partitionId(partition) + .basePath(basePath) + .partitionNumPerRange(1) + .partitionNum(partitionNum) + .blockIdBitmap(partitionToExpectBlocks.get(partition)) + .taskIdBitmap(taskIdBitmap) + .shuffleServerInfoList(shuffleServerInfoList) + .hadoopConf(hadoopConf) + .shuffleDataDistributionType(dataDistributionType) + .expectedTaskIdsBitmapFilterEnable(expectedTaskIdsBitmapFilterEnable) + .retryMax(retryMax) + .retryIntervalMax(retryIntervalMax) + .rssConf(rssConf) + .taskAttemptId(context.taskAttemptId()) + if (codec.isPresent && rssConf + .get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_ENABLED) + .asInstanceOf[Boolean]) { + builder + .overlappingDecompressionEnabled(true) + .codec(codec.get) + .overlappingDecompressionThreadNum( + rssConf.get(RssSparkConfig.RSS_READ_OVERLAPPING_DECOMPRESSION_THREADS)) + } + + ShuffleClientFactory.getInstance.createShuffleReadClient(builder) + } + + override def readAsShuffleBlockIterator(): CometShuffleBlockIterator = { + new CometUniffleShuffleBlockIterator() + } + + class CometUniffleShuffleBlockIterator extends CometShuffleBlockIterator with Logging { + private var currentPartition: Int = startPartition + private var current: ByteBuffer = _ + private var currentShuffleReadClient: ShuffleReadClient = createShuffleReadClient( + currentPartition) + + private val headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN) + private val INITIAL_BUFFER_SIZE = 128 * 1024 + private var dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) + private var currentBlockLength = 0 + + var currentBatchFieldCount: Int = _ + + override def hasNext: Int = { + if (current == null || !current.hasRemaining) { + val nextBlock = nextShuffleBlock() + if (nextBlock.isEmpty) { + return -1 + } + current = nextBlock.get + } + + val oldLimit = current.limit() + try { + current.limit(current.position() + headerBuf.capacity()) + headerBuf.clear() + headerBuf.put(current) + headerBuf.flip() + } finally { + current.limit(oldLimit) + } + + val compressedLength = headerBuf.getLong + currentBatchFieldCount = headerBuf.getLong.toInt + + // Subtract 8 because compressedLength includes the 8-byte field count we already read + val bytesToRead = compressedLength - 8 + if (bytesToRead > Integer.MAX_VALUE) { + throw new IllegalStateException( + "Native shuffle block size of " + bytesToRead + " exceeds maximum of " + + Integer.MAX_VALUE + ". Try reducing spark.comet.columnar.shuffle.batch.size.") + } + currentBlockLength = bytesToRead.toInt + + if (dataBuf.capacity < currentBlockLength) { + val newCapacity = Math.min(bytesToRead * 2L, Integer.MAX_VALUE).toInt + dataBuf = ByteBuffer.allocateDirect(newCapacity) + } + dataBuf.clear + dataBuf.limit(currentBlockLength) + + while (dataBuf.hasRemaining) { + if (!current.hasRemaining) { + val nextBlock = nextShuffleBlock() + if (nextBlock.isEmpty) { + throw new IllegalStateException( + "Unexpected end of shuffle data while reading block of length " + currentBlockLength) + } + current = nextBlock.get + } + val readLimit = dataBuf.remaining() + var oldLimit = -1 + try { + if (readLimit < current.remaining()) { + oldLimit = current.limit() + current.limit(current.position() + readLimit) + } + dataBuf.put(current) + } finally { + if (oldLimit != -1) { + current.limit(oldLimit) + } + } + } + + currentBlockLength + } + + private def nextShuffleBlock(): Option[ByteBuffer] = { Review Comment: `nextShuffleBlock()` uses tail recursion (`return nextShuffleBlock()`) to skip empty partitions. Adding `@tailrec` ensures the compiler will optimize it into a loop and prevents accidental non-tail changes from introducing deep recursion. ########## native/jni-bridge/src/lib.rs: ########## @@ -238,6 +240,8 @@ pub struct JVMClasses<'a> { pub comet_udf_bridge: Option<CometUdfBridge<'a>>, /// JNI handles for the CometS3CredentialDispatcher SPI and the CometS3Credentials POJO. pub comet_s3_credential_dispatcher: CometS3CredentialDispatcher<'a>, + /// TODO: add comment + pub shuffle_partition_pusher: ShufflePartitionPusher<'a>, Review Comment: `/// TODO: add comment` should be replaced with a real doc comment so it is clear what `shuffle_partition_pusher` is used for in JNI class caching. ########## spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala: ########## @@ -555,6 +555,14 @@ object CometShuffleExchangeExec return reasons.toSeq } + // TODO: code factor Review Comment: There is a leftover TODO comment (`code factor`) in production code. It would be clearer either to remove it or replace it with an explanation of why the Uniffle shuffle manager is treated specially here. ########## .github/workflows/pr_build_linux.yml: ########## @@ -582,3 +582,131 @@ jobs: SPARK_TPCDS_JOIN_CONF: | spark.sql.autoBroadcastJoinThreshold=-1 spark.sql.join.forceApplyShuffledHashJoin=true + + # Uniffle integration test - verifies Comet shuffle against a local Uniffle cluster + uniffle-integration-test: + needs: build-native + name: Uniffle Integration Test + runs-on: ubuntu-24.04 + container: + image: amd64/rust + env: + JAVA_TOOL_OPTIONS: --add-exports=java.base/sun.nio.ch=ALL-UNNAMED --add-exports=java.base/sun.util.calendar=ALL-UNNAMED --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.lang.invoke=ALL-UNNAMED --add-opens=java.base/java.lang.reflect=ALL-UNNAMED --add-opens=java.base/java.io=ALL-UNNAMED --add-opens=java.base/java.net=ALL-UNNAMED --add-opens=java.base/java.nio=ALL-UNNAMED --add-opens=java.base/java.util=ALL-UNNAMED --add-opens=java.base/java.util.concurrent=ALL-UNNAMED --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED --add-opens=java.base/jdk.internal.ref=ALL-UNNAMED --add-opens=java.base/sun.nio.ch=ALL-UNNAMED --add-opens=java.base/sun.nio.cs=ALL-UNNAMED --add-opens=java.base/sun.security.action=ALL-UNNAMED --add-opens=java.base/sun.util.calendar=ALL-UNNAMED + UNIFFLE_VERSION: 0.10.0 + HADOOP_VERSION: 2.10.2 + steps: + - uses: actions/checkout@v7 + + - name: Setup Rust & Java toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: ${{ env.RUST_VERSION }} + jdk-version: 17 + + - name: Download native library + uses: actions/download-artifact@v8 + with: + name: native-lib-linux + path: native/target/release/ + + - name: Cache Maven dependencies + uses: actions/cache@v6 + with: + path: | + ~/.m2/repository + /root/.m2/repository + key: ${{ runner.os }}-java-maven-${{ hashFiles('**/pom.xml') }} + restore-keys: | + ${{ runner.os }}-java-maven- + + - name: Cache TPC-H data + id: cache-uniffle-tpch + uses: actions/cache@v6 + with: + path: ./tpch + key: tpch-${{ hashFiles('.github/workflows/pr_build_linux.yml') }} + + - name: Cache TPC-DS data + id: cache-uniffle-tpcds + uses: actions/cache@v6 + with: + path: ./tpcds-sf-1 + key: tpcds-${{ hashFiles('.github/workflows/pr_build_linux.yml') }} + + - name: Install Uniffle + run: | + wget "https://www.apache.org/dyn/closer.lua/uniffle/${UNIFFLE_VERSION}/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz?action=download"; \ + -O "/opt/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz" + wget "https://www.apache.org/dyn/closer.lua/hadoop/common/hadoop-${HADOOP_VERSION}/hadoop-${HADOOP_VERSION}.tar.gz?action=download"; \ + -O "/opt/hadoop-${HADOOP_VERSION}.tar.gz" Review Comment: Using `https://www.apache.org/dyn/closer.lua/...` in CI can be flaky because it depends on mirror selection and can sometimes return an HTML redirect/maintenance page instead of the expected tarball. For deterministic CI, it is safer to download from `archive.apache.org` (or another fixed host) for both Uniffle and Hadoop artifacts. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
