diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 26281316f0..8fdb4ec8d2 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -581,3 +581,79 @@ 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.nio=ALL-UNNAMED --add-opens=java.base/java.lang=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: 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" + mkdir /opt/uniffle + tar xzf "/opt/apache-uniffle-${UNIFFLE_VERSION}-bin.tar.gz" \ + -C /opt/uniffle --strip-components=1 + tar xzf "/opt/hadoop-${HADOOP_VERSION}.tar.gz" -C /opt/ + mkdir /opt/uniffle/shuffle_data + printf 'XMX_SIZE=16g\nHADOOP_HOME=/opt/hadoop-%s\n' "${HADOOP_VERSION}" \ + > /opt/uniffle/conf/rss-env.sh + printf 'rss.coordinator.shuffle.nodes.max 1\nrss.rpc.server.port 19999\n' \ + > /opt/uniffle/conf/coordinator.conf + printf '%s\n' \ + 'rss.server.app.expired.withoutHeartbeat 7200000' \ + 'rss.server.heartbeat.delay 3000' \ + 'rss.rpc.server.port 19997' \ + 'rss.rpc.server.type GRPC_NETTY' \ + 'rss.jetty.http.port 19996' \ + 'rss.server.netty.port 19995' \ + 'rss.storage.basePath /opt/uniffle/shuffle_data' \ + 'rss.storage.type MEMORY_LOCALFILE' \ + 'rss.coordinator.quorum localhost:19999' \ + 'rss.server.flush.thread.alive 10' \ + 'rss.server.single.buffer.flush.threshold 64m' \ + > /opt/uniffle/conf/server.conf + cd /opt/uniffle + bash ./bin/start-coordinator.sh + bash ./bin/start-shuffle-server.sh + + - name: Build project + run: | + ./mvnw -B -Prelease install -DskipTests -Pspark-3.5,uniffle + + - name: Run Uniffle integration tests + run: | + export COMET_SHUFFLE_MANAGER=uniffle && ./mvnw test -Pspark-3.5,uniffle,uniffle-comet-test -pl spark diff --git a/native/Cargo.lock b/native/Cargo.lock index c54812c818..8f4ee80a87 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -198,7 +198,7 @@ dependencies = [ "snap", "strum", "strum_macros", - "thiserror 2.0.18", + "thiserror", "uuid", "zstd", ] @@ -1342,12 +1342,6 @@ dependencies = [ "shlex 2.0.1", ] -[[package]] -name = "cesu8" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" - [[package]] name = "cexpr" version = "0.6.0" @@ -2014,7 +2008,7 @@ dependencies = [ "iceberg", "iceberg-storage-opendal", "itertools 0.15.0", - "jni 0.22.4", + "jni", "lazy_static", "log", "log4rs", @@ -2050,7 +2044,7 @@ dependencies = [ "datafusion", "serde", "serde_json", - "thiserror 2.0.18", + "thiserror", ] [[package]] @@ -2076,14 +2070,14 @@ dependencies = [ "assertables", "datafusion", "datafusion-comet-common", - "jni 0.22.4", + "jni", "lazy_static", "once_cell", "parquet", "paste", "prost", "regex", - "thiserror 2.0.18", + "thiserror", ] [[package]] @@ -2125,10 +2119,11 @@ dependencies = [ "datafusion-comet-spark-expr", "futures", "itertools 0.15.0", - "jni 0.21.1", + "jni", "log", "lz4_flex", "parquet", + "paste", "simd-adler32", "snap", "tempfile", @@ -2150,7 +2145,7 @@ dependencies = [ "datafusion-comet-jni-bridge", "futures", "hex", - "jni 0.22.4", + "jni", "num", "rand 0.10.2", "regex", @@ -2779,7 +2774,7 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" dependencies = [ - "thiserror 2.0.18", + "thiserror", ] [[package]] @@ -3958,22 +3953,6 @@ dependencies = [ "jiff-tzdb", ] -[[package]] -name = "jni" -version = "0.21.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" -dependencies = [ - "cesu8", - "cfg-if", - "combine", - "jni-sys 0.3.1", - "log", - "thiserror 1.0.69", - "walkdir", - "windows-sys 0.45.0", -] - [[package]] name = "jni" version = "0.22.4" @@ -3984,11 +3963,11 @@ dependencies = [ "combine", "java-locator", "jni-macros", - "jni-sys 0.4.1", + "jni-sys", "libloading", "log", "simd_cesu8", - "thiserror 2.0.18", + "thiserror", "walkdir", "windows-link", ] @@ -4006,15 +3985,6 @@ dependencies = [ "syn 2.0.118", ] -[[package]] -name = "jni-sys" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" -dependencies = [ - "jni-sys 0.4.1", -] - [[package]] name = "jni-sys" version = "0.4.1" @@ -4270,7 +4240,7 @@ dependencies = [ "serde-value", "serde_json", "serde_yaml", - "thiserror 2.0.18", + "thiserror", "thread-id", "typemap-ors", "unicode-segmentation", @@ -4577,7 +4547,7 @@ dependencies = [ "serde", "serde_json", "serde_urlencoded", - "thiserror 2.0.18", + "thiserror", "tokio", "tracing", "url", @@ -5259,7 +5229,7 @@ dependencies = [ "spin 0.10.0", "symbolic-demangle", "tempfile", - "thiserror 2.0.18", + "thiserror", ] [[package]] @@ -5414,7 +5384,7 @@ dependencies = [ "rustc-hash 2.1.3", "rustls", "socket2", - "thiserror 2.0.18", + "thiserror", "tokio", "tracing", "web-time", @@ -5437,7 +5407,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.18", + "thiserror", "tinyvec", "tracing", "web-time", @@ -6002,7 +5972,7 @@ checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" dependencies = [ "core-foundation", "core-foundation-sys", - "jni 0.22.4", + "jni", "log", "once_cell", "rustls", @@ -6596,33 +6566,13 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl 2.0.18", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.118", + "thiserror-impl", ] [[package]] @@ -7342,22 +7292,13 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.45.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" -dependencies = [ - "windows-targets 0.42.2", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -7366,7 +7307,7 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -7378,67 +7319,34 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" -dependencies = [ - "windows_aarch64_gnullvm 0.42.2", - "windows_aarch64_msvc 0.42.2", - "windows_i686_gnu 0.42.2", - "windows_i686_msvc 0.42.2", - "windows_x86_64_gnu 0.42.2", - "windows_x86_64_gnullvm 0.42.2", - "windows_x86_64_msvc 0.42.2", -] - [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -7451,48 +7359,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.42.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d2c1a5df02..c56b9960cd 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1218,6 +1218,7 @@ pub extern "system" fn Java_org_apache_comet_Native_getRustThreadId( use crate::execution::columnar_to_row::ColumnarToRowContext; use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; +use datafusion_comet_shuffle::RssPartitionPusher; use datafusion_spark::function::math::bin::SparkBin; use datafusion_spark::function::string::soundex::SparkSoundex; @@ -1351,3 +1352,38 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( Ok(()) }) } + +#[no_mangle] +/// # Safety +/// TODO: add comment +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) + }) +} + +#[no_mangle] +/// # Safety +/// TODO: add comment +pub unsafe extern "system" fn Java_org_apache_comet_Native_releaseRssPartitionPusher( + e: EnvUnowned, + _class: JClass, + handle: jlong, +) { + try_unwrap_or_throw(&e, |_env| { + debug_assert!(handle != 0, "releaseRssPartitionPusher: handle is null"); + if handle != 0 { + let _pusher: Box = Box::from_raw(handle as *mut RssPartitionPusher); + // _pusher is dropped here, freeing the resources + } + Ok(()) + }) +} diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 7ed2b3331c..21be43c0f3 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -114,6 +114,9 @@ use datafusion::physical_plan::limit::GlobalLimitExec; use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; use datafusion_comet_proto::spark_expression::ListLiteral; use datafusion_comet_proto::spark_operator::SparkFilePartition; +use datafusion_comet_proto::spark_partitioning::partition_writer::PartitionWriterStruct::{ + LocalPartitionWriter, RssPartitionWriter, +}; use datafusion_comet_proto::{ spark_expression::{ self, agg_expr::ExprStruct as AggExprStruct, expr::ExprStruct, literal::Value, AggExpr, @@ -126,6 +129,7 @@ use datafusion_comet_proto::{ }, spark_partitioning::{partitioning::PartitioningStruct, Partitioning as SparkPartitioning}, }; +use datafusion_comet_shuffle::ShufflePartitionWriter; use datafusion_comet_spark_expr::{ jvm_udf::JvmScalarUdfExpr, ApproxPercentile, ArrayInsert, Avg, AvgDecimal, Cast, CheckOverflow, Correlation, Covariance, CreateNamedStruct, DecimalRescaleCheckOverflow, GetArrayStructFields, @@ -1630,12 +1634,27 @@ impl PhysicalPlanner { }?; let write_buffer_size = writer.write_buffer_size as usize; + let partition_writer_struct = writer + .partition_writer + .as_ref() + .and_then(|pw| pw.partition_writer_struct.clone()) + .ok_or_else(|| { + GeneralError("ShuffleWriter is missing partition_writer".to_string()) + })?; + let shuffle_partition_writer = match partition_writer_struct { + LocalPartitionWriter(local) => ShufflePartitionWriter::Local { + output_data_file: local.output_data_file, + output_index_file: local.output_index_file, + }, + RssPartitionWriter(rss) => ShufflePartitionWriter::Rss { + rss_partition_pusher_handle: rss.rss_partition_pusher, + }, + }; let shuffle_writer = Arc::new(ShuffleWriterExec::try_new( writer_input, partitioning, codec, - writer.output_data_file.clone(), - writer.output_index_file.clone(), + shuffle_partition_writer, writer.tracing_enabled, write_buffer_size, )?); diff --git a/native/jni-bridge/src/lib.rs b/native/jni-bridge/src/lib.rs index 8db4c07851..3e68368406 100644 --- a/native/jni-bridge/src/lib.rs +++ b/native/jni-bridge/src/lib.rs @@ -195,7 +195,9 @@ mod comet_s3_credential_dispatcher; mod comet_task_memory_manager; mod comet_udf_bridge; mod shuffle_block_iterator; +mod shuffle_partition_pusher; +use crate::shuffle_partition_pusher::ShufflePartitionPusher; use arrow_array_stream::ArrowArrayStream; pub use comet_metric_node::*; pub use comet_s3_credential_dispatcher::CometS3CredentialDispatcher; @@ -238,6 +240,8 @@ pub struct JVMClasses<'a> { pub comet_udf_bridge: Option>, /// 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>, } unsafe impl Send for JVMClasses<'_> {} @@ -316,6 +320,7 @@ impl JVMClasses<'_> { bridge }, comet_s3_credential_dispatcher: CometS3CredentialDispatcher::new(env).unwrap(), + shuffle_partition_pusher: ShufflePartitionPusher::new(env).unwrap(), } }); } diff --git a/native/jni-bridge/src/shuffle_partition_pusher.rs b/native/jni-bridge/src/shuffle_partition_pusher.rs new file mode 100644 index 0000000000..ae8bb787de --- /dev/null +++ b/native/jni-bridge/src/shuffle_partition_pusher.rs @@ -0,0 +1,48 @@ +// 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 jni::{ + errors::Result as JniResult, + objects::{JClass, JMethodID}, + signature::{Primitive, ReturnType}, + strings::JNIString, + Env, +}; + +pub struct ShufflePartitionPusher<'a> { + pub class: JClass<'a>, + pub method_push_partition_data: JMethodID, + pub method_push_partition_data_ret: ReturnType, +} + +impl<'a> ShufflePartitionPusher<'a> { + pub const JVM_CLASS: &'static str = "org/apache/comet/shuffle/ShufflePartitionPusher"; + + pub fn new(env: &mut Env<'a>) -> JniResult> { + let class = env.find_class(JNIString::new(Self::JVM_CLASS))?; + + Ok(ShufflePartitionPusher { + method_push_partition_data: env.get_method_id( + JNIString::new(Self::JVM_CLASS), + jni::jni_str!("pushPartitionData"), + jni::jni_sig!("(I[BI)I"), + )?, + method_push_partition_data_ret: ReturnType::Primitive(Primitive::Int), + class, + }) + } +} diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 2fcfe7f25b..6f02c62aed 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -324,8 +324,7 @@ enum CompressionCodec { message ShuffleWriter { spark.spark_partitioning.Partitioning partitioning = 1; - string output_data_file = 3; - string output_index_file = 4; + spark.spark_partitioning.PartitionWriter partition_writer = 2; CompressionCodec codec = 5; int32 compression_level = 6; bool tracing_enabled = 7; diff --git a/native/proto/src/proto/partitioning.proto b/native/proto/src/proto/partitioning.proto index e70b8264f0..9a9d9e6604 100644 --- a/native/proto/src/proto/partitioning.proto +++ b/native/proto/src/proto/partitioning.proto @@ -58,3 +58,19 @@ message RoundRobinPartition { // Maximum number of columns to hash. 0 means no limit (hash all columns). int32 max_hash_columns = 2; } + +message PartitionWriter { + oneof partition_writer_struct { + LocalPartitionWriter local_partition_writer = 1; + RssPartitionWriter rss_partition_writer = 2; + } +} + +message LocalPartitionWriter { + string output_data_file = 1; + string output_index_file = 2; +} + +message RssPartitionWriter { + int64 rss_partition_pusher = 1; +} \ No newline at end of file diff --git a/native/shuffle/Cargo.toml b/native/shuffle/Cargo.toml index 4d62181089..674f0bb54a 100644 --- a/native/shuffle/Cargo.toml +++ b/native/shuffle/Cargo.toml @@ -41,7 +41,8 @@ datafusion-comet-jni-bridge = { workspace = true } datafusion-comet-spark-expr = { workspace = true } futures = { workspace = true } itertools = "0.15.0" -jni = "0.21" +jni = "0.22.4" +paste = "1.0.14" log = "0.4" lz4_flex = { version = "0.13.1", default-features = false, features = ["frame"] } # parquet is only used by the shuffle_bench binary (shuffle-bench feature) @@ -61,6 +62,9 @@ tokio = { version = "1", features = ["macros", "rt-multi-thread"] } [features] shuffle-bench = ["clap", "parquet"] +[package.metadata.cargo-machete] +ignored = ["paste"] + [lib] name = "datafusion_comet_shuffle" path = "src/lib.rs" diff --git a/native/shuffle/benches/shuffle_writer.rs b/native/shuffle/benches/shuffle_writer.rs index 27abd919fa..29d2b9435b 100644 --- a/native/shuffle/benches/shuffle_writer.rs +++ b/native/shuffle/benches/shuffle_writer.rs @@ -30,7 +30,8 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_comet_shuffle::{ - CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShuffleWriterExec, + CometPartitioning, CompressionCodec, ShuffleBlockWriter, ShufflePartitionWriter, + ShuffleWriterExec, }; use itertools::Itertools; use std::io::Cursor; @@ -149,8 +150,10 @@ fn create_shuffle_writer_exec( ))), partitioning, compression_codec, - "/tmp/data.out".to_string(), - "/tmp/index.out".to_string(), + ShufflePartitionWriter::Local { + output_data_file: "/tmp/data.out".to_string(), + output_index_file: "/tmp/index.out".to_string(), + }, false, 1024 * 1024, ) diff --git a/native/shuffle/src/bin/shuffle_bench.rs b/native/shuffle/src/bin/shuffle_bench.rs index cc1c960848..50d617d8b9 100644 --- a/native/shuffle/src/bin/shuffle_bench.rs +++ b/native/shuffle/src/bin/shuffle_bench.rs @@ -45,7 +45,9 @@ use datafusion::physical_plan::common::collect; use datafusion::physical_plan::metrics::{MetricValue, MetricsSet}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{ParquetReadOptions, SessionContext}; -use datafusion_comet_shuffle::{CometPartitioning, CompressionCodec, ShuffleWriterExec}; +use datafusion_comet_shuffle::{ + CometPartitioning, CompressionCodec, ShufflePartitionWriter, ShuffleWriterExec, +}; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use std::fs; use std::path::{Path, PathBuf}; @@ -476,8 +478,10 @@ async fn execute_shuffle_write( input, partitioning, codec, - data_file, - index_file, + ShufflePartitionWriter::Local { + output_data_file: data_file, + output_index_file: index_file, + }, false, write_buffer_size, ) diff --git a/native/shuffle/src/lib.rs b/native/shuffle/src/lib.rs index 2263ae0dac..83ccd5aee0 100644 --- a/native/shuffle/src/lib.rs +++ b/native/shuffle/src/lib.rs @@ -28,5 +28,7 @@ pub(crate) mod writers; pub use comet_partitioning::CometPartitioning; pub use ipc::read_ipc_compressed; pub use schema_align::SchemaAlignExec; +pub use shuffle_writer::ShufflePartitionWriter; pub use shuffle_writer::ShuffleWriterExec; +pub use writers::RssPartitionPusher; pub use writers::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index de6989e52a..5a98f91530 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -22,8 +22,8 @@ use crate::partitioners::{ EmptySchemaShufflePartitioner, MultiPartitionShuffleRepartitioner, ShufflePartitioner, SinglePartitionShufflePartitioner, }; -use crate::writers::LocalPartitionWriter; -use crate::{CometPartitioning, CompressionCodec, ShuffleBlockWriter}; +use crate::writers::{LocalPartitionWriter, PartitionWriter, RssPartitionWriter}; +use crate::{CometPartitioning, CompressionCodec, RssPartitionPusher, ShuffleBlockWriter}; use async_trait::async_trait; use datafusion::common::exec_datafusion_err; use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; @@ -46,6 +46,17 @@ use std::{ sync::Arc, }; +#[derive(Debug, Clone)] +pub enum ShufflePartitionWriter { + Local { + output_data_file: String, + output_index_file: String, + }, + Rss { + rss_partition_pusher_handle: i64, + }, +} + /// The shuffle writer operator maps each input partition to M output partitions based on a /// partitioning scheme. No guarantees are made about the order of the resulting partitions. #[derive(Debug)] @@ -54,10 +65,8 @@ pub struct ShuffleWriterExec { input: Arc, /// Partitioning scheme to use partitioning: CometPartitioning, - /// Output data file path - output_data_file: String, - /// Output index file path - output_index_file: String, + /// Partition writer + partition_writer: ShufflePartitionWriter, /// Metrics metrics: ExecutionPlanMetricsSet, /// Cache for expensive-to-compute plan properties @@ -76,8 +85,7 @@ impl ShuffleWriterExec { input: Arc, partitioning: CometPartitioning, codec: CompressionCodec, - output_data_file: String, - output_index_file: String, + partition_writer: ShufflePartitionWriter, tracing_enabled: bool, write_buffer_size: usize, ) -> Result { @@ -92,8 +100,7 @@ impl ShuffleWriterExec { input, partitioning, metrics: ExecutionPlanMetricsSet::new(), - output_data_file, - output_index_file, + partition_writer, cache, codec, tracing_enabled, @@ -149,8 +156,7 @@ impl ExecutionPlan for ShuffleWriterExec { Arc::clone(&children[0]), self.partitioning.clone(), self.codec.clone(), - self.output_data_file.clone(), - self.output_index_file.clone(), + self.partition_writer.clone(), self.tracing_enabled, self.write_buffer_size, )?)), @@ -174,8 +180,7 @@ impl ExecutionPlan for ShuffleWriterExec { futures::stream::once(external_shuffle( input, partition, - self.output_data_file.clone(), - self.output_index_file.clone(), + self.partition_writer.clone(), self.partitioning.clone(), metrics, context, @@ -192,8 +197,7 @@ impl ExecutionPlan for ShuffleWriterExec { async fn external_shuffle( mut input: SendableRecordBatchStream, partition: usize, - output_data_file: String, - output_index_file: String, + partition_writer: ShufflePartitionWriter, partitioning: CometPartitioning, metrics: ShufflePartitionerMetrics, context: Arc, @@ -204,40 +208,67 @@ async fn external_shuffle( let schema = input.schema(); let shuffle_block_writer = ShuffleBlockWriter::try_new(schema.as_ref(), codec.clone())?; - let local_partition_writer = LocalPartitionWriter::try_new( - output_data_file, - output_index_file, - shuffle_block_writer, - partitioning.partition_count(), - context.session_config().batch_size(), - write_buffer_size, - context.runtime_env(), - )?; - - let mut repartitioner: Box = match &partitioning { - _ if schema.fields().is_empty() => { - log::debug!("found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner"); - Box::new(EmptySchemaShufflePartitioner::try_new( + + let mut repartitioner: Box = match partition_writer { + ShufflePartitionWriter::Local { + output_data_file, + output_index_file, + } => { + let local_partition_writer = LocalPartitionWriter::try_new( + output_data_file, + output_index_file, + shuffle_block_writer, + partitioning.partition_count(), + context.session_config().batch_size(), + write_buffer_size, + context.runtime_env(), + )?; + build_repartitioner( + partition, local_partition_writer, - Arc::clone(&schema), + partitioning, + &schema, + metrics, + &context, + tracing_enabled, + )? + } + ShufflePartitionWriter::Rss { + rss_partition_pusher_handle, + } => { + // `rss_partition_pusher_handle` is a raw pointer to a `Box` + // created by `createRssPartitionPusher` and released separately by + // `releaseRssPartitionPusher`, so we borrow it here without taking ownership. + debug_assert!( + rss_partition_pusher_handle != 0, + "external_shuffle: rss_partition_pusher_handle is null" + ); + let pusher = unsafe { + (rss_partition_pusher_handle as *mut RssPartitionPusher) + .as_mut() + .ok_or_else(|| { + exec_datafusion_err!( + "external_shuffle: rss_partition_pusher_handle is null" + ) + })? + }; + let rss_partition_writer = RssPartitionWriter::try_new( + pusher, + shuffle_block_writer, partitioning.partition_count(), + context.session_config().batch_size(), + write_buffer_size, + )?; + build_repartitioner( + partition, + rss_partition_writer, + partitioning, + &schema, metrics, - )?) + &context, + tracing_enabled, + )? } - any if any.partition_count() == 1 => Box::new(SinglePartitionShufflePartitioner::try_new( - local_partition_writer, - metrics, - context.session_config().batch_size(), - )?), - _ => Box::new(MultiPartitionShuffleRepartitioner::try_new( - partition, - local_partition_writer, - partitioning, - metrics, - context.runtime_env(), - context.session_config().batch_size(), - tracing_enabled, - )?), }; while let Some(batch) = input.next().await { @@ -259,6 +290,47 @@ async fn external_shuffle( Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) as SendableRecordBatchStream) } +/// Selects the appropriate [`ShufflePartitioner`] for the given `partitioning` and wires it to +/// `partition_writer`, which controls where the partitioned output is stored (local files or a +/// remote shuffle service). The partitioner choice is independent of the storage backend. +#[allow(clippy::too_many_arguments)] +fn build_repartitioner( + partition: usize, + partition_writer: W, + partitioning: CometPartitioning, + schema: &SchemaRef, + metrics: ShufflePartitionerMetrics, + context: &Arc, + tracing_enabled: bool, +) -> Result> { + let repartitioner: Box = match &partitioning { + _ if schema.fields().is_empty() => { + log::debug!("found empty schema, overriding {partitioning:?} partitioning with EmptySchemaShufflePartitioner"); + Box::new(EmptySchemaShufflePartitioner::try_new( + partition_writer, + Arc::clone(schema), + partitioning.partition_count(), + metrics, + )?) + } + any if any.partition_count() == 1 => Box::new(SinglePartitionShufflePartitioner::try_new( + partition_writer, + metrics, + context.session_config().batch_size(), + )?), + _ => Box::new(MultiPartitionShuffleRepartitioner::try_new( + partition, + partition_writer, + partitioning, + metrics, + context.runtime_env(), + context.session_config().batch_size(), + tracing_enabled, + )?), + }; + Ok(repartitioner) +} + #[cfg(test)] mod test { use super::*; @@ -531,8 +603,10 @@ mod test { ))), partitioning, CompressionCodec::Zstd(1), - "/tmp/data.out".to_string(), - "/tmp/index.out".to_string(), + ShufflePartitionWriter::Local { + output_data_file: "/tmp/data.out".to_string(), + output_index_file: "/tmp/index.out".to_string(), + }, false, 1024 * 1024, // write_buffer_size: 1MB default ) @@ -590,8 +664,10 @@ mod test { ))), CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), - data_file.clone(), - index_file.clone(), + ShufflePartitionWriter::Local { + output_data_file: data_file.clone(), + output_index_file: index_file.clone(), + }, false, 1024 * 1024, ) @@ -794,8 +870,10 @@ mod test { ))), CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), - data_file.to_str().unwrap().to_string(), - index_file.to_str().unwrap().to_string(), + ShufflePartitionWriter::Local { + output_data_file: data_file.to_str().unwrap().to_string(), + output_index_file: index_file.to_str().unwrap().to_string(), + }, false, 1024 * 1024, ) @@ -882,8 +960,10 @@ mod test { ))), CometPartitioning::RoundRobin(num_partitions, 0), CompressionCodec::Zstd(1), - data_file.to_str().unwrap().to_string(), - index_file.to_str().unwrap().to_string(), + ShufflePartitionWriter::Local { + output_data_file: data_file.to_str().unwrap().to_string(), + output_index_file: index_file.to_str().unwrap().to_string(), + }, false, 1024 * 1024, ) diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index 6d330fd12a..0d17056b9c 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -19,10 +19,13 @@ mod buf_batch_writer; mod checksum; mod local; mod partition_writer; +mod rss; mod shuffle_block_writer; pub(crate) use buf_batch_writer::BufBatchWriter; pub(crate) use checksum::Checksum; pub(crate) use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter; +pub use rss::rss_partition_pusher::RssPartitionPusher; +pub(crate) use rss::rss_partition_writer::RssPartitionWriter; pub use shuffle_block_writer::{CompressionCodec, ShuffleBlockWriter}; diff --git a/native/shuffle/src/writers/rss/mod.rs b/native/shuffle/src/writers/rss/mod.rs new file mode 100644 index 0000000000..267d8ba5c9 --- /dev/null +++ b/native/shuffle/src/writers/rss/mod.rs @@ -0,0 +1,19 @@ +// 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. + +pub mod rss_partition_pusher; +pub mod rss_partition_writer; diff --git a/native/shuffle/src/writers/rss/rss_partition_pusher.rs b/native/shuffle/src/writers/rss/rss_partition_pusher.rs new file mode 100644 index 0000000000..96c362e4fc --- /dev/null +++ b/native/shuffle/src/writers/rss/rss_partition_pusher.rs @@ -0,0 +1,73 @@ +// 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; + +#[derive(Debug)] +pub struct RssPartitionPusher { + pid: i32, + jobject: Arc>>, +} + +impl RssPartitionPusher { + pub fn try_new(jobject: Arc>>) -> datafusion::common::Result { + Ok(RssPartitionPusher { pid: -1, jobject }) + } + + pub fn clone_with_pid(&self, pid: i32) -> Self { + RssPartitionPusher { + pid, + jobject: self.jobject.clone(), + } + } + + pub fn push_partition_data(&mut self, buf: &[u8]) -> datafusion::common::Result { + 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) + }) + } +} + +impl Clone for RssPartitionPusher { + fn clone(&self) -> Self { + RssPartitionPusher { + pid: self.pid, + jobject: self.jobject.clone(), + } + } +} + +impl Write for RssPartitionPusher { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.push_partition_data(buf) + .map(|n| n as usize) + .map_err(std::io::Error::other) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} diff --git a/native/shuffle/src/writers/rss/rss_partition_writer.rs b/native/shuffle/src/writers/rss/rss_partition_writer.rs new file mode 100644 index 0000000000..dad1eba995 --- /dev/null +++ b/native/shuffle/src/writers/rss/rss_partition_writer.rs @@ -0,0 +1,100 @@ +// 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 crate::metrics::ShufflePartitionerMetrics; +use crate::writers::{BufBatchWriter, PartitionWriter, RssPartitionPusher}; +use crate::ShuffleBlockWriter; +use arrow::array::RecordBatch; + +pub(crate) struct RssPartitionWriter { + num_output_partitions: usize, + partition_writers: Vec>, +} + +impl RssPartitionWriter { + pub fn try_new( + pusher: &mut RssPartitionPusher, + shuffle_block_writer: ShuffleBlockWriter, + num_output_partitions: usize, + batch_size: usize, + write_buffer_size: usize, + ) -> datafusion::common::Result { + let partition_writers = (0..num_output_partitions) + .map(|pid| { + BufBatchWriter::new( + shuffle_block_writer.clone(), + pusher.clone_with_pid(pid as i32), + write_buffer_size, + batch_size, + ) + }) + .collect::>(); + Ok(Self { + num_output_partitions, + partition_writers, + }) + } +} + +impl PartitionWriter for RssPartitionWriter { + fn write( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> + where + I: Iterator>, + { + for batch in iter.by_ref() { + let batch = batch?; + self.partition_writers[pid].write(&batch, &metrics.encode_time, &metrics.write_time)?; + } + + // Flush the partition writer if there are multiple output partitions. + // This is to prevent partition writer buffers from holding too much memory. + if self.num_output_partitions > 1 { + self.partition_writers[pid].flush(&metrics.encode_time, &metrics.write_time)?; + } + + Ok(()) + } + + fn finish_partition( + &mut self, + pid: usize, + iter: &mut I, + metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> + where + I: Iterator>, + { + for batch in iter.by_ref() { + let batch = batch?; + self.partition_writers[pid].write(&batch, &metrics.encode_time, &metrics.write_time)?; + } + + self.partition_writers[pid].flush(&metrics.encode_time, &metrics.write_time) + } + + fn finish_all( + &mut self, + _metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> { + Ok(()) + } +} diff --git a/pom.xml b/pom.xml index 6739c01f29..dd45aae372 100644 --- a/pom.xml +++ b/pom.xml @@ -83,6 +83,7 @@ under the License. 4.13.6 2.0.17 33.2.1-jre + 0.10.0 1.21.4 2.31.51 ${project.basedir}/../native/target/debug @@ -870,6 +871,22 @@ under the License. + + + uniffle + + thirdparty/comet-uniffle + + + + + org.apache.uniffle + rss-client-spark3-shaded + ${uniffle.version} + + + + @@ -1252,6 +1269,22 @@ under the License. com.google.thirdparty.publicsuffix.PublicSuffixType + + org.apache.uniffle + rss-client-spark3-shaded + + + javax.servlet.* + META-INF/versions/9/javax/xml/bind/ModuleUtil + javax.activation.MailcapCommandMap + javax.activation.MimetypesFileTypeMap + javax.annotation.MatchesPattern$Checker + javax.annotation.Nonnegative$Checker + javax.annotation.RegEx$Checker + javax.annotation.Resource$AuthenticationType + META-INF/versions/11/org/roaringbitmap/ArraysShim + + true true diff --git a/spark-integration/pom.xml b/spark-integration/pom.xml index f9eb9605cf..5fdd73d492 100644 --- a/spark-integration/pom.xml +++ b/spark-integration/pom.xml @@ -44,21 +44,6 @@ under the License. org.apache.datafusion comet-spark-spark${spark.version.short}_${scala.binary.version} ${project.version} - - - - org.apache.datafusion - comet-common-spark${spark.version.short}_${scala.binary.version} - - - - org.apache.arrow - * - - diff --git a/spark/pom.xml b/spark/pom.xml index 26ca9f6dca..5db944ad5a 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -43,6 +43,7 @@ under the License. org.apache.datafusion comet-common-spark${spark.version.short}_${scala.binary.version} ${project.version} + true org.apache.spark @@ -59,6 +60,7 @@ under the License. org.apache.arrow arrow-vector + true org.scala-lang.modules @@ -76,6 +78,7 @@ under the License. com.google.protobuf protobuf-java + true org.scalatest @@ -115,6 +118,7 @@ under the License. com.google.guava guava + true * @@ -144,10 +148,12 @@ under the License. org.apache.arrow arrow-memory-unsafe + true org.apache.arrow arrow-c-data + true org.apache.hadoop @@ -358,6 +364,61 @@ under the License. + + uniffle-comet-test + + + org.apache.uniffle + rss-client-spark3-shaded + ${uniffle.version} + + + + + + org.apache.maven.plugins + maven-surefire-plugin + + + + ${project.basedir}/../thirdparty/comet-uniffle/target/comet-uniffle_${scala.binary.version}-${project.version}.jar + + + + + + org.scalatest + scalatest-maven-plugin + ${scalatest-maven-plugin.version} + + + ${project.basedir}/../thirdparty/comet-uniffle/target/comet-uniffle_${scala.binary.version}-${project.version}.jar + + ${project.build.directory}/surefire-reports + . + SparkTestSuite.txt + D + + org.apache.comet.IntegrationTestSuite + + + file:src/test/resources/log4j2.properties + true + ${project.build.directory}/tmp + + + + + test + + test + + + + + + + diff --git a/spark/src/main/java/org/apache/comet/CometLocalShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometLocalShuffleBlockIterator.java new file mode 100644 index 0000000000..9bd8b8ccde --- /dev/null +++ b/spark/src/main/java/org/apache/comet/CometLocalShuffleBlockIterator.java @@ -0,0 +1,142 @@ +/* + * 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.comet; + +import java.io.Closeable; +import java.io.EOFException; +import java.io.IOException; +import java.io.InputStream; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.Channels; +import java.nio.channels.ReadableByteChannel; + +/** + * Provides raw compressed shuffle blocks to native code via JNI. + * + *

Reads block headers (compressed length + field count) from a shuffle InputStream and loads the + * compressed body into a DirectByteBuffer. Native code pulls blocks by calling hasNext() and + * getBuffer(). + * + *

The DirectByteBuffer returned by getBuffer() is only valid until the next hasNext() call. + * Native code must fully consume it (via read_ipc_compressed which allocates new memory for the + * decompressed data) before pulling the next block. + */ +public class CometLocalShuffleBlockIterator implements CometShuffleBlockIterator, Closeable { + + private static final int INITIAL_BUFFER_SIZE = 128 * 1024; + + private final ReadableByteChannel channel; + private final InputStream inputStream; + private final ByteBuffer headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN); + private ByteBuffer dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE); + private boolean closed = false; + private int currentBlockLength = 0; + + public CometLocalShuffleBlockIterator(InputStream in) { + this.inputStream = in; + this.channel = Channels.newChannel(in); + } + + /** + * Reads the next block header and loads the compressed body into the internal buffer. Called by + * native code via JNI. + * + *

Header format: 8-byte compressedLength (includes field count but not itself) + 8-byte + * fieldCount (discarded, schema comes from protobuf). + * + * @return the compressed body length in bytes (codec prefix + compressed IPC), or -1 if EOF + */ + public int hasNext() throws IOException { + if (closed) { + return -1; + } + + // Read 16-byte header: clear() resets position=0, limit=capacity, + // preparing the buffer for channel.read() to fill it + headerBuf.clear(); + while (headerBuf.hasRemaining()) { + int bytesRead = channel.read(headerBuf); + if (bytesRead < 0) { + if (headerBuf.position() == 0) { + close(); + return -1; + } + throw new EOFException("Data corrupt: unexpected EOF while reading batch header"); + } + } + headerBuf.flip(); + long compressedLength = headerBuf.getLong(); + // Field count discarded - schema determined by ShuffleScan protobuf fields + headerBuf.getLong(); + + // Subtract 8 because compressedLength includes the 8-byte field count we already read + long 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 = (int) bytesToRead; + + if (dataBuf.capacity() < currentBlockLength) { + int newCapacity = (int) Math.min(bytesToRead * 2L, Integer.MAX_VALUE); + dataBuf = ByteBuffer.allocateDirect(newCapacity); + } + + dataBuf.clear(); + dataBuf.limit(currentBlockLength); + while (dataBuf.hasRemaining()) { + int bytesRead = channel.read(dataBuf); + if (bytesRead < 0) { + throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch"); + } + } + // Note: native side uses get_direct_buffer_address (base pointer) + currentBlockLength, + // not the buffer's position/limit. No flip needed. + + return currentBlockLength; + } + + /** + * Returns the DirectByteBuffer containing the current block's compressed bytes (4-byte codec + * prefix + compressed IPC data). Called by native code via JNI. + */ + public ByteBuffer getBuffer() { + return dataBuf; + } + + /** Returns the length of the current block in bytes. Called by native code via JNI. */ + public int getCurrentBlockLength() { + return currentBlockLength; + } + + @Override + public void close() throws IOException { + if (!closed) { + closed = true; + inputStream.close(); + } + } +} diff --git a/spark/src/main/java/org/apache/comet/CometNativeShuffleReader.java b/spark/src/main/java/org/apache/comet/CometNativeShuffleReader.java new file mode 100644 index 0000000000..4759b0e682 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/CometNativeShuffleReader.java @@ -0,0 +1,25 @@ +/* + * 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.comet; + +public interface CometNativeShuffleReader { + + CometShuffleBlockIterator readAsShuffleBlockIterator(); +} diff --git a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java index 9f72b20f51..49c02883d4 100644 --- a/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java +++ b/spark/src/main/java/org/apache/comet/CometShuffleBlockIterator.java @@ -19,124 +19,16 @@ package org.apache.comet; -import java.io.Closeable; -import java.io.EOFException; import java.io.IOException; -import java.io.InputStream; import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.channels.Channels; -import java.nio.channels.ReadableByteChannel; -/** - * Provides raw compressed shuffle blocks to native code via JNI. - * - *

Reads block headers (compressed length + field count) from a shuffle InputStream and loads the - * compressed body into a DirectByteBuffer. Native code pulls blocks by calling hasNext() and - * getBuffer(). - * - *

The DirectByteBuffer returned by getBuffer() is only valid until the next hasNext() call. - * Native code must fully consume it (via read_ipc_compressed which allocates new memory for the - * decompressed data) before pulling the next block. - */ -public class CometShuffleBlockIterator implements Closeable { - - private static final int INITIAL_BUFFER_SIZE = 128 * 1024; - - private final ReadableByteChannel channel; - private final InputStream inputStream; - private final ByteBuffer headerBuf = ByteBuffer.allocate(16).order(ByteOrder.LITTLE_ENDIAN); - private ByteBuffer dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE); - private boolean closed = false; - private int currentBlockLength = 0; - - public CometShuffleBlockIterator(InputStream in) { - this.inputStream = in; - this.channel = Channels.newChannel(in); - } - - /** - * Reads the next block header and loads the compressed body into the internal buffer. Called by - * native code via JNI. - * - *

Header format: 8-byte compressedLength (includes field count but not itself) + 8-byte - * fieldCount (discarded, schema comes from protobuf). - * - * @return the compressed body length in bytes (codec prefix + compressed IPC), or -1 if EOF - */ - public int hasNext() throws IOException { - if (closed) { - return -1; - } - - // Read 16-byte header: clear() resets position=0, limit=capacity, - // preparing the buffer for channel.read() to fill it - headerBuf.clear(); - while (headerBuf.hasRemaining()) { - int bytesRead = channel.read(headerBuf); - if (bytesRead < 0) { - if (headerBuf.position() == 0) { - close(); - return -1; - } - throw new EOFException("Data corrupt: unexpected EOF while reading batch header"); - } - } - headerBuf.flip(); - long compressedLength = headerBuf.getLong(); - // Field count discarded - schema determined by ShuffleScan protobuf fields - headerBuf.getLong(); - - // Subtract 8 because compressedLength includes the 8-byte field count we already read - long 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 = (int) bytesToRead; - - if (dataBuf.capacity() < currentBlockLength) { - int newCapacity = (int) Math.min(bytesToRead * 2L, Integer.MAX_VALUE); - dataBuf = ByteBuffer.allocateDirect(newCapacity); - } - - dataBuf.clear(); - dataBuf.limit(currentBlockLength); - while (dataBuf.hasRemaining()) { - int bytesRead = channel.read(dataBuf); - if (bytesRead < 0) { - throw new EOFException("Data corrupt: unexpected EOF while reading compressed batch"); - } - } - // Note: native side uses get_direct_buffer_address (base pointer) + currentBlockLength, - // not the buffer's position/limit. No flip needed. +public interface CometShuffleBlockIterator { - return currentBlockLength; - } + int hasNext() throws IOException; - /** - * Returns the DirectByteBuffer containing the current block's compressed bytes (4-byte codec - * prefix + compressed IPC data). Called by native code via JNI. - */ - public ByteBuffer getBuffer() { - return dataBuf; - } + ByteBuffer getBuffer(); - /** Returns the length of the current block in bytes. Called by native code via JNI. */ - public int getCurrentBlockLength() { - return currentBlockLength; - } + int getCurrentBlockLength(); - @Override - public void close() throws IOException { - if (!closed) { - closed = true; - inputStream.close(); - } - } + void close() throws IOException; } diff --git a/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java new file mode 100644 index 0000000000..30219e8259 --- /dev/null +++ b/spark/src/main/java/org/apache/comet/shuffle/ShufflePartitionPusher.java @@ -0,0 +1,25 @@ +/* + * 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.comet.shuffle; + +public interface ShufflePartitionPusher { + + int pushPartitionData(int partitionId, byte[] bytes, int length); +} diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index f117372b3c..a3304ea3d8 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -179,8 +179,10 @@ object CometSparkSessionExtensions extends Logging { COMET_EXEC_SHUFFLE_ENABLED.get(conf) && isCometShuffleManagerEnabled(conf) def isCometShuffleManagerEnabled(conf: SQLConf): Boolean = { - conf.contains("spark.shuffle.manager") && conf.getConfString("spark.shuffle.manager") == - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" + conf.contains("spark.shuffle.manager") && (conf.getConfString("spark.shuffle.manager") == + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" || conf.getConfString( + "spark.shuffle.manager") == + "org.apache.spark.sql.comet.uniffle.CometUniffleShuffleManager") } def isCometScan(op: SparkPlan): Boolean = { diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 3cfa51b6e1..a7fd514823 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -25,6 +25,7 @@ import org.apache.spark.{CometTaskMemoryManager, TaskContext} import org.apache.spark.sql.comet.CometMetricNode import org.apache.comet.parquet.CometFileKeyUnwrapper +import org.apache.comet.shuffle.ShufflePartitionPusher class Native extends NativeBase { @@ -254,4 +255,7 @@ class Native extends NativeBase { */ @native def columnarToRowClose(c2rHandle: Long): Unit + @native def createRssPartitionPusher(pusher: ShufflePartitionPusher): Long + + @native def releaseRssPartitionPusher(handle: Long): Long } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBaseNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBaseNativeShuffleWriter.scala new file mode 100644 index 0000000000..4f8023196c --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBaseNativeShuffleWriter.scala @@ -0,0 +1,310 @@ +/* + * 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.execution.shuffle + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition} +import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometScalarSubquery, PlanDataInjector} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.types.StructField + +import org.apache.comet.{CometConf, CometExecIterator} +import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde} +import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator} +import org.apache.comet.serde.PartitioningOuterClass.PartitionWriter +import org.apache.comet.serde.operator.schema2Proto + +trait CometBaseNativeShuffleWriter { + + def spec: NativeShuffleSpec + def outputPartitioning: Partitioning + def outputAttributes: Seq[Attribute] + def rangePartitionBounds: Option[Seq[InternalRow]] + def numParts: Int + def context: TaskContext + + def nativeWrite[K, V]( + inputs: Iterator[Product2[K, V]], + shuffleWriterSQLMetrics: Map[String, SQLMetric], + rssPartitionPusherHandle: Long): Unit = { + val unifiedPlan = buildUnifiedPlan(rssPartitionPusherHandle) + nativeWrite(inputs, shuffleWriterSQLMetrics, unifiedPlan) + } + + def nativeWrite[K, V]( + inputs: Iterator[Product2[K, V]], + shuffleWriterSQLMetrics: Map[String, SQLMetric], + dataFile: String, + indexFile: String): Unit = { + val unifiedPlan = buildUnifiedPlan(dataFile, indexFile) + nativeWrite(inputs, shuffleWriterSQLMetrics, unifiedPlan) + } + + private def nativeWrite[K, V]( + inputs: Iterator[Product2[K, V]], + shuffleWriterSQLMetrics: Map[String, SQLMetric], + unifiedPlan: Operator): Unit = { + val shuffleInputIter = inputs match { + case it: CometNativeShuffleInputIterator => it + case other => + throw new IllegalStateException( + "CometNativeShuffleWriter expects its input iterator to be a " + + "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + + s"${other.getClass.getName}") + } + val partitionIdx = shuffleInputIter.partitionIndex + val inputObjects = shuffleInputIter.inputObjects + val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators + + val ctx = spec.execContext + val finalNativePlan = if (ctx.commonByKey.nonEmpty) { + val partitionDataByKey = ctx.perPartitionByKey.map { case (k, arr) => + k -> arr(partitionIdx) + } + PlanDataInjector.injectPlanData(unifiedPlan, ctx.commonByKey, partitionDataByKey) + } else { + unifiedPlan + } + + // ShuffleWriter metrics at the root; child's metric tree underneath so the SQL UI's per-node + // breakdown matches what the split-driver flow showed. + val nativeMetrics = CometMetricNode(shuffleWriterSQLMetrics, Seq(spec.childMetricNode)) + + // The leaf scans execute inside this writer's single plan rather than a separate native + // stage RDD, so the usual CometExecRDD.compute() bridge (operators.scala) never runs for + // them. Report their bytes/rows to the task's input metrics here instead. + if (ctx.hasScanInput) { + Option(context).foreach(nativeMetrics.reportScanInputMetrics) + } + + val cometIter = new CometExecIterator( + CometExec.newIterId, + inputObjects, + outputAttributes.length, + CometExec.serializeNativePlan(finalNativePlan), + nativeMetrics, + numParts, + partitionIdx, + ctx.broadcastedHadoopConfForEncryption, + ctx.encryptedFilePaths, + shuffleBlockIters) + + // Register subqueries against the iterator id so native callbacks resolve them to values. + ctx.subqueries.foreach { sub => + CometScalarSubquery.setSubquery(cometIter.id, sub) + } + Option(context).foreach { taskCtx => + taskCtx.addTaskCompletionListener[Unit] { _ => + ctx.subqueries.foreach { sub => + CometScalarSubquery.removeSubquery(cometIter.id, sub) + } + } + } + + while (cometIter.hasNext) { + cometIter.next() + } + cometIter.close() + + } + + private def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = { + val partitionWriter = PartitionWriter + .newBuilder() + .setLocalPartitionWriter( + PartitioningOuterClass.LocalPartitionWriter + .newBuilder() + .setOutputDataFile(dataFile) + .setOutputIndexFile(indexFile) + .build()) + .build() + + buildUnifiedPlan(partitionWriter) + } + + protected def buildUnifiedPlan(rssPartitionPusherHandle: Long): Operator = { + val partitionWriter = PartitionWriter + .newBuilder() + .setRssPartitionWriter( + PartitioningOuterClass.RssPartitionWriter + .newBuilder() + .setRssPartitionPusher(rssPartitionPusherHandle) + .build()) + .build() + + buildUnifiedPlan(partitionWriter) + } + + private def buildUnifiedPlan(partitionWriter: PartitionWriter): Operator = { + val shuffleWriterBuilder = OperatorOuterClass.ShuffleWriter.newBuilder() + shuffleWriterBuilder.setPartitionWriter(partitionWriter) + + if (SparkEnv.get.conf.getBoolean("spark.shuffle.compress", true)) { + val codec = CometConf.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC.get() match { + case "zstd" => CompressionCodec.Zstd + case "lz4" => CompressionCodec.Lz4 + case "snappy" => CompressionCodec.Snappy + case other => throw new UnsupportedOperationException(s"invalid codec: $other") + } + shuffleWriterBuilder.setCodec(codec) + } else { + shuffleWriterBuilder.setCodec(CompressionCodec.None) + } + shuffleWriterBuilder.setCompressionLevel( + CometConf.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL.get) + shuffleWriterBuilder.setWriteBufferSize( + CometConf.COMET_SHUFFLE_WRITE_BUFFER_SIZE.get().min(Int.MaxValue).toInt) + + outputPartitioning match { + case p if isSinglePartitioning(p) => + val partitioning = PartitioningOuterClass.SinglePartition.newBuilder() + val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() + shuffleWriterBuilder.setPartitioning( + partitioningBuilder.setSinglePartition(partitioning).build()) + case _: HashPartitioning => + val hashPartitioning = outputPartitioning.asInstanceOf[HashPartitioning] + val partitioning = PartitioningOuterClass.HashPartition.newBuilder() + partitioning.setNumPartitions(outputPartitioning.numPartitions) + + val partitionExprs = hashPartitioning.expressions + .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + + if (partitionExprs.length != hashPartitioning.expressions.length) { + throw new UnsupportedOperationException( + s"Partitioning $hashPartitioning is not supported.") + } + + partitioning.addAllHashExpression(partitionExprs.asJava) + + val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() + shuffleWriterBuilder.setPartitioning( + partitioningBuilder.setHashPartition(partitioning).build()) + case _: RangePartitioning => + val rangePartitioning = outputPartitioning.asInstanceOf[RangePartitioning] + val partitioning = PartitioningOuterClass.RangePartition.newBuilder() + partitioning.setNumPartitions(outputPartitioning.numPartitions) + + // Detect duplicates by tracking expressions directly, similar to DataFusion's LexOrdering + // DataFusion will deduplicate identical sort expressions in LexOrdering, + // so we need to transform boundary rows to match the deduplicated structure + val seenExprs = mutable.HashSet[Expression]() + val deduplicationMap = mutable.ArrayBuffer[(Int, Boolean)]() // (originalIndex, isKept) + + rangePartitioning.ordering.zipWithIndex.foreach { case (sortOrder, idx) => + if (seenExprs.contains(sortOrder.child)) { + deduplicationMap += (idx -> false) + } else { + seenExprs += sortOrder.child + deduplicationMap += (idx -> true) + } + } + + { + val orderingExprs = rangePartitioning.ordering + .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) + if (orderingExprs.length != rangePartitioning.ordering.length) { + throw new UnsupportedOperationException( + s"Partitioning $rangePartitioning is not supported.") + } + partitioning.addAllSortOrders(orderingExprs.asJava) + } + + val boundarySchema = rangePartitioning.ordering.flatMap(e => Some(e.dataType)) + + // rangePartitionBounds holds Spark InternalRows of partitioning boundaries: each row is a + // boundary, each entry a value in that row (row-major, not column-major). Convert to + // Literals and keep only the entries whose ordering expression survived deduplication, so + // the boundary shape matches DataFusion's deduplicated LexOrdering. + val transformedBoundaryExprs: Seq[Seq[Literal]] = + rangePartitionBounds.get.map((row: InternalRow) => { + val allLiterals = + row.toSeq(boundarySchema).zip(boundarySchema).map { case (value, valueType) => + Literal(value, valueType) + } + allLiterals + .zip(deduplicationMap) + .filter(_._2._2) + .map(_._1) + }) + + { + val boundaryRows: Seq[PartitioningOuterClass.BoundaryRow] = transformedBoundaryExprs + .map((rowLiterals: Seq[Literal]) => { + val rowBuilder = PartitioningOuterClass.BoundaryRow.newBuilder(); + val serializedExprs = + rowLiterals.map(lit_value => + QueryPlanSerde.exprToProto(lit_value, outputAttributes).get) + rowBuilder.addAllPartitionBounds(serializedExprs.asJava) + rowBuilder.build() + }) + partitioning.addAllBoundaryRows(boundaryRows.asJava) + } + + val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() + shuffleWriterBuilder.setPartitioning( + partitioningBuilder.setRangePartition(partitioning).build()) + + case _: RoundRobinPartitioning => + val partitioning = PartitioningOuterClass.RoundRobinPartition.newBuilder() + partitioning.setNumPartitions(outputPartitioning.numPartitions) + partitioning.setMaxHashColumns( + CometConf.COMET_EXEC_SHUFFLE_WITH_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) + + val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() + shuffleWriterBuilder.setPartitioning( + partitioningBuilder.setRoundRobinPartition(partitioning).build()) + + case _ => + throw new UnsupportedOperationException( + s"Partitioning $outputPartitioning is not supported.") + } + + shuffleWriterBuilder.setTracingEnabled(CometConf.COMET_TRACING_ENABLED.get()) + + // Used by the native planner to cast the inlined child's output when DataFusion's + // declared return type drifts from Spark catalyst (see comet#4515). + val expectedFields = outputAttributes + .map(a => StructField(a.name, a.dataType, a.nullable, a.metadata)) + .toArray + schema2Proto(expectedFields).foreach(shuffleWriterBuilder.addExpectedOutputSchema) + + OperatorOuterClass.Operator + .newBuilder() + .setShuffleWriter(shuffleWriterBuilder) + .addChildren(spec.childNativeOp) + .build() + } + + private def isSinglePartitioning(p: Partitioning): Boolean = p match { + case SinglePartition => true + case rp: RangePartitioning => + // Spark sometimes generates RangePartitioning schemes with numPartitions == 1, + // or the computed bounds results in a single target partition. + // In this case Comet just serializes a SinglePartition scheme to native. + rp.numPartitions == 1 || rangePartitionBounds.forall(_.isEmpty) + case hp: HashPartitioning => hp.numPartitions == 1 + case _ => false + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala index 14e656f038..a26759c860 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometBlockStoreShuffleReader.scala @@ -35,7 +35,7 @@ import org.apache.spark.storage.BlockManagerId import org.apache.spark.storage.ShuffleBlockFetcherIterator import org.apache.spark.util.CompletionIterator -import org.apache.comet.{CometConf, Native} +import org.apache.comet.{CometConf, CometLocalShuffleBlockIterator, CometNativeShuffleReader, CometShuffleBlockIterator, Native} import org.apache.comet.vector.NativeUtil /** @@ -52,6 +52,7 @@ class CometBlockStoreShuffleReader[K, C]( mapOutputTracker: MapOutputTracker = SparkEnv.get.mapOutputTracker, shouldBatchFetch: Boolean = false) extends ShuffleReader[K, C] + with CometNativeShuffleReader with Logging { private val dep = handle.dependency.asInstanceOf[CometShuffleDependency[_, _, _]] @@ -157,7 +158,7 @@ class CometBlockStoreShuffleReader[K, C]( * Returns the raw concatenated InputStream of all shuffle blocks, bypassing the decode step. * Used by ShuffleScan direct read path. */ - def readAsRawStream(): InputStream = { + private def readAsRawStream(): InputStream = { val streams = fetchIterator.map(_._2) new java.io.SequenceInputStream(new java.util.Enumeration[InputStream] { override def hasMoreElements: Boolean = streams.hasNext @@ -165,6 +166,10 @@ class CometBlockStoreShuffleReader[K, C]( }) } + override def readAsShuffleBlockIterator(): CometShuffleBlockIterator = { + new CometLocalShuffleBlockIterator(this.readAsRawStream()) + } + private def fetchContinuousBlocksInBatch: Boolean = { val conf = SparkEnv.get.conf val serializerRelocatable = dep.serializer.supportsRelocationOfSerializedObjects diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala index a486e2e861..5d0827be27 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala @@ -22,24 +22,14 @@ package org.apache.spark.sql.comet.execution.shuffle import java.nio.{ByteBuffer, ByteOrder} import java.nio.file.{Files, Paths} -import scala.collection.mutable -import scala.jdk.CollectionConverters._ - import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.internal.Logging import org.apache.spark.scheduler.MapStatus import org.apache.spark.shuffle.{IndexShuffleBlockResolver, ShuffleWriteMetricsReporter, ShuffleWriter} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} -import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition} -import org.apache.spark.sql.comet.{CometExec, CometMetricNode, CometScalarSubquery, PlanDataInjector} +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.metric.SQLMetric -import org.apache.spark.sql.types.StructField - -import org.apache.comet.{CometConf, CometExecIterator} -import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass, QueryPlanSerde} -import org.apache.comet.serde.OperatorOuterClass.{CompressionCodec, Operator} -import org.apache.comet.serde.operator.schema2Proto /** * Drives the native shuffle write in a single [[CometExecIterator]] per partition. The plan is @@ -50,7 +40,7 @@ import org.apache.comet.serde.operator.schema2Proto * [[CometShuffleExchangeExec.prepareShuffleDependency]] convenience overload). Same handling * either way. */ -class CometNativeShuffleWriter[K, V]( +case class CometNativeShuffleWriter[K, V]( spec: NativeShuffleSpec, outputPartitioning: Partitioning, outputAttributes: Seq[Attribute], @@ -62,6 +52,7 @@ class CometNativeShuffleWriter[K, V]( metricsReporter: ShuffleWriteMetricsReporter, rangePartitionBounds: Option[Seq[InternalRow]] = None) extends ShuffleWriter[K, V] + with CometBaseNativeShuffleWriter with Logging { private val OFFSET_LENGTH = 8 @@ -79,32 +70,6 @@ class CometNativeShuffleWriter[K, V]( val tempDataFilePath = Paths.get(tempDataFilename) val tempIndexFilePath = Paths.get(tempIndexFilename) - // The dep's _rdd is always a CometNativeShuffleInputRDD on this path. Pattern-match instead - // of asInstanceOf so a future RDD-layering change produces a clear error here rather than a - // bare ClassCastException deeper in the stack. - val shuffleInputIter = inputs match { - case it: CometNativeShuffleInputIterator => it - case other => - throw new IllegalStateException( - "CometNativeShuffleWriter expects its input iterator to be a " + - "CometNativeShuffleInputIterator (produced by CometNativeShuffleInputRDD), got " + - s"${other.getClass.getName}") - } - val partitionIdx = shuffleInputIter.partitionIndex - val inputObjects = shuffleInputIter.inputObjects - val shuffleBlockIters = shuffleInputIter.shuffleBlockIterators - - val unifiedPlan = buildUnifiedPlan(tempDataFilename, tempIndexFilename) - val ctx = spec.execContext - val finalNativePlan = if (ctx.commonByKey.nonEmpty) { - val partitionDataByKey = ctx.perPartitionByKey.map { case (k, arr) => - k -> arr(partitionIdx) - } - PlanDataInjector.injectPlanData(unifiedPlan, ctx.commonByKey, partitionDataByKey) - } else { - unifiedPlan - } - val detailedMetrics = Seq( "elapsed_compute", "encode_time", @@ -120,45 +85,7 @@ class CometNativeShuffleWriter[K, V]( "write_time" -> metricsWriteTime) ++ metrics.filterKeys(detailedMetrics.contains) - // ShuffleWriter metrics at the root; child's metric tree underneath so the SQL UI's per-node - // breakdown matches what the split-driver flow showed. - val nativeMetrics = CometMetricNode(shuffleWriterSQLMetrics, Seq(spec.childMetricNode)) - - // The leaf scans execute inside this writer's single plan rather than a separate native - // stage RDD, so the usual CometExecRDD.compute() bridge (operators.scala) never runs for - // them. Report their bytes/rows to the task's input metrics here instead. - if (ctx.hasScanInput) { - Option(context).foreach(nativeMetrics.reportScanInputMetrics) - } - - val cometIter = new CometExecIterator( - CometExec.newIterId, - inputObjects, - outputAttributes.length, - CometExec.serializeNativePlan(finalNativePlan), - nativeMetrics, - numParts, - partitionIdx, - ctx.broadcastedHadoopConfForEncryption, - ctx.encryptedFilePaths, - shuffleBlockIters) - - // Register subqueries against the iterator id so native callbacks resolve them to values. - ctx.subqueries.foreach { sub => - CometScalarSubquery.setSubquery(cometIter.id, sub) - } - Option(context).foreach { taskCtx => - taskCtx.addTaskCompletionListener[Unit] { _ => - ctx.subqueries.foreach { sub => - CometScalarSubquery.removeSubquery(cometIter.id, sub) - } - } - } - - while (cometIter.hasNext) { - cometIter.next() - } - cometIter.close() + nativeWrite(inputs, shuffleWriterSQLMetrics, tempDataFilename, tempIndexFilename) // get partition lengths from shuffle write output index file var offset = 0L @@ -200,162 +127,6 @@ class CometNativeShuffleWriter[K, V]( MapStatus.apply(SparkEnv.get.blockManager.shuffleServerId, partitionLengths, mapId) } - private def isSinglePartitioning(p: Partitioning): Boolean = p match { - case SinglePartition => true - case rp: RangePartitioning => - // Spark sometimes generates RangePartitioning schemes with numPartitions == 1, - // or the computed bounds results in a single target partition. - // In this case Comet just serializes a SinglePartition scheme to native. - rp.numPartitions == 1 || rangePartitionBounds.forall(_.isEmpty) - case hp: HashPartitioning => hp.numPartitions == 1 - case _ => false - } - - /** - * Build the unified `ShuffleWriter(child = childNativeOp)` plan with the partitioning serde, - * compression settings, and output file paths. - */ - private def buildUnifiedPlan(dataFile: String, indexFile: String): Operator = { - val shuffleWriterBuilder = OperatorOuterClass.ShuffleWriter.newBuilder() - shuffleWriterBuilder.setOutputDataFile(dataFile) - shuffleWriterBuilder.setOutputIndexFile(indexFile) - - if (SparkEnv.get.conf.getBoolean("spark.shuffle.compress", true)) { - val codec = CometConf.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC.get() match { - case "zstd" => CompressionCodec.Zstd - case "lz4" => CompressionCodec.Lz4 - case "snappy" => CompressionCodec.Snappy - case other => throw new UnsupportedOperationException(s"invalid codec: $other") - } - shuffleWriterBuilder.setCodec(codec) - } else { - shuffleWriterBuilder.setCodec(CompressionCodec.None) - } - shuffleWriterBuilder.setCompressionLevel( - CometConf.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL.get) - shuffleWriterBuilder.setWriteBufferSize( - CometConf.COMET_SHUFFLE_WRITE_BUFFER_SIZE.get().min(Int.MaxValue).toInt) - - outputPartitioning match { - case p if isSinglePartitioning(p) => - val partitioning = PartitioningOuterClass.SinglePartition.newBuilder() - val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() - shuffleWriterBuilder.setPartitioning( - partitioningBuilder.setSinglePartition(partitioning).build()) - case _: HashPartitioning => - val hashPartitioning = outputPartitioning.asInstanceOf[HashPartitioning] - val partitioning = PartitioningOuterClass.HashPartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) - - val partitionExprs = hashPartitioning.expressions - .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) - - if (partitionExprs.length != hashPartitioning.expressions.length) { - throw new UnsupportedOperationException( - s"Partitioning $hashPartitioning is not supported.") - } - - partitioning.addAllHashExpression(partitionExprs.asJava) - - val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() - shuffleWriterBuilder.setPartitioning( - partitioningBuilder.setHashPartition(partitioning).build()) - case _: RangePartitioning => - val rangePartitioning = outputPartitioning.asInstanceOf[RangePartitioning] - val partitioning = PartitioningOuterClass.RangePartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) - - // Detect duplicates by tracking expressions directly, similar to DataFusion's LexOrdering - // DataFusion will deduplicate identical sort expressions in LexOrdering, - // so we need to transform boundary rows to match the deduplicated structure - val seenExprs = mutable.HashSet[Expression]() - val deduplicationMap = mutable.ArrayBuffer[(Int, Boolean)]() // (originalIndex, isKept) - - rangePartitioning.ordering.zipWithIndex.foreach { case (sortOrder, idx) => - if (seenExprs.contains(sortOrder.child)) { - deduplicationMap += (idx -> false) - } else { - seenExprs += sortOrder.child - deduplicationMap += (idx -> true) - } - } - - { - val orderingExprs = rangePartitioning.ordering - .flatMap(e => QueryPlanSerde.exprToProto(e, outputAttributes)) - if (orderingExprs.length != rangePartitioning.ordering.length) { - throw new UnsupportedOperationException( - s"Partitioning $rangePartitioning is not supported.") - } - partitioning.addAllSortOrders(orderingExprs.asJava) - } - - val boundarySchema = rangePartitioning.ordering.flatMap(e => Some(e.dataType)) - - // rangePartitionBounds holds Spark InternalRows of partitioning boundaries: each row is a - // boundary, each entry a value in that row (row-major, not column-major). Convert to - // Literals and keep only the entries whose ordering expression survived deduplication, so - // the boundary shape matches DataFusion's deduplicated LexOrdering. - val transformedBoundaryExprs: Seq[Seq[Literal]] = - rangePartitionBounds.get.map((row: InternalRow) => { - val allLiterals = - row.toSeq(boundarySchema).zip(boundarySchema).map { case (value, valueType) => - Literal(value, valueType) - } - allLiterals - .zip(deduplicationMap) - .filter(_._2._2) - .map(_._1) - }) - - { - val boundaryRows: Seq[PartitioningOuterClass.BoundaryRow] = transformedBoundaryExprs - .map((rowLiterals: Seq[Literal]) => { - val rowBuilder = PartitioningOuterClass.BoundaryRow.newBuilder(); - val serializedExprs = - rowLiterals.map(lit_value => - QueryPlanSerde.exprToProto(lit_value, outputAttributes).get) - rowBuilder.addAllPartitionBounds(serializedExprs.asJava) - rowBuilder.build() - }) - partitioning.addAllBoundaryRows(boundaryRows.asJava) - } - - val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() - shuffleWriterBuilder.setPartitioning( - partitioningBuilder.setRangePartition(partitioning).build()) - - case _: RoundRobinPartitioning => - val partitioning = PartitioningOuterClass.RoundRobinPartition.newBuilder() - partitioning.setNumPartitions(outputPartitioning.numPartitions) - partitioning.setMaxHashColumns( - CometConf.COMET_EXEC_SHUFFLE_WITH_ROUND_ROBIN_PARTITIONING_MAX_HASH_COLUMNS.get()) - - val partitioningBuilder = PartitioningOuterClass.Partitioning.newBuilder() - shuffleWriterBuilder.setPartitioning( - partitioningBuilder.setRoundRobinPartition(partitioning).build()) - - case _ => - throw new UnsupportedOperationException( - s"Partitioning $outputPartitioning is not supported.") - } - - shuffleWriterBuilder.setTracingEnabled(CometConf.COMET_TRACING_ENABLED.get()) - - // Used by the native planner to cast the inlined child's output when DataFusion's - // declared return type drifts from Spark catalyst (see comet#4515). - val expectedFields = outputAttributes - .map(a => StructField(a.name, a.dataType, a.nullable, a.metadata)) - .toArray - schema2Proto(expectedFields).foreach(shuffleWriterBuilder.addExpectedOutputSchema) - - OperatorOuterClass.Operator - .newBuilder() - .setShuffleWriter(shuffleWriterBuilder) - .addChildren(spec.childNativeOp) - .build() - } - override def stop(success: Boolean): Option[MapStatus] = { if (success) { Some(mapStatus) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 7629c3e59f..03ce1aca57 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/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 + if (conf.contains("spark.shuffle.manager") && + conf.getConfString("spark.shuffle.manager") == + "org.apache.spark.sql.comet.uniffle.CometUniffleShuffleManager") { + reasons += "CometUniffleShuffleManager does not support jvm columnar shuffle" + return reasons.toSeq + } + if (isShuffleOperator(s.child)) { reasons += s"Child ${s.child.getClass.getName} is a shuffle operator" return reasons.toSeq diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala index bd69e91898..42e267e2e7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala @@ -230,7 +230,7 @@ class CometShuffleManager(conf: SparkConf) extends ShuffleManager with Logging { handle match { case cometShuffleHandle: CometNativeShuffleHandle[K @unchecked, V @unchecked] => val dep = cometShuffleHandle.dependency.asInstanceOf[CometShuffleDependency[_, _, _]] - new CometNativeShuffleWriter( + CometNativeShuffleWriter( dep.nativeShuffleSpec.get, dep.outputPartitioning.get, dep.outputAttributes, diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala index 7604910b06..c8e162e222 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffledRowRDD.scala @@ -21,13 +21,14 @@ package org.apache.spark.sql.comet.execution.shuffle import org.apache.spark._ import org.apache.spark.rdd.RDD +import org.apache.spark.shuffle.ShuffleReader import org.apache.spark.shuffle.sort.SortShuffleManager import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLShuffleReadMetricsReporter} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.vectorized.ColumnarBatch -import org.apache.comet.CometShuffleBlockIterator +import org.apache.comet.{CometNativeShuffleReader, CometRuntimeException, CometShuffleBlockIterator} /** * Different from [[org.apache.spark.sql.execution.ShuffledRowRDD]], this RDD is specialized for @@ -92,9 +93,7 @@ class CometShuffledBatchRDD( } } - private def createReader( - split: Partition, - context: TaskContext): CometBlockStoreShuffleReader[_, _] = { + private def createReader(split: Partition, context: TaskContext): ShuffleReader[_, _] = { val tempMetrics = context.taskMetrics().createTempShuffleReadMetrics() // `SQLShuffleReadMetricsReporter` will update its own metrics for SQL exchange operator, // as well as the `tempMetrics` for basic shuffle metrics. @@ -108,7 +107,6 @@ class CometShuffledBatchRDD( endReducerIndex, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] case PartialReducerPartitionSpec(reducerIndex, startMapIndex, endMapIndex, _) => SparkEnv.get.shuffleManager @@ -120,7 +118,6 @@ class CometShuffledBatchRDD( reducerIndex + 1, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] case PartialMapperPartitionSpec(mapIndex, startReducerIndex, endReducerIndex) => SparkEnv.get.shuffleManager @@ -132,7 +129,6 @@ class CometShuffledBatchRDD( endReducerIndex, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] case CoalescedMapperPartitionSpec(startMapIndex, endMapIndex, numReducers) => SparkEnv.get.shuffleManager @@ -144,7 +140,6 @@ class CometShuffledBatchRDD( numReducers, context, sqlMetricsReporter) - .asInstanceOf[CometBlockStoreShuffleReader[_, _]] } } @@ -155,8 +150,14 @@ class CometShuffledBatchRDD( def computeAsShuffleBlockIterator( split: Partition, context: TaskContext): CometShuffleBlockIterator = { - val reader = createReader(split, context) - new CometShuffleBlockIterator(reader.readAsRawStream()) + createReader(split, context) match { + case nativeShuffleReader: CometNativeShuffleReader => + nativeShuffleReader.readAsShuffleBlockIterator() + case other => + throw new CometRuntimeException( + s"Unexpected shuffle reader type: ${other.getClass.getName}, " + + s"expected an implement of: ${classOf[CometNativeShuffleReader].getName}") + } } override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index 9a59c9c2eb..ffec8c308c 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -64,15 +64,33 @@ abstract class CometTestBase with CometPlanChecker { import testImplicits._ - protected val shuffleManager: String = - "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" + protected def shuffleManagerConfs: Map[String, String] = { + sys.env.getOrElse("COMET_SHUFFLE_MANAGER", "").toLowerCase match { + case "uniffle" => + Map( + SHUFFLE_MANAGER.key -> "org.apache.spark.sql.comet.uniffle.CometUniffleShuffleManager", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native", + // uniffle shuffle does not support jvm columnar shuffle, so we disable it for now + CometConf.COMET_EXEC_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.key -> "false", + "spark.rss.coordinator.quorum" -> "localhost:19999", + "spark.rss.storage.type" -> "MEMORY_LOCALFILE", + "spark.rss.client.type" -> "GRPC_NETTY", + "spark.serializer" -> "org.apache.spark.serializer.KryoSerializer", + "spark.shuffle.service.enabled" -> "false", + "spark.sql.adaptive.localShuffleReader.enabled" -> "false") + case _ => + Map( + SHUFFLE_MANAGER.key -> "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager", + CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") + } + } protected def sparkConf: SparkConf = { val conf = new SparkConf() conf.set("spark.hadoop.fs.file.impl", classOf[DebugFilesystem].getName) conf.set("spark.ui.enabled", "false") conf.set(SQLConf.SHUFFLE_PARTITIONS, 10) // reduce parallelism in tests - conf.set(SHUFFLE_MANAGER, shuffleManager) conf.set(MEMORY_OFFHEAP_ENABLED.key, "true") conf.set(MEMORY_OFFHEAP_SIZE.key, "2g") conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "1g") @@ -99,6 +117,11 @@ abstract class CometTestBase conf.set( "spark.sql.resultQueryStage.maxThreadThreshold", sys.env.getOrElse("SPARK_TEST_SQL_RESULT_QUERY_STAGE_MAX_THREAD_THRESHOLD", "1024")) + + shuffleManagerConfs.foreach { case (k, v) => + conf.set(k, v) + } + conf } diff --git a/thirdparty/comet-uniffle/pom.xml b/thirdparty/comet-uniffle/pom.xml new file mode 100644 index 0000000000..a0d553a9e8 --- /dev/null +++ b/thirdparty/comet-uniffle/pom.xml @@ -0,0 +1,63 @@ + + + + + 4.0.0 + + org.apache.datafusion + comet-parent-spark${spark.version.short}_${scala.binary.version} + 1.0.0-SNAPSHOT + ../../pom.xml + + + comet-uniffle_${scala.binary.version} + comet-uniffle + + + + org.apache.uniffle + rss-client-spark3-shaded + ${uniffle.version} + + + + org.apache.datafusion + comet-spark-spark${spark.version.short}_${scala.binary.version} + ${project.version} + + + + org.apache.spark + spark-sql_${scala.binary.version} + + + + + + + net.alchim31.maven + scala-maven-plugin + + + + + diff --git a/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleManager.scala b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleManager.scala new file mode 100644 index 0000000000..ba190a787a --- /dev/null +++ b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleManager.scala @@ -0,0 +1,227 @@ +/* + * 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.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 + +class CometUniffleShuffleManager(conf: SparkConf, isDriver: Boolean) + extends RssShuffleManager(conf, isDriver) { + + assert( + conf.get(CometConf.COMET_SHUFFLE_MODE.key, "") == "native", + "CometUniffleShuffleManager only supports native shuffle mode") + + override def getReaderImpl[K, C]( + handle: ShuffleHandle, + startMapIndex: Int, + endMapIndex: Int, + startPartition: Int, + endPartition: Int, + context: TaskContext, + metrics: ShuffleReadMetricsReporter, + taskIdBitmap: Roaring64NavigableMap): ShuffleReader[K, C] = { + if (!handle.isInstanceOf[RssShuffleHandle[K, _, C]]) { + throw new RssException("Unexpected ShuffleHandle:" + handle.getClass.getName) + } + val rssShuffleHandle = + handle.asInstanceOf[RssShuffleHandle[K, _, C]] + + rssShuffleHandle.getDependency match { + case dependency: CometShuffleDependency[_, _, _] + if dependency.shuffleType == CometNativeShuffle => + val partitionNum = rssShuffleHandle.getDependency.partitioner.numPartitions + val shuffleId = rssShuffleHandle.getShuffleId + var shuffleHandleInfo: ShuffleHandleInfo = null + if (this.shuffleManagerRpcServiceEnabled && this.rssStageRetryForWriteFailureEnabled) { + shuffleHandleInfo = this.getRemoteShuffleHandleInfoWithStageRetry( + context.stageId, + context.stageAttemptNumber, + shuffleId, + false) + } else if (this.shuffleManagerRpcServiceEnabled && this.partitionReassignEnabled) { + shuffleHandleInfo = this.getRemoteShuffleHandleInfoWithBlockRetry( + context.stageId, + context.stageAttemptNumber, + shuffleId, + false) + } else { + shuffleHandleInfo = new SimpleShuffleHandleInfo( + shuffleId, + rssShuffleHandle.getPartitionToServers, + rssShuffleHandle.getRemoteStorage) + } + val serverToPartitions = + getPartitionDataServers(shuffleHandleInfo, startPartition, endPartition) + val blockIdBitmap = this.getShuffleResultForMultiPart( + this.clientType, + serverToPartitions, + rssShuffleHandle.getAppId, + shuffleId, + context.stageAttemptNumber, + shuffleHandleInfo.createPartitionReplicaTracking) + val readMetrics = context.taskMetrics.shuffleReadMetrics + val shuffleRemoteStorageInfo = rssShuffleHandle.getRemoteStorage + val shuffleRemoteStoragePath = shuffleRemoteStorageInfo.getPath + val readerHadoopConf = RssSparkShuffleUtils.getRemoteStorageHadoopConf( + this.sparkConf, + shuffleRemoteStorageInfo) + new CometUniffleShuffleReader[K, C]( + startPartition, + endPartition, + startMapIndex, + endMapIndex, + context, + rssShuffleHandle, + shuffleRemoteStoragePath, + readerHadoopConf, + partitionNum, + RssUtils.generatePartitionToBitmap( + blockIdBitmap, + startPartition, + endPartition, + this.blockIdLayout), + taskIdBitmap, + readMetrics, + this.managerClientSupplier, + RssSparkConfig.toRssConf(this.sparkConf), + this.dataDistributionType, + shuffleHandleInfo.getAllPartitionServersForReader) + case _ => + super.getReaderImpl( + handle, + startMapIndex, + endMapIndex, + startPartition, + endPartition, + context, + metrics, + taskIdBitmap) + } + } + + private def getShuffleResultForMultiPart( + clientType: String, + serverToPartitions: util.Map[ShuffleServerInfo, util.Set[Integer]], + appId: String, + shuffleId: Int, + stageAttemptId: Int, + replicaRequirementTracking: PartitionDataReplicaRequirementTracking) = { + val failedPartitions: util.Set[Integer] = Sets.newHashSet[Integer]() + try { + this.shuffleWriteClient.getShuffleResultForMultiPart( + clientType, + serverToPartitions, + appId, + shuffleId, + failedPartitions, + replicaRequirementTracking) + } catch { + case e: RssFetchFailedException => + throw RssSparkShuffleUtils.reportRssFetchFailedException( + this.managerClientSupplier, + e, + this.sparkConf, + appId, + shuffleId, + stageAttemptId, + failedPartitions) + } + } + + private def getPartitionDataServers( + shuffleHandleInfo: ShuffleHandleInfo, + startPartition: Int, + endPartition: Int): util.Map[ShuffleServerInfo, util.Set[Integer]] = { + val allPartitionToServers = shuffleHandleInfo.getAllPartitionServersForReader.asScala + val requirePartitionToServers = + allPartitionToServers.filter(x => x._1 >= startPartition && x._1 < endPartition).asJava + val serverToPartitions = RssUtils.generateServerToPartitions(requirePartitionToServers) + serverToPartitions + } + + override def getWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + context: TaskContext, + metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = { + if (!handle.isInstanceOf[RssShuffleHandle[_, _, _]]) { + throw new RssException("Unexpected ShuffleHandle:" + handle.getClass.getName) + } + val rssHandle: RssShuffleHandle[K, V, V] = handle.asInstanceOf[RssShuffleHandle[K, V, V]] + + rssHandle.getDependency match { + case dependency: CometShuffleDependency[_, _, _] + if dependency.shuffleType == CometNativeShuffle => + setPusherAppId(rssHandle) + val taskId: String = context.taskAttemptId + "_" + context.attemptNumber + var writeMetrics: ShuffleWriteMetrics = null + if (metrics != null) { + writeMetrics = new RssShuffleManager.WriteMetrics(metrics) + } else { + writeMetrics = context.taskMetrics.shuffleWriteMetrics + } + // set rss.row.based to false to mark it as columnar shuffle + val conf: SparkConf = sparkConf.clone + .set(RssSparkConfig.SPARK_RSS_CONFIG_PREFIX + RssSparkConfig.RSS_ROW_BASED.key, "false") + .set("spark.rss.client.io.compression.codec", "NONE") + + CometUniffleShuffleWriter[K, V, V]( + dependency.nativeShuffleSpec.get, + dependency.outputPartitioning.get, + dependency.outputAttributes, + dependency.numParts, + dependency.shuffleWriteMetrics, + rssHandle.getAppId, + rssHandle.getShuffleId, + taskId, + context.taskAttemptId, + writeMetrics, + this, + conf, + shuffleWriteClient, + rssHandle, + (taskId: String) => markFailedTask(taskId), + context, + dependency.rangePartitionBounds) + case dependency: CometShuffleDependency[_, _, _] => + throw new RssException("Unexpected shuffleType:" + dependency.shuffleType) + case _ => + super.getWriter(handle, mapId, context, metrics) + } + + } + +} diff --git a/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala new file mode 100644 index 0000000000..f2cc0c3b7a --- /dev/null +++ b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleReader.scala @@ -0,0 +1,356 @@ +/* + * 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.{CometNativeShuffleReader, CometShuffleBlockIterator, Native} +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 + + override def hasNext: Boolean = { + shuffleBlockIterator.hasNext != -1 + } + + override def next(): (Int, ColumnarBatch) = { + 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 = { + 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] = { + val shuffleBlock: ShuffleBlock = if (currentShuffleReadClient != null) { + currentShuffleReadClient.readShuffleBlockData + } else { + null + } + val rawData = if (shuffleBlock != null) { + shuffleBlock.getByteBuffer + } else { + null + } + if (rawData == null) { + currentPartition += 1 + if (currentPartition >= endPartition) { + return None + } + if (currentShuffleReadClient != null) { + currentShuffleReadClient.checkProcessedBlockIds() + currentShuffleReadClient.logStatics() + currentShuffleReadClient.close() + } + if (partitionToExpectBlocks.get(currentPartition).isEmpty) { + currentShuffleReadClient = null + } else { + currentShuffleReadClient = createShuffleReadClient(currentPartition) + } + return nextShuffleBlock() + } + Some(rawData) + } + + override def getBuffer: ByteBuffer = dataBuf + + override def getCurrentBlockLength: Int = currentBlockLength + + override def close(): Unit = { + if (current != null) { + current = null + } + if (dataBuf != null) { + dataBuf = ByteBuffer.allocateDirect(INITIAL_BUFFER_SIZE) + dataBuf = null + } + if (currentShuffleReadClient != null) { + currentShuffleReadClient.close() + currentShuffleReadClient = null + } + } + } +} diff --git a/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleWriter.scala b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleWriter.scala new file mode 100644 index 0000000000..d315aad59a --- /dev/null +++ b/thirdparty/comet-uniffle/src/main/scala/org/apache/spark/sql/comet/uniffle/CometUniffleShuffleWriter.scala @@ -0,0 +1,151 @@ +/* + * 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 + sendCommit() + sendRestBlockAndWait() + 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() + length + } + + private var stopped = false + override def stop(success: Boolean): Option[MapStatus] = { + if (!stopped) { + stopped = true + releaseShufflePusher() + val result = super.stop(success) + result + } else { + Option.empty + } + } + + // release native shuffle pusher + private def releaseShufflePusher(): Unit = { + if (nativeShufflePusherHandle != -1) { + nativeLib.releaseRssPartitionPusher(nativeShufflePusherHandle) + nativeShufflePusherHandle = -1 + } + } + + private def sendRestBlockAndWait(): Unit = { + val shuffleBlockInfos = super.getBufferManager.clear() + super.processShuffleBlockInfos(shuffleBlockInfos); + // make checkBlockSendResult no arguments + super.internalCheckBlockSendResult(); + } + +}